a socks server implemented use twisted, support socks4/4a, socks5 protocol, only connect command support

socks-server.py

#!python

"""
 a socks server implemented by twisted
 support socks4/4a, socks5 protocol
 only support CONNECT command 
#
"""

from twisted.internet.protocol import Protocol
from twisted.internet.protocol import Factory, ClientFactory
import socket
import struct

HANDSHAKE = 1
CONNECTING = 2
CONNECTED = 3

CMD_CONNECT = 0x01

CMD_OK5 = 0x00
CMD_OK4 = 0x5a

NO_AUTH = 0x00

SOCKS4 = 0x04
SOCKS5 = 0x05

class SocksProtocol(Protocol):
    def __init__(self):
        self.downstream = None
        self.state = HANDSHAKE

    def dataReceived(self, data):
        if self.state == HANDSHAKE:
            self.handleHandshake(data)
            return

        if self.state == CONNECTING:
            self.handleCmd(data)
            return

        if self.state == CONNECTED:
            self.handleData(data)

    def connectionLost(self, reason):
        if self.downstream and self.downstream.proto:
            self.downstream.proto.transport.abortConnection()

    def handleCmd(self, data):
        proto, cmd = struct.unpack("!BB", data[:2])
        addr=''
        port = 0
        if cmd != CMD_CONNECT:
            print("unsupported command %d", cmd)
            self.transport.abortConnection()
            return

        if proto == SOCKS5:
            _, addrtype = struct.unpack("!BB", data[2:4])

            if addrtype == 0x01: # ipv4
                # 4 bytes
                addr = data[4:4+4]
                addr = socket.inet_ntop(socket.AF_INET, addr) # to string format
                port, = struct.unpack("!H", data[8:10])
            elif addrtype == 0x03: # domain name
                # get length
                l, = struct.unpack("!B", data[4])
                
                # domain name
                addr = data[5:5+l]

                # port
                port, = struct.unpack("!H", data[5+l:5+l+2])
            elif addrtype == 0x04: # ipv6
                # 16 bytes
                addr = data[4:4+16]
                addr = socket.inet_ntop(socket.AF_INET6, addr) # to string format
                port, = struct,unpack("!H", data[20:22])
            else:
                print("unknown addrtype %d", addrtype)
                self.trnasport.abortConnection()
                return
        elif proto == SOCKS4:
            port, ip = struct.unpack("!HI", data[2:8])
            addr = socket.inet_ntoa(struct.pack("!I", ip)) # to string format

            if addr.startswith("0.0.0"): # socks4a

                # skip username
                i = 8
                while True:
                    if data[i] == '\x00':
                        break
                    i += 1

                # point to the first byte of domain name
                i += 1

                # move to the last byte '\x00'
                j = i
                while True:
                    if data[j] == '\x00':
                        break
                    j += 1

                addr = data[i:j]
        else:
            print("unknown version %d", proto)
            self.transport.abortConnection()
            return

        # write ok response to client
        # if server connection failed, abort the client connection 
        #
        if proto == SOCKS5:
            r = struct.pack("!BBBBIH", SOCKS5, CMD_OK5, 0x00, 0x01, 0x00,0x1234)
            self.transport.write(r)
        elif proto == SOCKS4:
            r = struct.pack("!BBHI", 0x00, CMD_OK4, 0x1234, 0x00)
            self.transport.write(r)

        # connecting to the target
        f = clientFactory(self)
        reactor.connectTCP(addr, port, f)

        self.state = CONNECTED
        self.downstream = f

    def handleData(self, data):
        if self.downstream.proto:
            self.downstream.proto.transport.write(data)
        else:
            # server connection not ready, we cache the data
            self.downstream.data += data

    def handleHandshake(self, data):
        proto, = struct.unpack("!B", data[0])
        if proto == SOCKS5:
            authlen, = struct.unpack("!B", data[1])
            self.transport.write(struct.pack("!BB", SOCKS5, NO_AUTH))
            self.state = CONNECTING
        elif proto == SOCKS4:
            self.handleCmd(data)
        else:
            print("unknown protocol version %d", proto)
            self.transport.abortConnection()

class clientFactory(ClientFactory):
    def __init__(self, upstream):
        self.upstream = upstream
        self.proto = None
        self.data = ''

    def clientConnectionLost(self, connector, reason):
        self.upstream.transport.abortConnection()

    def clientConnectionFailed(self, connetor, reason):
        self.upstream.transport.abortConnection()

    def buildProtocol(self, addr):
        p = clientProtocol(self)
        self.proto = p
        return p

class clientProtocol(Protocol):
    def __init__(self, factory):
        self.factory = factory

    def dataReceived(self, data):
        self.factory.upstream.transport.write(data)

    def connectionMade(self):
        # send cached data to server
        if self.factory.data:
            self.transport.write(self.factory.data)
            self.factory.data = ''

class SocksFactory(Factory):
    protocol=SocksProtocol

if __name__ == "__main__":
    from twisted.internet import reactor
    reactor.listenTCP(1080, SocksFactory())
    reactor.run()