micro socat port forward in python

usocat.py

"""
python usocat.py <listening port> <downstream address>

Example:

python usocat.py 0:10001 192.168.1.5:8000
"""

import logging
import select
import socket
import sys


def forward(src, dst):
    data = src.recv(8*1024)
    if not data:
        raise ConnectionError(f"Connection closed from {src.getpeername()}")
    dst.sendall(data)


def handle_client(upstream, downstream):
    while True:
        socks, _, _ = select.select((upstream, downstream), (), ())
        for sock in socks:
            if sock is upstream:
                forward(upstream, downstream)
            elif sock is downstream:
                forward(downstream, upstream)


def main():
    import sys
    up = sys.argv[1].split(":")
    up = up[0], int(up[1])
    down = sys.argv[2].split(":")
    down = down[0], int(down[1])
    logging.basicConfig(level="INFO")
    serv = socket.socket()
    serv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    serv.bind(up)
    serv.listen(1)
    logging.info("Ready to accept request on %s...", up)
    while True:
        try:
            upstream, upstream_addr  = serv.accept()
            logging.info("New connection from %r", upstream_addr)
            logging.info("Connecting to %r", down)
            downstream = socket.create_connection(down)
            with upstream, downstream:
                handle_client(upstream, downstream)
        except OSError as error:
            logging.error("Communication error: %r", error)


if __name__ == "__main__":
    main()