forward-port.py

#!/usr/bin/env python3
"""Forward a local tcp port to host:port.

Usage: %(prog)s local_port:host:port

Example:

  $ python3 forward-port.py 26992:icanhazip.com:80 # start server

and in another window:

  $ curl localhost:26992 # connect to it
"""
import asyncio
import logging
import sys

info = logging.getLogger('forward-port').info

@asyncio.coroutine
def copy_stream(reader, writer, bufsize=1<<16):
    while True:
        data = yield from reader.read(bufsize)
        if not data:
            break
        writer.write(data)
        yield from writer.drain()
    writer.close()

def port_forwarder(host, port, *, loop):
    @asyncio.coroutine
    def forward(local_reader, local_writer):
        client = local_writer.get_extra_info('peername')
        info('connected client %s %s', *client)
        remote_reader, remote_writer = yield from asyncio.open_connection(host, port, loop=loop)
        yield from asyncio.wait([copy_stream(local_reader, remote_writer),
                                 copy_stream(remote_reader, local_writer)],
                                loop=loop)
        info('disconnected client %s %s', *client)

    return forward

# main
logging.basicConfig(level=logging.INFO,
                    format="%(asctime)-15s %(message)s", datefmt="%F %T")
if len(sys.argv) != 2:
    sys.exit(__doc__)
local_port, host, port = sys.argv[1].split(':') # e.g., 12345:localhost:22

loop = asyncio.get_event_loop()
server = loop.run_until_complete(asyncio.start_server(port_forwarder(host, int(port), loop=loop),
                                                      'localhost', int(local_port), loop=loop))
info('listening on: %s %s', *server.sockets[0].getsockname())
for closing in range(2):
    try:
        loop.run_until_complete(server.wait_closed())
    except KeyboardInterrupt:
        if not closing:
            server.close()
            info('closing server')
    else:
        break
    info('done')
loop.close()