Simple SOCKS5 server on Python

socks_proxy.py

import sys, os, select, socket

from struct import pack, unpack
from socketserver import ThreadingMixIn, TCPServer, StreamRequestHandler
from optparse import OptionParser, make_option

from config import Conf

BUFF_SIZE = 0x1000

class ThreadingTCPServer(ThreadingMixIn, TCPServer):

    log_file = None    
    allow_reuse_address = True

class SocksProxy(StreamRequestHandler):

    SOCKS_VERSION = 5

    SOCKS_E_SUCCESS = 0
    SOCKS_E_CONN_ERROR = 4
    SOCKS_E_BAD_COMMAND = 7
    SOCKS_E_BAD_ADDRESS = 8

    SOCKS_ADDR_IPv4 = 1
    SOCKS_ADDR_DOMAIN = 3

    def receive(self, size):

        ret = ''

        while len(ret) < size:

            # receive needed amount of data
            data = self.connection.recv(size - len(ret))
            if len(data) <= 0:

                raise(Exception('Client disconnected'))

            ret += data

        return ret

    def send(self, data):

        self.connection.sendall(data)

    def send_reply(self, command, addr = 0, port = 0):

        self.send(pack('!BBBBIH', self.SOCKS_VERSION, command, 0, self.SOCKS_ADDR_IPv4, addr, port))

    def log(self, message):

        message = '[%s:%d] %s' % (self.client_address[0], self.client_address[1], message)

        print(message)

        if self.server.log_file is not None:

            # write message into the log file
            self.server.log_file.write(message + '\n')
            self.server.log_file.flush()

    def handle(self):        

        self.log('Incoming connection accepted')

        try:

            # receive method request
            method_request = self.receive(1 + 1)
            version, method_num = unpack('BB', method_request)

            if version != self.SOCKS_VERSION:

                raise(Exception('Unsupported protocol version (%d)' % version))
            
            if method_num > 0:

                # receive methods list
                methods = self.receive(method_num)

            # send method reply
            self.send(pack('BB', self.SOCKS_VERSION, 0))

            # receive request
            request = self.receive(1 + 1 + 1 + 1)
            version, command, _, addr_type = unpack('BBBB', request)

            if command == 1:

                if addr_type in [ self.SOCKS_ADDR_IPv4, self.SOCKS_ADDR_DOMAIN ]:

                    success = False

                    if addr_type == self.SOCKS_ADDR_IPv4:

                        # reveice address and port
                        addr, port = unpack('!IH', self.receive(4 + 2))

                        # convert IP address to string
                        addr = socket.inet_ntoa(addr)

                        self.log('Connecting to the %s:%d' % (addr, port))

                        success = True

                    elif addr_type == self.SOCKS_ADDR_DOMAIN:

                        # receive domain name
                        addr_len, = unpack('B', self.receive(1))
                        addr = self.receive(addr_len)

                        # receive port
                        port, = unpack('!H', self.receive(2))

                        self.log('Connecting to the %s:%d' % (addr, port))

                        try:

                            addr = socket.gethostbyname(addr)
                            success = True

                        except socket.gaierror:   

                            self.log('ERROR: Unable to resolve hostname')    

                            # unable to resolve hostname
                            self.send_reply(self.SOCKS_E_CONN_ERROR)               

                    if success:

                        success = False

                        try:

                            # connect to the remote host
                            remote = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                            remote.connect(( addr, port ))

                            bind_addr, bind_port = remote.getsockname()
                            bind_addr, = unpack('!I', socket.inet_aton(bind_addr))

                            # send reply
                            self.send_reply(self.SOCKS_E_SUCCESS, bind_addr, bind_port)

                            success = True

                        except socket.error:

                            self.log('ERROR: Unable to connect to the remote server')

                            # unable to connect to the remote server
                            self.send_reply(self.SOCKS_E_CONN_ERROR)                        

                        if success:

                            self.log('Connection established')

                            # transfer data between client and remote host
                            self.exchange_loop(self.connection, remote)

                else:

                    # bad address type
                    self.send_reply(self.SOCKS_E_BAD_ADDRESS)

            else:

                # bad command
                self.send_reply(self.SOCKS_E_BAD_COMMAND)

        except Exception, why:

            self.log('ERROR: %s' % str(why))

        # disconnect client
        self.server.close_request(self.request)

        self.log('Connection closed')

    def exchange_loop(self, client, remote):

        while True:

            # wait until client or remote is available for read
            r, w, e = select.select([ client, remote ], [], [])

            if client in r:

                data = client.recv(BUFF_SIZE)
                if len(data) == 0:

                    break

                if remote.send(data) <= 0:

                    break

            if remote in r:

                data = remote.recv(BUFF_SIZE)
                if len(data) == 0:

                    break

                if client.send(data) <= 0:

                    break

def main():

    option_list = [

        make_option('-a', '--address', dest = 'address', default = Conf.SOCKS_ADDR,
            help = 'address to listen on'),

        make_option('-p', '--port', dest = 'port', default = Conf.SOCKS_PORT,
            help = 'TCP port to listen on'),

        make_option('-l', '--log-file', dest = 'log_file', default = None,
            help = 'port mapping source port') ] 

    options, _ = OptionParser(option_list = option_list).parse_args()        

    addr = ( options.address, options.port )

    print('[+] Starting SOCKS server at %s:%d' % addr)

    # create server instance
    server = ThreadingTCPServer(addr, SocksProxy)

    # open log file
    log_file = options.log_file if options.log_file is not None else None
    server.log_file = log_file if log_file is None else open(log_file, 'wb')

    try:

        # run server
        server.serve_forever()

    except KeyboardInterrupt:

        print('\nEXIT\n')

    # cleanup
    server.shutdown()

    return 0

if __name__ == '__main__':
    
    exit(main())

#
# EoF
#