medium-socks5/main.py

try:
    import gevent
    import gevent.monkey
    gevent.monkey.patch_all(dns=gevent.version_info[0] >= 1)
except ImportError:
    print("Warning: gevent not found, fallback to traditional implementation.")

import socket
import socketserver
import struct
import select
from sys import argv
import chnroutes_data


DEFAULT_LISTEN_HOST = "127.0.0.1"
DEFAULT_LISTEN_PORT = 8964

DEFAULT_UPSTREAM_HOST = "127.0.0.1"
DEFAULT_UPSTREAM_PORT = 1984

SOCKS5_VER = 5
SOCKS5_METHOD = 0
SOCKS5_CMD_CONNECT = 1
SOCKS5_CMD_BIND = 2
SOCKS5_CMD_UDP_ASSOCIATE = 3
SOCKS5_ATYP_IPV4 = 1
SOCKS5_ATYP_DOMAINNAME = 3
SOCKS5_ATYP_IPV6 = 4
SOCKS5_REP_SUCCESS = 0
SOCKS5_RESERVED = 0


_dns = {}


def dns(addr):
    ip = _dns.get(addr, "")
    if not ip:
        ip = socket.getaddrinfo(addr, None)[0][4][0]
        _dns[addr] = ip
    return ip


class UpstreamProxyServer():

    def __init__(self, addr, port):
        self._sock = socket.create_connection((addr, port))

    def connect(self, target):
        hello = bytes((SOCKS5_VER, 1, SOCKS5_METHOD))
        self._sock.sendall(hello)
        ver, method = self._sock.recv(2)

        connect = bytes((SOCKS5_VER, SOCKS5_CMD_CONNECT, SOCKS5_RESERVED,
                         target.atyp)) + target.raw_addr + target.raw_port
        self._sock.sendall(connect)

        reply = self._sock.recv(4 + 4 + 2)
        assert reply[1] == 0

    def recv(self, buf):
        return self._sock.recv(buf)

    def send(self, buf):
        return self._sock.send(buf)

    def close(self):
        return self._sock.close()

    def fileno(self):
        return self._sock.fileno()


class DirectConnectProxyServer(UpstreamProxyServer):

    def __init__(self, addr, port):
        pass

    def connect(self, target):
        self._sock = socket.create_connection((target.addr, target.port))


class Target():

    def __init__(self, atyp, addr):
        self.atyp = atyp

        if self.atyp == SOCKS5_ATYP_IPV4:
            self.length = 4
            self.raw_addr = addr[:self.length]
            self.addr = socket.inet_ntoa(self.raw_addr)
            self.raw_port = addr[self.length:self.length + 2]
        elif self.atyp == SOCKS5_ATYP_IPV6:
            self.length = 16
            self.raw_addr = addr[:self.length]
            self.addr = socket.inet_ntop(socket.AF_INET6, self.raw_addr)
            self.raw_port = addr[self.length:self.length + 2]
        elif self.atyp == SOCKS5_ATYP_DOMAINNAME:
            self.length = addr[0]
            self.raw_addr = addr[0:1 + self.length]
            self.addr = self.raw_addr[1:].decode("UTF-8")
            self.raw_port = addr[1 + self.length:1 + self.length + 2]

        self.port = struct.unpack("!H", self.raw_port)[0]


class ThreadingTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    allow_reuse_address = True


