Python 作业#1

作业题目

使用asyncio的streams（coroutine based API）实现SOCKS5服务器。

协议参考：RFC 1928 - SOCKS Protocol Verison 5

只需要实现CMD X‘01’（即CONNECT）

只需要实现METHOD X‘00’（即NO AUTHENTICATION REQUIRED）

import asyncio
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(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)

    else:

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


async def main_logic():
    server = await asyncio.start_server(handle, '127.0.0.1', 8765)

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

    async with server:
        await server.serve_forever()

def main():
    loop = asyncio.get_event_loop()
    loop.run_until_complete(main_logic())

if __name__ == "__main__":
    main()