socks5server

#!/usr/bin/env python

from socket import *
import select
import struct
import errno


streams = {} # fileno -> Socks5Handler


class TCPStream(object):
    def __init__(self, sock, poll, event=None):
        self.sock = sock
        self.poll = poll
        self.recved_data = b''
        self.to_write_data = b''
        self.event = event
        self.read_callback = None # (size, call_back)
        self.write_callback = None # callback
        self.error_callback = None

    def process(self, evt):
        try:
            if evt & select.EPOLLIN:
                self._on_read()
            if evt & select.EPOLLOUT:
                self._on_write()
            if evt & select.EPOLLERR:
                if self.error_callback is not None:
                    self.error_callback()
                self.close()
        except OSError as e:
            print('error:', e)
            self.close()

    def close(self):
        if self.sock.fileno() > 0:
            self.poll.unregister(self.sock.fileno())
            streams.pop(self.sock.fileno())
            self.sock.close()

    def _add_event(self, evt):
        if self.event is None:
            self.event = evt | select.EPOLLERR
            self.poll.register(self.sock.fileno(), self.event)
        elif self.event & evt == 0:
            self.event |= evt
            self.poll.modify(self.sock.fileno(), self.event)

    def _rm_event(self, event):
        if self.event is None:
            return
        if self.event & event:
            self.event &= ~event
            self.poll.modify(self.sock.fileno(), self.event)

    def _on_read(self):
        try:
            # while读到错误也可以，为了简单我就只读一次
            data = self.sock.recv(65535)
            if not data:
                return self.close()
        except OSError as e:
            if e.errno != errno.EAGAIN:
                raise
        self.recved_data += data

        if self.read_callback is not None:
            size, callback = self.read_callback
            if size == 0:
                data = self.recved_data
                self.recved_data = b''
                self.read_callback = None
                callback(data)
            elif len(self.recved_data) >= size:
                data = self.recved_data[:size]
                self.recved_data = self.recved_data[size:]
                self.read_callback = None
                callback(data)

    def _on_write(self):
        write_callback = self.write_callback
        self._rm_event(select.EPOLLOUT)
        if self.to_write_data:
            try:
                n = self.sock.send(self.to_write_data)
                self.to_write_data = self.to_write_data[n:]
            except OSError as e:
                if e.errno != errno.EAGAIN:
                    raise
                return
            if not self.to_write_data:
                if write_callback is not None:
                    self.write_callback = None
                    write_callback()
        else:
            if write_callback is not None:
                self.write_callback = None
                write_callback()

    def read_n_bytes(self, n, callback):
        if len(self.recved_data) >= n:
            data = self.recved_data[:n]
            self.recved_data = self.recved_data[n:]
            callback(data)
        else:
            self.read_callback = (n, callback)
            self._add_event(select.EPOLLIN)

    def add_err_callback(self, cb):
        self.error_callback = cb

    def write_data(self, data, callback):
        if data:
            try:
                n = self.sock.send(data)
                self.to_write_data = data[n:]
            except OSError as e:
                if e.errno == errno.EBADF:
                    # socket closed
                    return
                elif e.errno != errno.EAGAIN:
                    raise
                self.to_write_data = data
            else:
                if n == len(data):
                    callback()
        if data is None or self.to_write_data:
            self.write_callback = callback
            self._add_event(select.EPOLLOUT)


    def read_all(self, callback):
        if len(self.recved_data) > 0:
            data = self.recved_data
            self.recved_data = b''
            callback(data)
        else:
            self.read_callback = 0, callback
            self._add_event(select.EPOLLIN)


class Socks5Handler(TCPStream):
    def __init__(self, sock, addr, poll):
        self.sock = sock
        self.addr = addr
        super(Socks5Handler, self).__init__(sock, poll)
        self.read_n_bytes(2, self.on_auth_req)

    def on_auth_req(self, data):
        if data[0] != 5:
            raise Exception('invalid socks5 version')
        if data[1] == 0:
            raise Exception('invalid socks5 request')
        self.read_n_bytes(data[1], self.do_auth_response)

    def do_auth_response(self, data):
        if 0 not in data:
            print('auth fail')
            raise Exception('noauth only!')
        self.write_data(
            b'\x05\x00', lambda: self.read_n_bytes(4, self.on_sock_req))

    def on_sock_req(self, data):
        if data[0] != 5:
            raise Exception('invalid socks5 version')
        if data[1] != 1:
            self.write_data(b'\x05\x07\x00\x01\x00\x00\x00\x00\x00\x00',
                self.close)
        if data[3] == 1:
            self.read_n_bytes(6, self.on_ipv4_req)
        elif data[3] == 3:
            self.read_n_bytes(1,
                lambda data: self.read_n_bytes(data[0] + 2, self.on_domain_req))
        else:
            self.write_data(b'\x05\x08\x00\x01\x00\x00\x00\x00\x00\x00',
                self.close)

    def on_ipv4_req(self, data):
        ip = '%d.%d.%d.%d' % data[:4]
        port = struct.unpack('>H', data[4:6])[0]
        self.add_stream((ip, port))

    def on_domain_req(self, data):
        doamin = data[:-2]
        port = struct.unpack('>H', data[-2:])[0]
        self.add_stream((doamin, port))

    def add_stream(self, addr):
        remote_sock = socket(AF_INET, SOCK_STREAM)
        remote_sock.setblocking(0)
        self.remote_stream = TCPStream(remote_sock, self.poll)
        streams[remote_sock.fileno()] = self.remote_stream
        try:
            remote_sock.connect(addr)
            self.on_connect_remote_ok()
        except OSError as e:
            if e.errno != errno.EINPROGRESS:
                print('  connect remote err:', e)
                return self.write_data(
                    b'\x05\x04\x00\x01\x00\x00\x00\x00\x00\x00', self.close)
            else:
                self.remote_stream.write_data(None, self.on_connect_remote_ok)

    def on_connect_remote_error(self):
        self.write_data(
            b'\x05\x04\x00\x01\x00\x00\x00\x00\x00\x00', self.close)
        print('on_connect_remote_error')

    def on_connect_remote_ok(self):
        self.write_data(b'\x05\x00\x00\x01\x00\x00\x00\x00\x00\x00',
                self.start_copy)

    def start_copy(self):
        self.new_copy(self, self.remote_stream)
        self.new_copy(self.remote_stream, self)

    def new_copy(self, f, t):
        # f.on_read f.on_write
        def on_read(data):
            t.write_data(data, lambda: f.read_all(on_read))
        f.read_all(on_read)


def main(ip, port):
    sock = socket(AF_INET, SOCK_STREAM)
    sock.setblocking(0)
    sock.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
    sock.bind((ip, port))
    sock.listen(64)

    poll = select.epoll()
    poll.register(sock.fileno(), select.EPOLLIN)
    while True:
        for fileno, event in poll.poll():
            if fileno == sock.fileno():
                cli, addr = sock.accept()
                cli.setblocking(0)
                streams[cli.fileno()] = Socks5Handler(cli, addr, poll)
            else:
                stream = streams.get(fileno, None)
                if stream is not None:
                    stream.process(event)
                else:
                    # never happen
                    print('wtf?')
                    poll.unregister(fileno)


if __name__ == '__main__':
    main('127.0.0.1', 1081)