firewall.py: clean up repeated calls to ssubprocess.call().

And make sshuttle exit with a well-defined exit code (111) if it needs to
reboot.

4 files changed, 24 insertions, 19 deletions

diff --git a/client.py b/client.py
index 7310c25..06fc573 100644
--- a/client.py
+++ b/client.py
@@ -171,7 +171,9 @@ class FirewallClient:
     def done(self):
         self.pfile.close()
         rv = self.p.wait()
-        if rv:
+        if rv == EXITCODE_NEEDS_REBOOT:
+            raise FatalNeedsReboot()
+        elif rv:
             raise Fatal('cleanup: %r returned %d' % (self.argv, rv))
 
 
diff --git a/firewall.py b/firewall.py
index 83ff6dc..de85bd0 100644
--- a/firewall.py
+++ b/firewall.py
@@ -20,6 +20,14 @@ def nonfatal(func, *args):
         log('error: %s\n' % e)
 
 
+def _call(argv):
+    debug1('>> %s\n' % ' '.join(argv))
+    rv = ssubprocess.call(argv)
+    if rv:
+        raise Fatal('%r returned %d' % (argv, rv))
+    return rv
+
+
 def ipt_chain_exists(name):
     argv = ['iptables', '-t', 'nat', '-nL']
     p = ssubprocess.Popen(argv, stdout = ssubprocess.PIPE)
@@ -33,10 +41,7 @@ def ipt_chain_exists(name):
 
 def ipt(*args):
     argv = ['iptables', '-t', 'nat'] + list(args)
-    debug1('>> %s\n' % ' '.join(argv))
-    rv = ssubprocess.call(argv)
-    if rv:
-        raise Fatal('%r returned %d' % (argv, rv))
+    _call(argv)
 
 
 _no_ttl_module = False
@@ -159,15 +164,9 @@ def _defaults_write_kernel_flags(flags):
     flagstr = ' '.join(flags)
     argv = ['defaults', 'write', KERNEL_FLAGS_PATH, KERNEL_FLAGS_NAME,
             flagstr]
-    debug1('>> %s\n' % ' '.join(argv))
-    rv = ssubprocess.call(argv)
-    if rv:
-        raise Fatal('%r returned %d' (argv, rv))
+    _call(argv)
     argv = ['plutil', '-convert', 'xml1', KERNEL_FLAGS_PATH + '.plist']
-    debug1('>> %s\n' % ' '.join(argv))
-    rv = ssubprocess.call(argv)
-    if rv:
-        raise Fatal('%r returned %d' (argv, rv))
+    _call(argv)
     
 
 
@@ -253,10 +252,7 @@ def _handle_diversion(divertsock, dnsport):
 
 def ipfw(*args):
     argv = ['ipfw', '-q'] + list(args)
-    debug1('>> %s\n' % ' '.join(argv))
-    rv = ssubprocess.call(argv)
-    if rv:
-        raise Fatal('%r returned %d' % (argv, rv))
+    _call(argv)
 
 
 def do_ipfw(port, dnsport, subnets):
@@ -296,8 +292,7 @@ def do_ipfw(port, dnsport, subnets):
                 "to work around a bug in MacOS 10.7 Lion.  You will need\n"
                 "to reboot before it takes effect.  You only have to\n"
                 "do this once.\n\n")
-            sys.exit(1)
-                
+            sys.exit(EXITCODE_NEEDS_REBOOT)
 
         ipfw('add', sport, 'check-state', 'ip',
              'from', 'any', 'to', 'any')
diff --git a/helpers.py b/helpers.py
index af49788..45a028b 100644
--- a/helpers.py
+++ b/helpers.py
@@ -30,6 +30,11 @@ class Fatal(Exception):
     pass
 
 
+EXITCODE_NEEDS_REBOOT = 111
+class FatalNeedsReboot(Fatal):
+    pass
+
+
 def list_contains_any(l, sub):
     for i in sub:
         if i in l:
diff --git a/main.py b/main.py
index a88ca9d..daf3b87 100755
--- a/main.py
+++ b/main.py
@@ -126,6 +126,9 @@ try:
                              parse_subnets(includes),
                              parse_subnets(excludes),
                              opt.syslog, opt.daemon, opt.pidfile))
+except FatalNeedsReboot, e:
+    log('You must reboot before using sshuttle.\n')
+    sys.exit(EXITCODE_NEEDS_REBOOT)
 except Fatal, e:
     log('fatal: %s\n' % e)
     sys.exit(99)