# -*- coding: utf-8 -*-
import socket
import threading
import sys
import time
import inspect

TARGET_HOST = 'zentao.sunyard.com.cn'
TARGET_PORT = 9788

def log_info(source, message):
    """Log info message"""
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
    frame = inspect.currentframe().f_back
    line = frame.f_lineno
    print("[%s] [INFO] [%s] (line %d) %s" % (timestamp, source, line, message))

def log_error(source, message):
    """Log error message"""
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
    frame = inspect.currentframe().f_back
    line = frame.f_lineno
    print("[%s] [ERROR] [%s] (line %d) %s" % (timestamp, source, line, message))

def handle_request(b_socket, request_data):
    """Handle a single request from device B"""
    # Parse HTTP request
    request_lines = request_data.split(b"\r\n")
    if request_lines:
        first_line = request_lines[0].decode('utf-8', errors='ignore')
        parts = first_line.split()
        if len(parts) >= 2:
            method = parts[0]
            url = parts[1]
            log_info("System", "Received request: %s %s" % (method, url))
    
    # Connect to target
    target_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    target_sock.settimeout(30)
    
    try:
        target_sock.connect((TARGET_HOST, TARGET_PORT))
        log_info("System", "Connected to target: %s:%d" % (TARGET_HOST, TARGET_PORT))
        
        if method == "CONNECT":
            # Send CONNECT response
            response = b"HTTP/1.1 200 Connection Established\r\n\r\n"
            b_socket.sendall(response)
            log_info("System", "Sent CONNECT response")
            
            # Establish tunnel
            def forward(src, dst):
                try:
                    while True:
                        data = src.recv(8192)
                        if not data:
                            break
                        dst.sendall(data)
                except:
                    pass
            
            t1 = threading.Thread(target=forward, args=(b_socket, target_sock))
            t2 = threading.Thread(target=forward, args=(target_sock, b_socket))
            t1.start()
            t2.start()
            t1.join()
            t2.join()
        
        else:
            # Forward request and response
            target_sock.sendall(request_data)
            
            while True:
                response = target_sock.recv(8192)
                if not response:
                    break
                b_socket.sendall(response)
        
        return True
    
    except ConnectionRefusedError:
        log_error("System", "Cannot connect to target: %s:%d" % (TARGET_HOST, TARGET_PORT))
        error_msg = b"HTTP/1.1 502 Bad Gateway\r\nContent-Type: text/plain\r\n\r\nCannot connect to target\r\n"
        b_socket.sendall(error_msg)
        return False
    except socket.timeout:
        log_error("System", "Target connection timeout")
        error_msg = b"HTTP/1.1 504 Gateway Timeout\r\nContent-Type: text/plain\r\n\r\nConnection timeout\r\n"
        b_socket.sendall(error_msg)
        return False
    except Exception as e:
        log_error("System", "Error handling request: %s" % str(e))
        return False
    finally:
        target_sock.close()

def listen_for_requests(sock):
    """Listen for requests from device B"""
    try:
        while True:
            # Wait for request with long timeout
            sock.settimeout(600)  # 10 minutes
            
            request_data = b""
            while b"\r\n\r\n" not in request_data:
                chunk = sock.recv(8192)
                if not chunk:
                    log_info("System", "Device B disconnected")
                    return
                request_data += chunk
            
            # Handle request in a separate thread
            t = threading.Thread(target=handle_request, args=(sock, request_data))
            t.start()
            
    except socket.timeout:
        log_info("System", "Connection idle timeout")
    except Exception as e:
        log_error("System", "Connection error: %s" % str(e))

def connect_to_b_device(b_host, b_port):
    """Connect to device B and handle requests"""
    while True:
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(10)
            sock.connect((b_host, b_port))
            
            # Wait for connection confirmation from server
            sock.settimeout(5)
            response = sock.recv(1024)
            if not response or b"PROXY_CONNECTED" not in response:
                log_error("System", "Failed to receive connection confirmation")
                sock.close()
                time.sleep(3)
                continue
            
            log_info("System", "Successfully connected to device B: %s:%d" % (b_host, b_port))
            
            listen_for_requests(sock)
            
        except socket.error as e:
            if e.errno == 111 or e.errno == 10061:
                log_info("System", "Device B not ready (%s:%d), retrying in 3 seconds..." % (b_host, b_port))
            else:
                log_error("System", "Connection error: %s" % str(e))
            time.sleep(3)
        except socket.timeout:
            log_info("System", "Connection timeout, retrying in 3 seconds...")
            time.sleep(3)
        except Exception as e:
            log_error("System", "Connection error: %s" % str(e))
            time.sleep(3)

def main():
    if len(sys.argv) < 3:
        print("=" * 60)
        print("           Proxy Client (Run on Computer A)")
        print("=" * 60)
        print("Target website:", TARGET_HOST, ":", TARGET_PORT)
        print("Usage: python proxy_client.py <Device B IP> <Device B Port>")
        print("Example: python proxy_client.py 192.168.1.100 8080")
        print("=" * 60)
        sys.exit(1)
    
    b_host = sys.argv[1]
    b_port = int(sys.argv[2])
    
    print("=" * 60)
    print("           Proxy Client (Run on Computer A)")
    print("=" * 60)
    print("Target website:", TARGET_HOST, ":", TARGET_PORT)
    print("Device B address:", b_host, ":", b_port)
    print("=" * 60)
    
    try:
        connect_to_b_device(b_host, b_port)
    except KeyboardInterrupt:
        log_info("System", "Received Ctrl+C, exiting...")
        sys.exit(0)

if __name__ == '__main__':
    main()