commit 4c09c008e3f634a12fd971263d8ade77d87565db Author: Arturo Filastò <art@fuffa.org> Date: Wed Aug 13 16:02:34 2014 +0200 If no country code is specified lookup the country code of the user. --- bin/oonideckgen | 7 +++- ooni/deckgen/cli.py | 104 +++++++++++++++++++++++++++++++-------------------- 2 files changed, 68 insertions(+), 43 deletions(-) diff --git a/bin/oonideckgen b/bin/oonideckgen index c990160..74ea8de 100755 --- a/bin/oonideckgen +++ b/bin/oonideckgen @@ -1,6 +1,7 @@ #!/usr/bin/env python import os import sys +import exceptions sys.path[:] = map(os.path.abspath, sys.path) sys.path.insert(0, os.path.abspath(os.getcwd())) @@ -11,8 +12,10 @@ from ooni.utils import log from ooni.deckgen import cli def failed(failure): - log.err("Failed to run oonideckgen") - log.exception(failure) + r = failure.trap(exceptions.SystemExit) + if r != exceptions.SystemExit: + log.err("Failed to run oonideckgen") + log.exception(failure) reactor.stop() def done(result): diff --git a/ooni/deckgen/cli.py b/ooni/deckgen/cli.py index 406417d..95e0c98 100644 --- a/ooni/deckgen/cli.py +++ b/ooni/deckgen/cli.py @@ -5,16 +5,19 @@ import errno import yaml +from twisted.internet import defer, reactor from twisted.python import usage -from . import __version__ -from ooni.resources import inputs +from ooni.geoip import ProbeIP from ooni.settings import config +from ooni.deckgen import __version__ +from ooni.resources import inputs + class Options(usage.Options): synopsis = """%s [options] - """ + """ % sys.argv[0] optParameters = [ ["country-code", "c", @@ -32,23 +35,6 @@ class Options(usage.Options): print("oonideckgen version: %s" % __version__) sys.exit(0) - def postOptions(self): - if not self['output'] or not self['country-code']: - raise usage.UsageError( - "Both --output and --country-code are required" - ) - if len(self['country-code']) != 2: - raise usage.UsageError("--country-code must be 2 characters") - if not os.path.isdir(self['output']): - raise usage.UsageError("%s is not a directory" % self['output']) - - self['country-code'] = self['country-code'].lower() - - output_dir = os.path.abspath(self['output']) - output_dir = os.path.join(output_dir, - "deck-%s" % self['country-code']) - self['output'] = output_dir - class Deck(): _base_entry = { @@ -82,26 +68,7 @@ class Deck(): f.write(yaml.safe_dump(self.deck)) -def usage(): - print "%s <two letter country code> <output dir>" % sys.argv[0] - - -def run(): - options = Options() - try: - options.parseOptions() - except usage.UsageError as error_message: - print "%s: %s" % (sys.argv[0], error_message) - print "%s: Try --help for usage details." % (sys.argv[0]) - sys.exit(1) - - config.read_config_file() - - try: - os.makedirs(options['output']) - except OSError as exception: - if exception.errno != errno.EEXIST: - raise +def generate_deck(options): dns_servers_processor = inputs['namebench-dns-servers.csv']['processor'] url_lists_processor = inputs['citizenlab-test-lists.zip']['processor'] @@ -147,5 +114,60 @@ def run(): print "Run ooniprobe like so:" print "ooniprobe -i %s" % deck_filename + +@defer.inlineCallbacks +def get_user_country_code(): + probe_ip = ProbeIP() + yield probe_ip.lookup() + defer.returnValue(probe_ip.geodata['countrycode']) + + +@defer.inlineCallbacks +def run(): + options = Options() + try: + options.parseOptions() + except usage.UsageError as error_message: + print "%s: %s" % (sys.argv[0], error_message) + print options + sys.exit(1) + + if not options['output']: + options['output'] = os.getcwd() + + if not options['country-code']: + options['country-code'] = yield get_user_country_code() + + if len(options['country-code']) != 2: + print "%s: --country-code must be 2 characters" % sys.argv[0] + sys.exit(2) + + if not os.path.isdir(options['output']): + print "%s: %s is not a directory" % (sys.argv[0], + options['output']) + sys.exit(3) + + options['country-code'] = options['country-code'].lower() + + output_dir = os.path.abspath(options['output']) + output_dir = os.path.join(output_dir, + "deck-%s" % options['country-code']) + options['output'] = output_dir + + config.read_config_file() + + try: + os.makedirs(options['output']) + except OSError as exception: + if exception.errno != errno.EEXIST: + raise + + generate_deck(options) + if __name__ == "__main__": - run() + d = run() + + @d.addBoth + def cb(_): + reactor.stop() + reactor.start()