Add -N (--auto-nets) option for auto-discovering subnets.

Now if you do

	./sshuttle -Nr username@myservername

It'll automatically route the "local" subnets (ie., stuff in the routing
table) from myservername.  This is (hopefully a reasonable default setting
for most people.

6 files changed, 131 insertions, 22 deletions

diff --git a/assembler.py b/assembler.py
index 1dce7c2..c478e37 100644
--- a/assembler.py
+++ b/assembler.py
@@ -7,7 +7,7 @@ while 1:
     if name:
         nbytes = int(sys.stdin.readline())
         if verbosity >= 2:
-            sys.stderr.write('remote assembling %r (%d bytes)\n'
+            sys.stderr.write('server: assembling %r (%d bytes)\n'
                              % (name, nbytes))
         content = z.decompress(sys.stdin.read(nbytes))
         exec compile(content, name, "exec")
diff --git a/client.py b/client.py
index 718410f..c9c3255 100644
--- a/client.py
+++ b/client.py
@@ -22,11 +22,11 @@ def original_dst(sock):
 class FirewallClient:
     def __init__(self, port, subnets):
         self.port = port
+        self.auto_nets = []
         self.subnets = subnets
-        subnets_str = ['%s/%d' % (ip,width) for ip,width in subnets]
         argvbase = ([sys.argv[0]] +
                     ['-v'] * (helpers.verbose or 0) +
-                    ['--firewall', str(port)] + subnets_str)
+                    ['--firewall', str(port)])
         argv_tries = [
             ['sudo'] + argvbase,
             ['su', '-c', ' '.join(argvbase)],
@@ -66,6 +66,9 @@ class FirewallClient:
             raise Fatal('%r returned %d' % (self.argv, rv))
 
     def start(self):
+        self.pfile.write('ROUTES\n')
+        for (ip,width) in self.subnets+self.auto_nets:
+            self.pfile.write('%s,%d\n' % (ip, width))
         self.pfile.write('GO\n')
         self.pfile.flush()
         line = self.pfile.readline()
@@ -80,7 +83,7 @@ class FirewallClient:
             raise Fatal('cleanup: %r returned %d' % (self.argv, rv))
 
 
-def _main(listener, fw, use_server, remotename):
+def _main(listener, fw, use_server, remotename, auto_nets):
     handlers = []
     if use_server:
         if helpers.verbose >= 1:
@@ -102,9 +105,22 @@ def _main(listener, fw, use_server, remotename):
             raise Fatal('expected server init string %r; got %r'
                             % (expected, initstring))
 
-    # we definitely want to do this *after* starting ssh, or we might end
-    # up intercepting the ssh connection!
-    fw.start()
+    def onroutes(routestr):
+        if auto_nets:
+            for line in routestr.strip().split('\n'):
+                (ip,width) = line.split(',', 1)
+                fw.auto_nets.append((ip,int(width)))
+
+        # we definitely want to do this *after* starting ssh, or we might end
+        # up intercepting the ssh connection!
+        #
+        # Moreover, now that we have the --auto-nets option, we have to wait
+        # for the server to send us that message anyway.  Even if we haven't
+        # set --auto-nets, we might as well wait for the message first, then
+        # ignore its contents.
+        mux.got_routes = None
+        fw.start()
+    mux.got_routes = onroutes
 
     def onaccept():
         sock,srcip = listener.accept()
@@ -149,7 +165,7 @@ def _main(listener, fw, use_server, remotename):
             mux.check_fullness()
 
 
-def main(listenip, use_server, remotename, subnets):
+def main(listenip, use_server, remotename, auto_nets, subnets):
     debug1('Starting sshuttle proxy.\n')
     listener = socket.socket()
     listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
@@ -179,6 +195,6 @@ def main(listenip, use_server, remotename, subnets):
     fw = FirewallClient(listenip[1], subnets)
     
     try:
-        return _main(listener, fw, use_server, remotename)
+        return _main(listener, fw, use_server, remotename, auto_nets)
     finally:
         fw.done()
diff --git a/firewall.py b/firewall.py
index b4bef1f..8ac5b9a 100644
--- a/firewall.py
+++ b/firewall.py
@@ -140,7 +140,7 @@ def program_exists(name):
 # exit.  In case that fails, it's not the end of the world; future runs will
 # supercede it in the transproxy list, at least, so the leftover rules
 # are hopefully harmless.
-def main(port, subnets):
+def main(port):
     assert(port > 0)
     assert(port <= 65535)
 
@@ -173,8 +173,21 @@ def main(port, subnets):
     line = sys.stdin.readline(128)
     if not line:
         return  # parent died; nothing to do
-    if line != 'GO\n':
-        raise Fatal('firewall: expected GO but got %r' % line)
+
+    subnets = []
+    if line != 'ROUTES\n':
+        raise Fatal('firewall: expected ROUTES but got %r' % line)
+    while 1:
+        line = sys.stdin.readline(128)
+        if not line:
+            raise Fatal('firewall: expected route but got %r' % line)
+        elif line == 'GO\n':
+            break
+        try:
+            (ip,width) = line.strip().split(',', 1)
+        except:
+            raise Fatal('firewall: expected route or GO but got %r' % line)
+        subnets.append((ip, int(width)))
     try:
         if line:
             debug1('firewall manager: starting transproxy.\n')
diff --git a/main.py b/main.py
index 15eb4cc..2811ff7 100755
--- a/main.py
+++ b/main.py
@@ -50,6 +50,7 @@ sshuttle --firewall <port> <subnets...>
 sshuttle --server
 --
 l,listen=  transproxy to this ip address and port number [default=0]
+N,auto-nets automatically determine subnets to route
 r,remote=  ssh hostname (and optional username) of remote sshuttle server
 v,verbose  increase debug message verbosity
 noserver   don't use a separate server process (mostly for debugging)
@@ -65,19 +66,19 @@ try:
     if opt.server:
         sys.exit(server.main())
     elif opt.firewall:
-        if len(extra) < 1:
-            o.fatal('at least one argument expected')
-        sys.exit(firewall.main(int(extra[0]),
-                               parse_subnets(extra[1:])))
+        if len(extra) != 1:
+            o.fatal('exactly one argument expected')
+        sys.exit(firewall.main(int(extra[0])))
     else:
-        if len(extra) < 1:
-            o.fatal('at least one subnet expected')
+        if len(extra) < 1 and not opt.auto_nets:
+            o.fatal('at least one subnet (or -N) expected')
         remotename = opt.remote
         if remotename == '' or remotename == '-':
             remotename = None
         sys.exit(client.main(parse_ipport(opt.listen or '0.0.0.0:0'),
                              not opt.noserver,
                              remotename,
+                             opt.auto_nets,
                              parse_subnets(extra)))
 except Fatal, e:
     log('fatal: %s\n' % e)
diff --git a/server.py b/server.py
index 1da22a1..4b3f55d 100644
--- a/server.py
+++ b/server.py
@@ -1,15 +1,83 @@
-import struct, socket, select
+import re, struct, socket, select, subprocess
 if not globals().get('skip_imports'):
     import ssnet, helpers
     from ssnet import SockWrapper, Handler, Proxy, Mux, MuxWrapper
     from helpers import *
 
 
+def _ipmatch(ipstr):
+    if ipstr == 'default':
+        ipstr = '0.0.0.0/0'
+    m = re.match(r'^(\d+(\.\d+(\.\d+(\.\d+)?)?)?)(?:/(\d+))?$', ipstr)
+    if m:
+        g = m.groups()
+        ips = g[0]
+        width = int(g[4] or 32)
+        if g[1] == None:
+            ips += '.0.0.0'
+            width = min(width, 8)
+        elif g[2] == None:
+            ips += '.0.0'
+            width = min(width, 16)
+        elif g[3] == None:
+            ips += '.0'
+            width = min(width, 24)
+        return (struct.unpack('!I', socket.inet_aton(ips))[0], width)
+
+
+def _ipstr(ip, width):
+    if width >= 32:
+        return ip
+    else:
+        return "%s/%d" % (ip, width)
+
+
+def _maskbits(netmask):
+    if not netmask:
+        return 32
+    for i in range(32):
+        if netmask[0] & (1<<i):
+            return 32-i
+    return 0
+
+
+def _list_routes():
+    argv = ['netstat', '-rn']
+    p = subprocess.Popen(argv, stdout=subprocess.PIPE)
+    routes = []
+    for line in p.stdout:
+        cols = re.split(r'\s+', line)
+        ipw = _ipmatch(cols[0])
+        if not ipw:
+            continue  # some lines won't be parseable; never mind
+        maskw = _ipmatch(cols[2])  # linux only
+        mask = _maskbits(maskw)   # returns 32 if maskw is null
+        width = min(ipw[1], mask)
+        ip = ipw[0] & (((1<<width)-1) << (32-width))
+        routes.append((socket.inet_ntoa(struct.pack('!I', ip)), width))
+    rv = p.wait()
+    if rv != 0:
+        raise Fatal('%r returned %d' % (argv, rv))
+    return routes
+
+
+def list_routes():
+    for (ip,width) in _list_routes():
+        if not ip.startswith('0.') and not ip.startswith('127.'):
+            yield (ip,width)
+        
+
+
 def main():
     if helpers.verbose >= 1:
         helpers.logprefix = ' s: '
     else:
         helpers.logprefix = 'server: '
+
+    routes = list(list_routes())
+    debug1('available routes:\n')
+    for r in routes:
+        debug1('  %s/%d\n' % r)
         
     # synchronization header
     sys.stdout.write('SSHUTTLE0001')
@@ -21,6 +89,9 @@ def main():
               socket.fromfd(sys.stdout.fileno(),
                             socket.AF_INET, socket.SOCK_STREAM))
     handlers.append(mux)
+    routepkt = ''.join('%s,%d\n' % r
+                       for r in routes)
+    mux.send(0, ssnet.CMD_ROUTES, routepkt)
 
     def new_channel(channel, data):
         (dstip,dstport) = data.split(',', 1)
diff --git a/ssnet.py b/ssnet.py
index f13bc9a..782df98 100644
--- a/ssnet.py
+++ b/ssnet.py
@@ -12,6 +12,7 @@ CMD_CONNECT = 0x4203
 CMD_CLOSE = 0x4204
 CMD_EOF = 0x4205
 CMD_DATA = 0x4206
+CMD_ROUTES = 0x4207
 
 cmd_to_name = {
     CMD_EXIT: 'EXIT',
@@ -21,6 +22,7 @@ cmd_to_name = {
     CMD_CLOSE: 'CLOSE',
     CMD_EOF: 'EOF',
     CMD_DATA: 'DATA',
+    CMD_ROUTES: 'ROUTES',
 }
     
 
@@ -220,7 +222,7 @@ class Mux(Handler):
         Handler.__init__(self, [rsock, wsock])
         self.rsock = rsock
         self.wsock = wsock
-        self.new_channel = None
+        self.new_channel = self.got_routes = None
         self.channels = {}
         self.chani = 0
         self.want = 0
@@ -259,12 +261,13 @@ class Mux(Handler):
         p = struct.pack('!ccHHH', 'S', 'S', channel, cmd, len(data)) + data
         self.outbuf.append(p)
         debug2(' > channel=%d cmd=%s len=%d (fullness=%d)\n'
-               % (channel, cmd_to_name[cmd], len(data), self.fullness))
+               % (channel, cmd_to_name.get(cmd,hex(cmd)),
+                  len(data), self.fullness))
         self.fullness += len(data)
 
     def got_packet(self, channel, cmd, data):
         debug2('<  channel=%d cmd=%s len=%d\n' 
-               % (channel, cmd_to_name[cmd], len(data)))
+               % (channel, cmd_to_name.get(cmd,hex(cmd)), len(data)))
         if cmd == CMD_PING:
             self.send(0, CMD_PONG, data)
         elif cmd == CMD_PONG:
@@ -277,6 +280,11 @@ class Mux(Handler):
             assert(not self.channels.get(channel))
             if self.new_channel:
                 self.new_channel(channel, data)
+        elif cmd == CMD_ROUTES:
+            if self.got_routes:
+                self.got_routes(data)
+            else:
+                raise Exception('weird: got CMD_ROUTES without got_routes?')
         else:
             callback = self.channels[channel]
             callback(cmd, data)