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 192.168.1.100 5000 127.0.0.1 8080     # Proxy port 5000 to localhost:8080
    Open http://192.168.1.100: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.')
        # Place your replace, prints etc. here
        writer.write(data)
        await writer.drain()
    writer.close()
    print(name, 'Closed connection to writer.')


async def handle_proxy(dst_port, dst_ip, src_reader, src_writer):
    "Open connection to destination and proxy with source."
    dst_reader, dst_writer = await asyncio.open_connection(dst_ip, 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, src_ip, dst_ip):
    "Run server with proxy handler forever."
    handler = functools.partial(handle_proxy, dst_port, dst_ip)
    server = await asyncio.start_server(handler, src_ip, 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_ip', type=str)
    parser.add_argument('src_port', type=int)
    parser.add_argument('dst_ip', type=str)
    parser.add_argument('dst_port', type=int)
    args = parser.parse_args()
    asyncio.run(serve(args.src_port, args.dst_port, args.src_ip, args.dst_ip ))


if __name__ == '__main__':
    main()