iptables.py: completely replace ipt script.

Doing it in python instead of shell makes the code a bit less error prone.
Plus we can parse the iptables output and avoid triggering iptables errors.

4 files changed, 65 insertions, 41 deletions

diff --git a/client.py b/client.py
index 364bcf8..54fd64b 100644
--- a/client.py
+++ b/client.py
@@ -21,16 +21,7 @@ def iptables_setup(port, subnets):
         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(listenip)
-    listener.listen(10)
-    log('Listening on %r.\n' % (listener.getsockname(),))
-
-    iptables_setup(listenip[1], subnets)
-
+def _main(listener, remotename, subnets):
     handlers = []
     def onaccept():
         sock,srcip = listener.accept()
@@ -61,3 +52,19 @@ def main(listenip, remotename, subnets):
         for s in handlers:
             if s.socks & ready:
                 s.callback()
+
+
+def main(listenip, remotename, subnets):
+    log('Starting sshuttle proxy.\n')
+    listener = socket.socket()
+    listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    listener.bind(listenip)
+    listener.listen(10)
+    log('Listening on %r.\n' % (listener.getsockname(),))
+
+    iptables_setup(listenip[1], subnets)
+
+    try:
+        return _main(listener, remotename, subnets)
+    finally:
+        iptables_setup(listenip[1], [])
diff --git a/ipt b/ipt
deleted file mode 100755
index 6ec37db..0000000
--- a/ipt
+++ /dev/null
@@ -1,28 +0,0 @@
-#!/bin/bash -x
-PORT="$1"
-shift
-
-if [ -z "$PORT" ] || ! [ "$PORT" -gt 0 ]; then
-	echo "'$PORT' is not a valid port number"
-	exit 1
-fi
-
-# basic cleanup/setup
-C=sshuttle-$PORT
-iptables -t nat -D OUTPUT -j $C
-iptables -t nat -F $C
-iptables -t nat -X $C
-
-if [ -z "$*" ]; then
-    # just delete existing rules
-    exit 0
-fi
-iptables -t nat -N $C
-iptables -t nat -I OUTPUT 1 -j $C
-iptables -t nat -D $C -j REDIRECT -p tcp --to-ports $PORT
-
-# create new subnet entries
-for subnet in "$@"; do
-	iptables -t nat -A $C -j REDIRECT --dest "$subnet" -p tcp \
-		--to-ports "$PORT" -m ttl \! --ttl 42
-done
diff --git a/iptables.py b/iptables.py
index 270423d..a65ae9d 100644
--- a/iptables.py
+++ b/iptables.py
@@ -1,5 +1,50 @@
-import subprocess
+import subprocess, re
+from helpers import *
 
+
+def chain_exists(name):
+    argv = ['iptables', '-t', 'nat', '-nL']
+    p = subprocess.Popen(argv, stdout = subprocess.PIPE)
+    for line in p.stdout:
+        if line.startswith('Chain %s ' % name):
+            return True
+    rv = p.wait()
+    if rv:
+        raise Exception('%r returned %d' % (argv, rv))
+
+
+def ipt(*args):
+    argv = ['iptables', '-t', 'nat'] + list(args)
+    log('>> %s\n' % ' '.join(argv))
+    rv = subprocess.call(argv)
+    if rv:
+        raise Exception('%r returned %d' % (argv, rv))
+
+
+# FIXME: this prints scary-looking errors
 def main(port, subnets):
+    assert(port > 0)
+    assert(port <= 65535)
+
+    chain = 'sshuttle-%s' % port
+
+    # basic cleanup/setup of chains
+    if chain_exists(chain):
+        ipt('-D', 'OUTPUT', '-j', chain)
+        ipt('-F', chain)
+        ipt('-X', chain)
+
+    if subnets:
+        ipt('-N', chain)
+        ipt('-F', chain)
+        ipt('-I', 'OUTPUT', '1', '-j', chain)
+
+        # create new subnet entries
+        for snet,swidth in subnets:
+            ipt('-A', chain, '-j', 'REDIRECT',
+                '--dest', '%s/%s' % (snet,swidth),
+                '-p', 'tcp',
+                '--to-ports', str(port),
+                '-m', 'ttl', '!', '--ttl', '42'  # to prevent infinite loops
+                )
     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 8d7d24e..0beaf0b 100755
--- a/main.py
+++ b/main.py
@@ -8,7 +8,7 @@ import options, client, iptables
 def parse_subnets(subnets_str):
     subnets = []
     for s in subnets_str:
-        m = re.match(r'(\d+)\.(\d+)\.(\d+)\.(\d+)(?:/(\d+))?$', s)
+        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()