iptables to socks forwarder

ipt2socks.py

#!/usr/bin/python


from __future__ import print_function

import os,sys
import asyncore
import socket
import struct

debug=os.environ.get("DEBUG")
quiet=os.environ.get("QUIET")

def get_orig_dst(sock):
    SO_ORIGINAL_DST=getattr(socket, "SO_ORIGINAL_DST", 80)
    try: port,host=struct.unpack(">2xH4s8x", sock.getsockopt(socket.SOL_IP, SO_ORIGINAL_DST, 16))
    except socket.error: return None
    host=socket.inet_ntoa(host)
    return host,port

class async_bufsend(asyncore.dispatcher):
    ignore_log_types={}
    buffer=None
    close_after_send=False
    def log(self, message):
        return asyncore.dispatcher.log(self, "%s %x %s"%(self.__class__.__name__, id(self), message))
    def close(self):
        if debug or not quiet:
            self.log("close: %s"%("later" if self.buffer else "now"))
        if not self.buffer:
            asyncore.dispatcher.close(self)
        else: self.close_after_send=True
    def handle_write(self):
        data=self.buffer.pop(0)
        num_sent=asyncore.dispatcher.send(self, data)
        if num_sent<len(data): self.buffer.insert(0, data[num_sent:])
        if self.close_after_send and not self.buffer: self.close()
    def writable(self):
        return True if self.buffer else False
    def send(self, data):
        if self.buffer is None: self.buffer=[]
        if debug: self.log("send: %r"%(data,))
        self.buffer.append(data)
        return len(data)

class Server(asyncore.dispatcher):
    ignore_log_types={}
    def log(self, message):
        return asyncore.dispatcher.log(self, "%s %x %s"%(self.__class__.__name__, id(self), message))
    def __init__(self, addr, **attrs):
        for k in attrs: setattr(self, k, attrs[k])
        asyncore.dispatcher.__init__(self)
        self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
        self.set_reuse_addr()
        self.bind(addr)
        if debug or not quiet: self.log("bound to: %r"%(self.socket.getsockname(),))
        self.listen(5)
    def handle_accept(self):
        clsock,claddr=self.accept()
        if debug or not quiet: self.log("Accepted from: %r"%(claddr,))
        handler=ClientHandler(clsock, client_addr=claddr)  # @UnusedVariable

class ClientHandler(async_bufsend):
    forwarder=None
    socks_addr=(os.environ.get("SOCKS_HOST","127.0.0.1"), int(os.environ.get("SOCKS_PORT", "1080")))
    def __init__(self, sock, **attrs):
        for k in attrs: setattr(self, k, attrs[k])
        async_bufsend.__init__(self, sock)
        orig_dst=get_orig_dst(sock)
        self.forwarder=SocksHandler(orig_dst, self.socks_addr, forwarder=self)
    def handle_close(self):
        if debug: self.log("closing")
        if self.forwarder: self.forwarder.close()
    def handle_read(self):
        data=self.recv(8192)
        if debug: self.log("read: %r"%(data,))
        if not data:
            if self.forwarder: self.forwarder.close()
            return self.close()
        if self.forwarder:
            self.forwarder.send(data)

class SocksHandler(ClientHandler):
    socks_connected=False
    auth_resp=None
    conn_resp=None
    queue_buffer=None
    def __init__(self, remote_addr, socks_addr, **attrs):
        for k in attrs: setattr(self, k, attrs[k])
        self.remote_host,self.remote_port=remote_addr
        async_bufsend.__init__(self)
        self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socks_addr=socks_addr
        self.connect(socks_addr)
        self.send(b"\x05\x02\x00\x02", init_seq=True)
    def handle_read(self):
        if self.socks_connected: return ClientHandler.handle_read(self)
        else: return self.initiate_socks()
    def send(self, data, init_seq=False):
        if self.socks_connected or init_seq: return ClientHandler.send(self, data)
        else:
            if self.queue_buffer is None: self.queue_buffer=[]
            self.queue_buffer.append(data)
            return len(data)
    def handle_connect(self): pass
    def initiate_socks(self):
        if self.auth_resp is None:
            self.auth_resp=self.recv(2)
            if self.auth_resp!=b"\x05\x00":
                self.buffer=None
                self.close()
                raise ValueError("Socks server didn't accept null auth: %r"%(self.auth_resp,))
            try: addr=b"\x01"+struct.pack("4B",*map(int,self.remote_host.split(".")))
            except (ValueError,struct.error):
                addr=struct.pack("BB", 3, len(self.remote_host))+self.remote_host
            if debug or not quiet: self.log("connecting to %s:%d"%(self.remote_host,self.remote_port))
            self.send(b"\x05\x01\x00"+addr+struct.pack(">H", self.remote_port), init_seq=True)
        elif self.conn_resp is None:
            self.conn_resp=self.recv(10)
            if len(self.conn_resp)==10 and self.conn_resp[:4]==b"\x05\x00\x00\x01":
                d=struct.unpack(">4BH",self.conn_resp[4:])
                if debug or not quiet: self.log("connected to %s:%s%s"%(self.remote_host, self.remote_port, " as %s:%s"%(".".join(map(str,d[:4])),d[4]) if any(d) or d[4] else ""))
            else:
                self.buffer=None
                self.close()
                raise ValueError("failed connect to %s:%d"%(self.remote_host, self.remote_port),self.conn_resp)
            self.socks_connected=True
            if self.queue_buffer:
                self.buffer=self.queue_buffer
        else:
            print("Connect() response: %r"%(self.conn_resp))
            raise RuntimeError("Don't know how I ended up here...")

if __name__ == '__main__':
    srv=Server((sys.argv[2] if len(sys.argv)>2 else os.environ.get("LISTEN_ADDR","127.0.0.1"),
        int(sys.argv[1] if len(sys.argv)>1 else os.environ.get("LISTEN_PORT", 0))))
    print("iptables -t nat -F via_socks && iptables -t nat -A via_socks -p tcp -j REDIRECT --to-ports %d"%srv.socket.getsockname()[1])
    try: asyncore.loop()
    except KeyboardInterrupt:
        print("Interrupted.")
        print("iptables -t nat -F via_socks")