python实现tcp代理

1.代理流程图

2.实现代码

#! usr/bin/python2
import sys
import socket
import thread

#handle local buffer
def request_handler(buffer):
    return buffer
#handle remote buffer
def response_handler(buffer):
    return buffer

#receive data from remote_server or remote_client
def receive_from(connection):
    buffer=""
    connection.settimeout(2)
    try:
        while True:
            data=connection.recv(4096)
            if not data:
                break
            buffer+=data
    except:
        pass
    return buffer

def server_loop(local_host,local_port,remote_host,remote_port,receive_first):
#作为服务器监听并接受remote_client连接
    server=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
    server.bind((local_host,local_port))
    server.listen(5)

    while True:
        client_socket,addr=server.accept()
        print("other client has connected to me")
        #打开一个线程处理与remote_client的连接，并作为客户端连接remote_server
        thread.start_new_thread(proxy_handler, (client_socket,remote_host,remote_port,receive_first))
        
def proxy_handler(client_socket,remote_host,remote_port,receive_first):
#作为客户端连接remote_server
    remote_socket=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
    remote_socket.connect((remote_host,remote_port))
    #判断是否先接受来自remote_server的信息
    if receive_first:
        remote_buffer=receive_from(remote_socket)
        print("from remote_server:
"+remote_buffer)
        if len(remote_buffer):
            print("sending to client")
            client_socket.send(remote_buffer)
    while True:
    #接受来自remote_client的信息并存储在local_buffer
    #将local_buffer的信息再发送到remote_server
        local_buffer=receive_from(client_socket)
        if len(local_buffer):
            print("from client:
"+local_buffer)
            local_buffer=request_handler(local_buffer)
            remote_socket.send(local_buffer)
            print("sent to remote_server")
    #接受来自remote_server的信息并存储在remote_buffer
    #将remote_buffer的信息再发送到remote_client
        remote_buffer=receive_from(remote_socket)
        if len(remote_buffer):
            print("from remote_server:
"+remote_buffer)
            remote_buffer=request_handler(remote_buffer)
            client_socket.send(remote_buffer)
            print("sent to client")
        
        
            
def main():
    if len(sys.argv[1:]) != 5:
        print("failed to input")
        sys.exit(0)
        
    local_host=sys.argv[1]
    local_port=int(sys.argv[2])
    
    remote_host=sys.argv[3]
    remote_port=int(sys.argv[4])
    
    receive_first=sys.argv[5]
    
    if "true" in receive_first:
        receive_first=True
    else:
        receive_first=False
        
    server_loop(local_host,local_port,remote_host,remote_port,receive_first)
main()