Parse options correctly; call ./ipt automatically.

3 files changed, 70 insertions, 9 deletions

diff --git a/client.py b/client.py
index a663b1a..364bcf8 100644
--- a/client.py
+++ b/client.py
@@ -1,4 +1,4 @@
-import struct, socket, select
+import struct, socket, select, subprocess
 from ssnet import SockWrapper, Handler, Proxy
 from helpers import *
 
@@ -13,14 +13,24 @@ def original_dst(sock):
     return (ip,port)
 
 
-def main(remotename, subnets):
+def iptables_setup(port, subnets):
+    subnets_str = ['%s/%d' % (ip,width) for ip,width in subnets]
+    argv = ['sudo', sys.argv[0], '--iptables', str(port)] + subnets_str
+    rv = subprocess.call(argv)
+    if rv != 0:
+        raise Exception('%r returned %d' % (argv, rv))
+
+
+def main(listenip, remotename, subnets):
     log('Starting sshuttle proxy.\n')
     listener = socket.socket()
     listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-    listener.bind(('0.0.0.0',1234))
+    listener.bind(listenip)
     listener.listen(10)
     log('Listening on %r.\n' % (listener.getsockname(),))
 
+    iptables_setup(listenip[1], subnets)
+
     handlers = []
     def onaccept():
         sock,srcip = listener.accept()
diff --git a/iptables.py b/iptables.py
new file mode 100644
index 0000000..270423d
--- /dev/null
+++ b/iptables.py
@@ -0,0 +1,5 @@
+import subprocess
+
+def main(port, subnets):
+    subnets_str = ['%s/%d' % (ip,width) for ip,width in subnets]
+    subprocess.call(['./ipt', str(port)] + subnets_str)
diff --git a/main.py b/main.py
index cf34bec..8d7d24e 100755
--- a/main.py
+++ b/main.py
@@ -1,9 +1,52 @@
 #!/usr/bin/env python
-import sys
-import options, client
+import sys, re
+import options, client, iptables
+
+
+# list of:
+# 1.2.3.4/5 or just 1.2.3.4
+def parse_subnets(subnets_str):
+    subnets = []
+    for s in subnets_str:
+        m = re.match(r'(\d+)\.(\d+)\.(\d+)\.(\d+)(?:/(\d+))?$', s)
+        if not m:
+            raise Exception('%r is not a valid IP subnet format' % s)
+        (a,b,c,d,width) = m.groups()
+        (a,b,c,d) = (int(a or 0), int(b or 0), int(c or 0), int(d or 0))
+        if width == None:
+            width = 32
+        else:
+            width = int(width)
+        if a > 255 or b > 255 or c > 255 or d > 255:
+            raise Exception('%d.%d.%d.%d has numbers > 255' % (a,b,c,d))
+        if width > 32:
+            raise Exception('*/%d is greater than the maximum of 32' % width)
+        subnets.append(('%d.%d.%d.%d' % (a,b,c,d), width))
+    return subnets
+
+
+# 1.2.3.4:567 or just 1.2.3.4 or just 567
+def parse_ipport(s):
+    s = str(s)
+    m = re.match(r'(?:(\d+)\.(\d+)\.(\d+)\.(\d+))?(?::)?(?:(\d+))?$', s)
+    if not m:
+        raise Exception('%r is not a valid IP:port format' % s)
+    (a,b,c,d,port) = m.groups()
+    (a,b,c,d,port) = (int(a or 0), int(b or 0), int(c or 0), int(d or 0),
+                      int(port or 0))
+    if a > 255 or b > 255 or c > 255 or d > 255:
+        raise Exception('%d.%d.%d.%d has numbers > 255' % (a,b,c,d))
+    if port > 65535:
+        raise Exception('*:%d is greater than the maximum of 65535' % port)
+    if a == None:
+        a = b = c = d = 0
+    return ('%d.%d.%d.%d' % (a,b,c,d), port)
+
 
 optspec = """
 sshuttle [-l [ip:]port] [-r [username@]sshserver] <subnets...>
+sshuttle --iptables <port> <subnets...>
+sshuttle --server
 --
 l,listen=  transproxy to this ip address and port number [default=0]
 r,remote=  ssh hostname (and optional username) of remote sshuttle server
@@ -17,13 +60,16 @@ if opt.server:
     o.fatal('server mode not implemented yet')
     sys.exit(1)
 elif opt.iptables:
-    o.fatal('iptables mode not implemented yet')
-    sys.exit(1)
+    if len(extra) < 1:
+        o.fatal('at least one argument expected')
+    sys.exit(iptables.main(int(extra[0]),
+                           parse_subnets(extra[1:])))
 else:
     if len(extra) < 1:
         o.fatal('at least one argument expected')
     remotename = extra[0]
     if remotename == '' or remotename == '-':
         remotename = None
-    subnets = extra[1:]
-    sys.exit(client.main(remotename, subnets))
+    sys.exit(client.main(parse_ipport(opt.listen or '0.0.0.0:1234'),
+                         remotename,
+                         parse_subnets(extra)))