Really basic transproxying on localhost.

When regenerating outgoing connections, we set TTL=42 to prevent re-proxying
of requests.  That's a little hacky, but at least it avoids infinite loops.

3 files changed, 159 insertions, 7 deletions

diff --git a/client.py b/client.py
index 34e8b93..e5e8609 100644
--- a/client.py
+++ b/client.py
@@ -1,8 +1,111 @@
-import struct
+import struct, select, errno
 from socket import *
 from helpers import *
 
 
+def _nb_clean(func, *args):
+    try:
+        return func(*args)
+    except error, e:
+        if e.args[0] in (errno.EWOULDBLOCK, errno.EAGAIN):
+            return None
+        raise
+
+
+class SockWrapper:
+    def __init__(self, sock):
+        self.sock = sock
+        self.shut_read = self.shut_write = False
+        self.buf = []
+
+    def noread(self):
+        if not self.shut_read:
+            log('%r: setting noread\n' % self)
+            self.shut_read = True
+            #self.sock.shutdown(SHUT_RD)  # doesn't do anything anyway
+        
+    def nowrite(self):
+        if not self.shut_write:
+            log('%r: setting nowrite\n' % self)
+            self.shut_write = True
+            self.sock.shutdown(SHUT_WR)
+        
+    def write(self, buf):
+        assert(buf)
+        self.sock.setblocking(False)
+        return _nb_clean(self.sock.send, buf)
+
+    def fill(self):
+        if self.shut_read:
+            return
+        self.sock.setblocking(False)
+        rb = _nb_clean(self.sock.recv, 65536)
+        if rb:
+            self.buf.append(rb)
+        if rb == '':  # empty string means EOF; None means nothing available
+            self.noread()
+
+    def maybe_fill(self):
+        if not self.buf:
+            self.fill()
+
+    def copy_to(self, outwrap):
+        if self.buf and self.buf[0]:
+            wrote = outwrap.sock.send(self.buf[0])
+            self.buf[0] = self.buf[0][wrote:]
+        while self.buf and not self.buf[0]:
+            self.buf.pop(0)
+        if not self.buf and self.shut_read:
+            outwrap.nowrite()
+
+
+class Handler:
+    def __init__(self, socks = None, callback = None):
+        self.ok = True
+        self.socks = set(socks or [])
+        if callback:
+            self.callback = callback
+
+    def pre_select(self, r, w, x):
+        r |= self.socks
+
+    def callback(self):
+        log('--no callback defined-- %r\n' % self)
+        (r,w,x) = select.select(self.socks, [], [], 0)
+        for s in r:
+            v = s.recv(4096)
+            if not v:
+                log('--closed-- %r\n' % self)
+                self.socks = set()
+                self.ok = False
+
+
+class Proxy(Handler):
+    def __init__(self, sock1, sock2):
+        Handler.__init__(self, [sock1, sock2])
+        self.wrap1 = SockWrapper(sock1)
+        self.wrap2 = SockWrapper(sock2)
+
+    def pre_select(self, r, w, x):
+        if self.wrap1.buf:
+            w.add(self.wrap2.sock)
+        elif not self.wrap1.shut_read:
+            r.add(self.wrap1.sock)
+        if self.wrap2.buf:
+            w.add(self.wrap1.sock)
+        elif not self.wrap2.shut_read:
+            r.add(self.wrap2.sock)
+
+    def callback(self):
+        self.wrap1.maybe_fill()
+        self.wrap2.maybe_fill()
+        self.wrap1.copy_to(self.wrap2)
+        self.wrap2.copy_to(self.wrap1)
+        if (self.wrap1.shut_read and self.wrap2.shut_read and
+            not self.wrap1.buf and not self.wrap2.buf):
+            self.ok = False
+
+
 def original_dst(sock):
     SO_ORIGINAL_DST = 80
     SOCKADDR_MIN = 16
@@ -20,8 +123,30 @@ def main(remotename, subnets):
     listener.bind(('0.0.0.0',1234))
     listener.listen(10)
     log('Listening on %r.\n' % (listener.getsockname(),))
-    while 1:
-        s,srcip = listener.accept()
-        dstip = original_dst(s)
+
+    handlers = []
+    def onaccept():
+        sock,srcip = listener.accept()
+        dstip = original_dst(sock)
         print 'Incoming connection from %r to %r.' % (srcip,dstip)
-        
+        outsock = socket()
+        outsock.setsockopt(SOL_IP, IP_TTL, 42)
+        outsock.connect(dstip)
+        handlers.append(Proxy(sock, outsock))
+    handlers.append(Handler([listener], onaccept))
+    
+    while 1:
+        r = set()
+        w = set()
+        x = set()
+        handlers = filter(lambda s: s.ok, handlers)
+        for s in handlers:
+            s.pre_select(r,w,x)
+        log('\nWaiting: %d[%d,%d,%d]...\n' 
+            % (len(handlers), len(r), len(w), len(x)))
+        (r,w,x) = select.select(r,w,x)
+        log('r=%r w=%r x=%r\n' % (r,w,x))
+        ready = set(r) | set(w) | set(x)
+        for s in handlers:
+            if s.socks & ready:
+                s.callback()
diff --git a/ipt b/ipt
index 04bb599..6ec37db 100755
--- a/ipt
+++ b/ipt
@@ -24,5 +24,5 @@ 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"
+		--to-ports "$PORT" -m ttl \! --ttl 42
 done
diff --git a/main.py b/main.py
index e038e22..cf34bec 100755
--- a/main.py
+++ b/main.py
@@ -1,2 +1,29 @@
 #!/usr/bin/env python
-import ssh, options
+import sys
+import options, client
+
+optspec = """
+sshuttle [-l [ip:]port] [-r [username@]sshserver] <subnets...>
+--
+l,listen=  transproxy to this ip address and port number [default=0]
+r,remote=  ssh hostname (and optional username) of remote sshuttle server
+server     [internal use only]
+iptables   [internal use only]
+"""
+o = options.Options('sshuttle', optspec)
+(opt, flags, extra) = o.parse(sys.argv[1:])
+
+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)
+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))