[tor-commits] [ooni-probe/master] If no country code is specified lookup the country code of the user.

[art at torproject.org]

commit 4c09c008e3f634a12fd971263d8ade77d87565db
Author: Arturo Filastò <art at 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
 
+
+ at defer.inlineCallbacks
+def get_user_country_code():
+    probe_ip = ProbeIP()
+    yield probe_ip.lookup()
+    defer.returnValue(probe_ip.geodata['countrycode'])
+
+
+ at 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()