How to make a simple proxy in python?

I wanted to make a very simple proxy using python (mostly to understand how it works). I am talking about a generic TCP proxy, not only http. I have built the following code, however, it seems to only work one way : ie the request is sent but I never get the answer. Here is the code :

import socket
import argparse

#Args
parser = argparse.ArgumentParser(description='ProxyDescription')
parser.add_argument('-l', '--listen', action='store', help='Listening port', default=80, type=int)
parser.add_argument('destination', action='store', help='Destination host')
parser.add_argument('port', action='store', help='Destination port', type=int)
args = parser.parse_args()

#Server
s1 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s1.bind(('', args.listen))
s1.listen(1)
conn1, addr1 = s1.accept()

#Client
s2 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s2.connect((args.destination, args.port))
s2.setblocking(0)

print "Connected by ", addr1
while 1:
    datato = conn1.recv(1024)
    if not datato: 
        print "breaking to"
        break
    s2.send(datato)
    print "data send : " + datato
    try:
        datafrom = s2.recv(1024)
        print "reveived data : " + datafrom   
        if not datafrom:
            print "breakinf from"
            break
        print "datafrom: " + datafrom
        conn1.send(datafrom)
    except socket.error, msg:
        print "error rcving: " + str(socket.error) + " || " + str(msg) 
        continue
print "the end ... closing"
conn1.close()
s2.close()

My test is simply launching this script and telnet through it. If I look with wireshark, I see that the request is fully understood by the server and I do get an answer, however, I never get the answer in the telnet. (testing with simple GET / on google) I feel the problem is related to "blocking" / "non-blocking" socket but I can't really understand where it is.

I think this is the simplest TCP Proxy you can do using the socket module. No need for multiple threads thanks to kqueue.

import socket
import selectors
import logging


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class TCPProxy:
    def __init__(
        self,
        listen_addr: str,
        listen_port: int,
        destination_addr: str,
        destination_port: int,
    ):
        self.listen_addr = listen_addr
        self.listen_port = listen_port
        self.destination_addr = destination_addr
        self.destination_port = destination_port
        self.selector = selectors.DefaultSelector()
        self.connections = {}

    def on_accept(self, sock: socket.socket):
        conn, addr = sock.accept()
        logger.info("New Connection: %s", addr)

        # For each client connection we create a new server connection
        # as well. This lets us support multiple clients and not get the
        # data mixed up.
        send_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        send_sock.connect((self.destination_addr, self.destination_port))
        send_sock.setblocking(False)

        self.connections[conn] = send_sock

        # Start listening on the client -> proxy and the proxy -> server
        # sockets.
        self.selector.register(send_sock, selectors.EVENT_READ, self.on_data_from_client)
        self.selector.register(conn, selectors.EVENT_READ, self.on_data_from_server)

    def on_data_from_server(self, conn: socket.socket):
        data = conn.recv(1024)
        out_sock = self.connections[conn]

        if data:
            out_sock.send(data)
        else:
            # If we stop getting data close both halves of the connection.
            self.selector.unregister(conn)
            self.selector.unregister(out_sock)
            conn.close()
            out_sock.close()
            del self.connections[conn]

    def on_data_from_client(self, conn: socket.socket):
        data = conn.recv(1024)

        # This is just a reverse lookup on the connections dict.
        return_sock = next(x for x, y in self.connections.items() if y is conn)
        
        if data:
            return_sock.send(data)
        else:
            # If we stop getting data close both halves of the connection.
            self.selector.unregister(conn)
            self.selector.unregister(return_sock)
            conn.close()
            return_sock.close()
            del self.connections[return_sock]

    def start(self):
        # Create a socket for the proxy and have it listen on the address/port you set.
        self.proxy_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.proxy_sock.bind((self.listen_addr, self.listen_port))
        self.proxy_sock.setblocking(False)
        self.proxy_sock.listen()

        # Start listening for incoming data on the socket.
        self.selector.register(self.proxy_sock, selectors.EVENT_READ, self.on_accept)

        while True:
            # Wait until one of the sockets has data.
            events = self.selector.select()
            for key, _ in events:
                # Call the on_* method associated with it.
                callback = key.data
                callback(key.fileobj)


proxy = TCPProxy("0.0.0.0", 4000, "172.253.63.101", 80)
proxy.start()

References:

Recommended topics

Hot tags