class MySocksServer(socketserver.StreamRequestHandler):

    @staticmethod
    def _is_private(target):
        if target.atyp == SOCKS5_ATYP_DOMAINNAME:
            try:
                ip = dns(target.addr)
            except socket.gaierror:
                return False
        else:
            ip = target.addr

        f = struct.unpack('!I', socket.inet_pton(socket.AF_INET, ip))[0]
        private = (
            [2130706432, 4278190080],  # 127.0.0.0,   255.0.0.0   http://tools.ietf.org/html/rfc3330
            [3232235520, 4294901760],  # 192.168.0.0, 255.255.0.0 http://tools.ietf.org/html/rfc1918
            [2886729728, 4293918720],  # 172.16.0.0,  255.240.0.0 http://tools.ietf.org/html/rfc1918
            [167772160, 4278190080],   # 10.0.0.0,    255.0.0.0   http://tools.ietf.org/html/rfc1918
        )
        for net in private:
            if (f & net[1] == net[0]):
                return True
        return False

    @staticmethod
    def _in_prc(target):
        def in_subnet(ip, subnet, netmask):
            return (ip & netmask) == (subnet & netmask)

        def str_ip_to_num(ip):
            return struct.unpack('!I', socket.inet_aton(ip))[0]

        if target.atyp == SOCKS5_ATYP_DOMAINNAME:
            try:
                ip = str_ip_to_num(dns(target.addr))
            except socket.gaierror:
                return False
        else:
            ip = struct.unpack("!I", target.raw_addr)[0]

        for i in chnroutes_data.CHINESE_SUBNETS:
            subnet = str_ip_to_num(i[0])
            netmask = str_ip_to_num(i[1])
            if in_subnet(ip, subnet, netmask):
                return True
        return False

    @staticmethod
    def _sendall(sock, data):
        sent_bytes = 0

        while sent_bytes < len(data):
            sent = sock.send(data[sent_bytes:])
            if sent < 0:
                return sent
            sent_bytes += sent
        return sent_bytes

    def handle_proxy_tcp(self, target):
        connect_msg = "Connect to %s:%d %s"
        connect_fmt = [target.addr, target.port]

        if self._in_prc(target) or self._is_private(target):
            proxy_server = DirectConnectProxyServer
            connect_fmt.append("(PRC)")
        else:
            proxy_server = UpstreamProxyServer
            connect_fmt.append("(Not in PRC)")

        print(connect_msg % tuple(connect_fmt))

        proxy = proxy_server(DEFAULT_UPSTREAM_HOST, DEFAULT_UPSTREAM_PORT)
        proxy.connect(target)

        fdset = (self.request, proxy)

        try:
            while True:
                r, w, e = select.select(fdset, (), ())

                if self.request in r:
                    data = self.request.recv(4096)
                    if len(data) <= 0:
                        break
                    sent = self._sendall(proxy, data)
                    if sent < len(data):
                        raise IOError("Failed to sent all data.")

                if proxy in r:
                    data = proxy.recv(4096)
                    if len(data) <= 0:
                        break
                    sent = self._sendall(self.request, data)
                    if sent < len(data):
                        raise IOError("Failed to sent all data.")
        except Exception as e:
            print(e)
        finally:
            self.request.close()
            proxy.close()
            pass

    def handle(self):
        ver, nmethods, methods = self.request.recv(3)
        if ver != SOCKS5_VER:
            self.request.close()
            return

        self.request.sendall(bytes((SOCKS5_VER, SOCKS5_METHOD)))

        ver, cmd, rsv, atyp = self.request.recv(4)
        if ver != SOCKS5_VER:
            self.request.close()
            return

        if cmd == SOCKS5_CMD_CONNECT:
            pass
        elif cmd == SOCKS5_CMD_BIND:
            self.request.close()
            return
        elif cmd == SOCKS5_CMD_UDP_ASSOCIATE:
            self.request.close()
            return

        if atyp not in (SOCKS5_ATYP_IPV4, SOCKS5_ATYP_IPV6, SOCKS5_ATYP_DOMAINNAME):
            self.request.close()
            return

        dst_addr = self.request.recv(256)  # 1 (domain length) + 253 (max domain length) + 2 (port)
        target = Target(atyp, dst_addr)

        reply = bytes((SOCKS5_VER, SOCKS5_REP_SUCCESS, SOCKS5_RESERVED, SOCKS5_ATYP_IPV4))
        reply += socket.inet_aton('0.0.0.0') + struct.pack("!H", 0)
        self.request.sendall(reply)

        self.handle_proxy_tcp(target)


if __name__ == "__main__":
    host = DEFAULT_LISTEN_HOST
    port = DEFAULT_LISTEN_PORT
    if len(argv) == 1:
        pass
    elif len(argv) == 2:
        port = int(argv[1])
    elif len(argv) == 3:
        host = argv[1]
        port = int(argv[2])

    server = ThreadingTCPServer((host, port), MySocksServer)
    server.serve_forever()