commit 0452fa3e05e8d0d01dc2c135d166927653dd0cdd Author: Arturo Filastò <arturo@filasto.net> Date: Sun Oct 7 16:06:30 2012 +0000 Implement geodata lookups to be included in every report. --- ooni/inputunit.py | 13 +++++++++++++ ooni/nettest.py | 18 +++++++++++++++++- ooni/reporter.py | 41 +++++++++++++++++++++++++++++++---------- ooni/runner.py | 16 +++++++++++----- ooni/utils/geodata.py | 7 ++++--- 5 files changed, 76 insertions(+), 19 deletions(-) diff --git a/ooni/inputunit.py b/ooni/inputunit.py index 69507b4..157f038 100644 --- a/ooni/inputunit.py +++ b/ooni/inputunit.py @@ -1,3 +1,16 @@ +from twisted.trial import unittest + +class PatchedPyUnitResultAdapter(unittest.PyUnitResultAdapter): + def __init__(self, original): + """ + Here we patch PyUnitResultAdapter to support our reporterFactory to + properly write headers to reports. + """ + self.original = original + self.reporterFactory = original.reporterFactory + +unittest.PyUnitResultAdapter = PatchedPyUnitResultAdapter + class InputUnitFactory(object): """ This is a factory that takes the size of input units to be generated a set diff --git a/ooni/nettest.py b/ooni/nettest.py index 89dd279..835a5a0 100644 --- a/ooni/nettest.py +++ b/ooni/nettest.py @@ -1,6 +1,7 @@ import itertools from twisted.python import log from twisted.trial import unittest, itrial +from twisted.internet import defer pyunit = __import__('unittest') @@ -71,12 +72,27 @@ class TestCase(unittest.TestCase): inputs = [None] inputFile = None - report = {} report['errors'] = [] optParameters = None + def deferSetUp(self, ignored, result): + """ + If we have the reporterFactory set we need to write the header. If such + method is not present we will only run the test skipping header + writing. + """ + if result.reporterFactory.firstrun: + print "Running both!!" + d1 = result.reporterFactory.writeHeader() + d2 = unittest.TestCase.deferSetUp(self, ignored, result) + dl = defer.DeferredList([d1, d2]) + return dl + else: + print "Only one :P" + return unittest.TestCase.deferSetUp(self, ignored, result) + def inputProcessor(self, fp): for x in fp.readlines(): yield x.strip() diff --git a/ooni/reporter.py b/ooni/reporter.py index 86c1fc1..b335738 100644 --- a/ooni/reporter.py +++ b/ooni/reporter.py @@ -6,7 +6,8 @@ import itertools from datetime import datetime from twisted.python.util import OrderedDict, untilConcludes from twisted.trial import unittest, reporter, runner -from ooni.utils import date +from twisted.internet import defer +from ooni.utils import date, log, geodata try: from scapy.all import packet @@ -24,6 +25,8 @@ class OReporter(pyunit.TestResult): This is an extension of the unittest TestResult. It adds support for reporting to yaml format. """ + reporterFactory = None + def __init__(self, stream=sys.stdout, tbformat='default', realtime=False, publisher=None, testSuite=None): super(OReporter, self).__init__() @@ -62,6 +65,8 @@ class ReporterFactory(OReporter): This is a reporter factory. It emits new instances of Reports. It is also responsible for writing the OONI Report headers. """ + firstrun = True + def __init__(self, stream=sys.stdout, tbformat='default', realtime=False, publisher=None, testSuite=None): super(ReporterFactory, self).__init__(stream=stream, @@ -70,25 +75,41 @@ class ReporterFactory(OReporter): self._testSuite = testSuite self._reporters = [] - def writeHeader(self, options, geodata={}): + @defer.inlineCallbacks + def writeHeader(self): + self.firstrun = False + options = self.options self._writeln("###########################################") self._writeln("# OONI Probe Report for %s test" % options['name']) self._writeln("# %s" % date.pretty_date()) self._writeln("###########################################") - address = {'asn': 'unknown', - 'ip': 'unknown'} - if 'ip' in geodata: - address['ip'] = geodata['ip'] + client_geodata = {} + log.msg("Running geo IP lookup via check.torproject.org") + + client_ip = yield geodata.myIP() + try: + import txtorcon + client_location = txtorcon.util.NetLocation(client_ip) + except: + log.err("txtorcon is not installed. Geolocation lookup is not"\ + "supported") - if 'asn' in geodata: - address['asn'] = geodata['asn'] + client_geodata['ip'] = client_ip + client_geodata['asn'] = client_location.asn + client_geodata['city'] = client_location.city + client_geodata['countrycode'] = client_location.countrycode test_details = {'startTime': repr(date.now()), - 'probeASN': address['asn'], + 'probeASN': client_geodata['asn'], + 'probeCC': client_geodata['countrycode'], + 'probeIP': client_geodata['ip'], + 'probeLocation': {'city': client_geodata['city'], + 'countrycode': + client_geodata['countrycode']}, 'testName': options['name'], 'testVersion': options['version'], - 'probeIP': address['ip']} + } self.writeYamlLine(test_details) self._writeln('') diff --git a/ooni/runner.py b/ooni/runner.py index a753d02..81272eb 100644 --- a/ooni/runner.py +++ b/ooni/runner.py @@ -16,7 +16,7 @@ from ooni.reporter import ReporterFactory from ooni.inputunit import InputUnitFactory from ooni.nettest import InputTestSuite from ooni import nettest -from ooni.utils import log +from ooni.utils import log, geodata from ooni.plugoo import tests as oonitests def isTestCase(thing): @@ -97,13 +97,16 @@ def processTest(obj, config): optParameters = obj.optParameters inputFile = obj.inputFile - Options.optParameters.append(inputFile) + if inputFile: + Options.optParameters.append(inputFile) options = Options() options.parseOptions(config['subArgs']) obj.localOptions = options - obj.inputFile = options[inputFile[0]] + + if inputFile: + obj.inputFile = options[inputFile[0]] return obj @@ -186,7 +189,10 @@ class ORunner(object): def runWithInputUnit(self, inputUnit): idx = 0 result = self.reporterFactory.create() + for input in inputUnit: + result.reporterFactory = self.reporterFactory + suite = self.baseSuite(self.cases) suite.input = input suite(result, idx) @@ -204,9 +210,9 @@ class ORunner(object): def run(self): #log.startLogging(sys.stdout) log.start() - self.reporterFactory.writeHeader(self.options) + + self.reporterFactory.options = self.options for inputUnit in InputUnitFactory(self.inputs): self.runWithInputUnit(inputUnit) - diff --git a/ooni/utils/geodata.py b/ooni/utils/geodata.py index 3e5202e..9f782e8 100644 --- a/ooni/utils/geodata.py +++ b/ooni/utils/geodata.py @@ -17,11 +17,13 @@ class BodyReceiver(protocol.Protocol): def myIP(): target_site = 'https://check.torproject.org/' regexp = "Your IP address appears to be: <b>(.+?)<\/b>" - myAgent = Agent(reactor) + result = yield myAgent.request('GET', target_site) + finished = defer.Deferred() result.deliverBody(BodyReceiver(finished)) + body = yield finished match = re.search(regexp, body) @@ -30,6 +32,5 @@ def myIP(): except: myip = "unknown" - return myip - + defer.returnValue(myip)