Python 作业#2

作业题目

实现localProxy双协议（SOCKS5和HTTP tunnel）本地代理。

支持（SOCKS5代理）基于#1作业的成果。

支持HTTP tunnel（ 即HTTP CONNECT method）可用于HTTPS代理。

关于HTTP tunnel可以参见：https://www.zhihu.com/question/21955083

import asyncio
import getopt
import sys
from struct import unpack, pack


async def msg_send(reader, writer, host):
    while reader.at_eof:
        try:
            data = await reader.read(1024 * 64)
            if not data:
                writer.close()
                break
        except (ConnectionAbortedError, ConnectionResetError) as e:
            writer.close()
            print(f"{host} exit exceptionally {repr(e)}")
            break

        try:
            writer.write(data)
            await writer.drain()
        except (ConnectionAbortedError, ConnectionResetError) as e:
            writer.close()
            print(f"{host} exit exceptionally {repr(e)}")
            break

    print(f"{host} exit")


async def handle_socks5(reader, writer):
    data = await reader.read(1024 * 64)
    addr = writer.get_extra_info('peername')
    print(f"connect from {addr!r}")

    # message at least 3 bytes
    if len(data) < 3:
        print('message is too short')
        writer.close()

    writer.write(b'\x05\x00')
    await writer.drain()
    data = await reader.read(1024 * 64)
    info = unpack('!4B', data[:4])
    ver, cmd, atyp = info[0], info[1], info[3]

    # domain name
    if ver == 5 and cmd == 1 and atyp == 3:

        addr_len = unpack('!B', data[4:5])[0]
        dst_addr = data[5:addr_len + 5].decode()
        dst_port = unpack('!H', data[addr_len + 5:])[0]
        print(f'destination domain name {dst_addr}，destination port {dst_port}')

        try:
            reader_remote, writer_remote = await asyncio.open_connection(dst_addr, dst_port)
            writer.write(pack('!5B', 5, 0, 0, 3, addr_len) + dst_addr.encode() + pack('!H', dst_port))
            await writer.drain()
            print(f'connect to {dst_addr} success ')
        except (TimeoutError, ConnectionRefusedError) as e:
            print(f'connect to {dst_addr} failed ')
            writer.write(pack('!5B', 5, 3, 0, 3, addr_len) + dst_addr.encode() + pack('!H', dst_port))
            await writer.drain()
            writer.close()
            return

        client_to_remote = msg_send(reader, writer_remote, dst_addr)
        remote_to_client = msg_send(reader_remote, writer, dst_addr)
        await asyncio.gather(client_to_remote, remote_to_client)

    # ipv4
    elif ver == 5 and cmd == 1 and atyp == 1:

        dst_address = '.'.join([str(a) for a in unpack('!BBBB', data[4:8])])
        dst_port = unpack('H', data[8:10])[0]
        print(f'destination address {dst_address}，destination port {dst_port}')

        try:
            reader_remote, writer_remote = await asyncio.open_connection(dst_address, dst_port)
            writer.write(pack('!8B', 5, 0, 0, 1, *unpack('!BBBB', data[4:8])) + pack('!H', dst_port))
            await writer.drain()
            print(f'connect to {dst_addr} success ')
        except (TimeoutError, ConnectionRefusedError) as e:
            print(f'connect to {dst_addr} failed ')
            writer.write(pack('!8B', 5, 3, 0, 1, *unpack('!BBBB', data[4:8])) + pack('!H', dst_port))
            await writer.drain()
            writer.close()
            return

        client_to_remote = msg_send(reader, writer_remote, dst_address)
        remote_to_client = msg_send(reader_remote, writer, dst_address)
        await asyncio.gather(client_to_remote, remote_to_client)

    # ipv6
    elif ver == 5 and cmd == 1 and atyp == 4:
        pass

    else:

        print('unsuppoted request')
        writer.close()
        return


async def handle_http(reader, writer):
    data = await reader.read(1024 * 64)
    addr = writer.get_extra_info('peername')
    print(f"connect from {addr!r}")
    print(data)

    info = data.decode().split("\r\n")
    request = info[0].split(" ")

    method = request[0]
    host = request[1].split(":")[0]
    port = request[1].split(":")[1]
    ver = request[2]

    print(f'destination domain name {host}，destination port {port}')

    if method == 'CONNECT':

        try:
            reader_remote, writer_remote = await asyncio.open_connection(host, port)
            reply = ver + " 200 Connection established\r\n"
            reply += "Proxy-agent: wyx\r\n\r\n"
            writer.write(reply.encode())
            await writer.drain()
            print(f'connect to {host} success ')
        except (TimeoutError, ConnectionRefusedError) as e:
            print(f'connect to {host} failed ')
            writer.close()
            return

        client_to_remote = msg_send(reader, writer_remote, host)
        remote_to_client = msg_send(reader_remote, writer, host)
        await asyncio.gather(client_to_remote, remote_to_client)



async def main_logic(mode, port):
    if(mode.lower() == "socks5"):
        server = await asyncio.start_server(handle_socks5, '127.0.0.1', port)
    else:
        server = await asyncio.start_server(handle_http, '127.0.0.1', port)

    addr = server.sockets[0].getsockname()
    print(f'Serving on {addr}')

    async with server:
        await server.serve_forever()


def main():
    try:
        opts, args = getopt.getopt(sys.argv[1:], "hm:p:v", ["help", "mode=", "port="])
    except getopt.GetoptError as e:
        print(str(e))
        sys.exit(2)
    mode = "socks5"
    port = "8765"
    verbose = False

    for key, value in opts:
        if key == "-v":
            verbose = True
        elif key in ("-h", "--help"):
            print("localProxy.py -m <mode> -p <port>")
            print("localProxy.py --mode <mode> --port <port>")

        elif key in ("-m", "--mode"):
            mode = value
        elif key in ("-p", "--port"):
            port = value

    if(mode.lower() != "socks5" and mode.lower() != "https"):
        print("wrong mode option")

    loop = asyncio.get_event_loop()
    loop.run_until_complete(main_logic(mode, port))

if __name__ == "__main__":
    main()