...

Text file src/github.com/datawire/ambassador/v2/python/tests/_forward.py

Documentation: github.com/datawire/ambassador/v2/python/tests

     1import socket
     2import sys
     3from threading import Thread
     4
     5
     6def server(dst_host: str, dst_port: int, src_host: str, src_port: int) -> None:
     7    listen_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     8    listen_sock.bind((src_host, src_port))
     9    listen_sock.listen(5)
    10
    11    while True:
    12        src_sock = listen_sock.accept()[0]
    13        dst_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    14        dst_sock.connect((dst_host, dst_port))
    15        Thread(target=forward, args=(src_sock, dst_sock), daemon=True).start()
    16        Thread(target=forward, args=(dst_sock, src_sock), daemon=True).start()
    17
    18
    19def forward(dst: socket.socket, src: socket.socket) -> None:
    20    while True:
    21        data = src.recv(64 * 1024)
    22        if data:
    23            dst.sendall(data)
    24        else:
    25            try:
    26                # Close destination first as origin is likely already closed
    27                dst.shutdown(socket.SHUT_WR)
    28                src.shutdown(socket.SHUT_RD)
    29            except OSError:
    30                pass
    31            return
    32
    33
    34if __name__ == "__main__":
    35    dst_host, dst_port_str, src_host, src_port_str = sys.argv[1:]
    36    dst_port = int(dst_port_str)
    37    src_port = int(src_port_str)
    38    server(dst_host, dst_port, src_host, src_port)

View as plain text