Simple TCP Proxy

tcp_proxy.py

"""Simple TCP Proxy

Receive connections on one port and forward them to another port.

Example:

    $ python3 -m http.server 8080     # Start an http server on port 8080
    $ python3 tcp_proxy.py 5000 8080  # Proxy port 5000 to 8080
    Open http://localhost:5000/ in Chrome

"""

import argparse
import asyncio
import functools


async def reader_to_writer(name, reader, writer):
    "Forward data from reader to writer."
    while True:
        data = await reader.read(2 ** 12)
        if not data:
            break
        print(name, 'Received', len(data), 'bytes from reader.')
        writer.write(data)
        await writer.drain()
    writer.close()
    print(name, 'Closed connection to writer.')


async def handle_proxy(dst_port, src_reader, src_writer):
    "Open connection to destination and proxy with source."
    dst_reader, dst_writer = await asyncio.open_connection('localhost', dst_port)
    print('Opened connection to', dst_port)
    src_dst = reader_to_writer('SRC->DST', src_reader, dst_writer)
    dst_src = reader_to_writer('DST->SRC', dst_reader, src_writer)
    await asyncio.gather(src_dst, dst_src)


async def serve(src_port, dst_port):
    "Run server with proxy handler forever."
    handler = functools.partial(handle_proxy, dst_port)
    server = await asyncio.start_server(handler, 'localhost', src_port)
    print('Started server on', src_port)
    async with server:
        await server.serve_forever()


def main():
    "Parse source and destination ports then start server."
    parser = argparse.ArgumentParser()
    parser.add_argument('src_port', type=int)
    parser.add_argument('dst_port', type=int)
    args = parser.parse_args()
    asyncio.run(serve(args.src_port, args.dst_port))


if __name__ == '__main__':
    main()