TCP proxy in Python - some IP-restricted programs requre a bit of TCP magic...

tcpproxy.py

#!/usr/bin/env python
# -*- coding: UTF-8 -*-

from optparse import OptionParser
import os
from select import select
import socket

def main():
    op = OptionParser()
    options, args = op.parse_args()
    if len(args) != 3:
        raise Exception("Arguments expected: <bind port> <host> <port>")
    tcp_proxy(args[0], args[1], args[2])

def tcp_proxy(bind_port, connect_host, connect_port):
    ls = get_listening_socket(bind_port)
    print "Ready (%r)" % ls
    while True:
        conn, addr = ls.accept()
        print "Accepted connection %s by %s" % (conn, addr)
        pid = os.fork()
        if pid == 0:
            # child
            ls.close()
            try:
                process_client_connection(conn, connect_host, connect_port)
            finally:
                os._exit(0)
        else:
            # parent
            conn.close()

def process_client_connection(conn, connect_host, connect_port):
    c2 = get_connected_socket(connect_host, connect_port)
    while True:
        r, w, x = select([conn, c2], [], [])
        if conn in r:
            data = conn.recv(4096)
            print "Req: %r" % data
            if data == "":
                break
            c2.sendall(data)
        if c2 in r:
            data = c2.recv(4096)
            if data == "":
                break
            print "Resp: %r" % data
            conn.sendall(data)
    conn.close()
    c2.close()
    print "Closed"

def get_connected_socket(connect_host, connect_port):
    s = None
    ares = socket.getaddrinfo(connect_host, connect_port,
        socket.AF_UNSPEC, socket.SOCK_STREAM)
    for res in ares:
        af, socktype, proto, canonname, sa = res
        try:
            s = socket.socket(af, socktype, proto)
        except socket.error as msg:
            s = None
            continue
        try:
            s.connect(sa)
        except socket.error as msg:
            s.close()
            s = None
            continue
        break
    if s is None:
        raise Exception("could not open socket")
    return s


def get_listening_socket(bind_port):
    ares = socket.getaddrinfo(None, bind_port,
                              socket.AF_UNSPEC, socket.SOCK_STREAM,
                              0, socket.AI_PASSIVE)
    for res in ares:
        af, socktype, proto, canonname, sa = res
        try:
            s = socket.socket(af, socktype, proto)
        except socket.error as msg:
            s = None
            continue
        try:
            s.bind(sa)
            s.listen(1)
        except socket.error as msg:
            s.close()
            s = None
            continue
        break
    if s is None:
        raise Exception("could not open listening socket")
    return s



if __name__ == "__main__":
    main()