python异步tcp代理服务器

import asyncio



class TcpServer:
    def __init__(self, addr, proxy_addr):
        self.addr = addr
        self.proxy_addr = proxy_addr
        self.server = None

    async def get_connect(self):
        reader, writer = await asyncio.open_connection(*self.proxy_addr)
        return reader, writer

    async def handle_data(self, data, flg):

        return data  # 这里简单返回原数据

    async def handle_rw(self, r, w, flg):
        try:
            while True:
                data = await r.read(1024)  # 读取数据，这里假设每次最多读取1024字节
                if not data:
                    break  # 连接关闭或没有更多数据
                processed_data = await self.handle_data(data, flg)
                w.write(processed_data)  # 发送处理后的数据
                await w.drain()  # 等待数据发送完毕
        except Exception as e:
            print(f"Error: {e}")
        finally:
            w.close()
            await asyncio.sleep(0.1)  # 确保关闭操作完成


    async def recv(self, r):
        return await r.read(1024)  # 示例：读取最多1024字节的数据

    async def send(self, w, data):
        w.write(data)
        await w.drain()

    def start_th(self, coro):
        asyncio.create_task(coro)  # 在事件循环中创建一个新的任务来运行协程

    async def handle_tcp(self, reader, writer):
        try:
            r2, w2 = await self.get_connect()  # 连接到代理服务器
            self.start_th(self.handle_rw(reader, w2, True))  # 启动读写协程处理客户端到代理的数据
            await self.handle_rw(r2, writer, False)  # 启动读写协程处理代理到客户端的数据
        except Exception as e:
            print(f"Handle TCP error: {e}")

    def run(self):
        
        self.server = asyncio.start_server(self.handle_tcp, *self.addr)
        asyncio.get_event_loop().run_until_complete(self.server)
        print(f"Serving on {self.addr}")
        asyncio.get_event_loop().run_forever()

# 使用示例
if __name__ == "__main__":
    server = TcpServer(('127.0.0.1', 8888), ('127.0.0.1', 9999))  # 监听地址和代理地址
    try:
        server.run()
    except KeyboardInterrupt:
        print("Server stopping...")
    finally:
        server.server.close()
        asyncio.get_event_loop().run_until_complete(server.server.wait_closed())
        print("Server closed.")