tcp2proxy

tcp2proxy.py

# -*- coding: utf-8 -*-
# 2020/3/4
# create by: snower

import sys
import time
import logging
import struct
import socket
import traceback
import sevent

def format_data_len(date_len):
    if date_len < 1024:
        return "%dB" % date_len
    elif date_len < 1024*1024:
        return "%.3fK" % (date_len/1024.0)
    elif date_len < 1024*1024*1024:
        return "%.3fM" % (date_len/(1024.0*1024.0))

def warp_write(conn, status, key):
    origin_write = conn.write
    def _(data):
        status[key] += len(data)
        return origin_write(data)
    return _

def socks5_build_protocol(remote_host, remote_port):
    try:
        protocol_data = b"".join(
            [b"\x05\x01\x00\x01", socket.inet_aton(remote_host), struct.pack(">H", remote_port)])
    except:
        try:
            protocol_data = b"".join([b"\x05\x01\x00\x04", socket.inet_pton(socket.AF_INET6, remote_host),
                                      struct.pack(">H", remote_port)])
        except:
            protocol_data = b"".join([b"\x05\x01\x00\x03", struct.pack(">B", len(remote_host)),
                                      bytes(remote_host, "utf-8") if isinstance(remote_host, str) else remote_host,
                                      struct.pack(">H", remote_port)])
    return protocol_data

def socks5_read_protocol(buffer):
    cmd = buffer.read(1)
    if cmd == b'\x01':
        return buffer.read(6)
    if cmd == b'\x04':
        return buffer.read(18)
    if cmd == b'\x03':
        addr_len = ord(buffer.read(1))
        return buffer.read(addr_len + 2)
    return False

async def socks5_proxy(conn, proxy_host, proxy_port, remote_host, remote_port):
    start_time = time.time()
    status = {"recv_len": 0, "send_len": 0}
    conn.write, pconn = warp_write(conn, status, "recv_len"), None

    try:
        pconn = sevent.tcp.Socket()
        await pconn.connectof((proxy_host, proxy_port))
        await pconn.send(b"\x05\x01\x00")
        buffer = await pconn.recv()
        if buffer.read() != b'\x05\00':
            logging.info("protocol hello error")
            return

        protocol_data = socks5_build_protocol(remote_host, remote_port)
        await pconn.send(protocol_data)
        buffer = await pconn.recv()
        if buffer.read(3) != b'\x05\x00\x00':
            logging.info("protocol error")
            return
        if not socks5_read_protocol(buffer):
            logging.info("protocol error")
            return
        pconn.write = warp_write(pconn, status, "send_len")
        logging.info("socks5 proxy connect %s:%d -> %s:%s -> %s:%s", conn.address[0], conn.address[1], proxy_host, proxy_port,
                     remote_host, remote_port)
        await pconn.linkof(conn)
    except sevent.errors.SocketClosed:
        pass
    except Exception as e:
        logging.info("socks5 proxy error %s:%d -> %s:%s -> %s:%s %s %.2fms\r%s", conn.address[0], conn.address[1], proxy_host, proxy_port,
                     remote_host, remote_port, e, (time.time() - start_time) * 1000, traceback.format_exc())
        return
    finally:
        conn.close()
        if pconn: pconn.close()

    logging.info("socks5 proxy connected %s:%d -> %s:%s -> %s:%s %s %s %.2fms", conn.address[0], conn.address[1], proxy_host, proxy_port,
                 remote_host, remote_port, format_data_len(status["send_len"]), format_data_len(status["recv_len"]), (time.time() - start_time) * 1000)

def http_build_protocol(remote_host, remote_port):
    remote = (bytes(remote_host, "utf-8") if isinstance(remote_host, str) else remote_host) + b':' + bytes(str(remote_port), "utf-8")
    return b"CONNECT " + remote + b" HTTP/1.1\r\nHost: " + remote + b"\r\nUser-Agent: sevent\r\nProxy-Connection: Keep-Alive\r\n\r\n"

async def http_proxy(conn, proxy_host, proxy_port, remote_host, remote_port):
    start_time = time.time()
    status = {"recv_len": 0, "send_len": 0}
    conn.write, pconn = warp_write(conn, status, "recv_len"), None

    try:
        pconn = sevent.tcp.Socket()
        await pconn.connectof((proxy_host, proxy_port))

        protocol_data = http_build_protocol(remote_host, remote_port)
        await pconn.send(protocol_data)
        buffer = await pconn.recv()
        if buffer.read(12).lower() != b"http/1.1 200":
            logging.info("protocol error")
            return
        buffer.read()
        pconn.write = warp_write(pconn, status, "send_len")
        logging.info("http proxy connect %s:%d -> %s:%s -> %s:%s", conn.address[0], conn.address[1], proxy_host, proxy_port,
                     remote_host, remote_port)
        await pconn.linkof(conn)
    except sevent.errors.SocketClosed:
        pass
    except Exception as e:
        logging.info("http proxy error %s:%d -> %s:%s -> %s:%s %s %.2fms\r%s", conn.address[0], conn.address[1], proxy_host, proxy_port,
                     remote_host, remote_port, e, (time.time() - start_time) * 1000, traceback.format_exc())
        return
    finally:
        conn.close()
        if pconn: pconn.close()

    logging.info("http proxy connected %s:%d -> %s:%s -> %s:%s %s %s %.2fms", conn.address[0], conn.address[1], proxy_host, proxy_port,
                 remote_host, remote_port, format_data_len(status["send_len"]), format_data_len(status["recv_len"]),
                 (time.time() - start_time) * 1000)


async def tcp_accept(server):
    proxy_url = sys.argv[2]
    proxy_info = proxy_url.strip().split("://")
    if len(proxy_info) == 1:
        protocol = "http"
        proxy_info = proxy_info[0].split(":")
        proxy_host, proxy_port = proxy_info[0], int(proxy_info[1]) if len(proxy_info) >= 2 else 8088
    else:
        protocol = proxy_info[0]
        proxy_info = proxy_info[1].split(":")
        proxy_host, proxy_port = proxy_info[0], int(proxy_info[1]) if len(proxy_info) >= 2 else 8088

    logging.info("starting server at port %s for %s proxy %s:%s ...", sys.argv[1], protocol, proxy_host, proxy_port)
    while True:
        conn = await server.accept()
        if protocol == "http":
            sevent.current().call_async(http_proxy, conn, proxy_host, proxy_port, sys.argv[3], int(sys.argv[4]))
        else:
            sevent.current().call_async(socks5_proxy, conn, proxy_host, proxy_port, sys.argv[3], int(sys.argv[4]))

if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)1.1s %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S', filemode='a+')
    server = sevent.tcp.Server()
    server.enable_reuseaddr()
    server.listen(("0.0.0.0", int(sys.argv[1])))
    sevent.run(tcp_accept, server)