tcp proxy

proxy.py

import socket
import threading
import sys


def receive_from(connection):
    buffer = b''

    # first respond allow 2 second timeout
    connection.settimeout(2)
    try:
        while True:
            data = connection.recv(4096)
            if not data:
                break
            buffer += data
            # expect subsequent data to be fast
            connection.settimeout(0.05)
    except:
        pass
    return buffer


def proxy_handler(client_socket, remote_host, remote_port, receive_first):
    # Define the remote socket used for forwarding requests
    remote_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    # Establish a connection to the remote host
    remote_socket.connect((remote_host, remote_port))

    while True:
        print("client => proxy")
        data = receive_from(client_socket)
        if (len(data) > 0):
            print(data)
            remote_socket.sendall(data)

            print("proxy => client")
            data = receive_from(remote_socket)
            if (len(data) > 0):
                print(data)
                client_socket.sendall(data)
            else:
                break
        else:
            break

    client_socket.close()
    remote_socket.close()


def server_loop(local_host, local_port, remote_host, remote_port, receive_first):
    # Define a server socket to listen on
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    try:
        # Bind the socket to the defined local address and port
        server.bind((local_host, local_port))
    except:
        print("[!!] Failed to connect to {0}:{1}".format(local_host, local_port))
        print("[!!] Check for other listening sockets or correct")
        sys.exit(0)

    print("Successfully listening on {0}:{1}".format(local_host, local_port))

    # Listen for a maximum of 5 connections
    server.listen(5)

    # Loop infinitely for incoming connections
    while True:
        client_socket, addr = server.accept()

        print("[==>] Received incoming connection from {0}:{1}".format(addr[0], addr[1])
              )
        # Start a new thread for any incoming connections
        proxy_thread = threading.Thread(target=proxy_handler, args=(
            client_socket, remote_host, remote_port, receive_first))
        proxy_thread.setDaemon(True)
        proxy_thread.start()


server_loop('0.0.0.0', int(sys.argv[1]), 'localhost', 8000, True)