from time import time
import trio  # python3 -m pip install --upgrade trio

async def forward_from_a_to_b(a, b):
    async for chunk in a:
        print(f"=> {chunk}", flush=True)
        await b.send_all(chunk)

async def proxy(listen_stream):
    peer = listen_stream.socket.getpeername()
    if len(peer) == 2:
        r_ip, r_port = peer
    else:
        # ('2a01:1111:2345:1337::1', 52345, 0, 0)
        assert len(peer) == 4, peer
        r_ip, r_port, _, __ = peer
    print(f"Conn from {r_ip}:{r_port} ({int(time())})")
    target_stream = await trio.open_tcp_stream(r_ip, 22)
    async with trio.open_nursery() as nursery:
        nursery.start_soon(forward_from_a_to_b, listen_stream, target_stream)
        nursery.start_soon(forward_from_a_to_b, target_stream, listen_stream)

async def main():
    while True:
        try:
            await trio.serve_tcp(proxy, 22)
        except KeyboardInterrupt:
            break
        except (trio.BrokenResourceError, OSError):
            print("Conn died")

if __name__ == '__main__':
    trio.run(main)
