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

def tunnel_thread(src, dst):
    """Tunnel thread: continuously forward data"""
    try:
        while True:
            data = src.recv(8192)
            if not data:
                break
            dst.sendall(data)
    except Exception as e:
        pass

def handle_b_device(sock):
    """Handle communication with device B, establish tunnel to target website"""
    print("\n[System] Connected to device B")
    
    try:
        while True:
            request = b""
            while True:
                chunk = sock.recv(8192)
                if not chunk:
                    print("[Device B] Disconnected")
                    return
                request += chunk
                if b"\r\n\r\n" in request:
                    break
            
            print("[System] Received request:", len(request), "bytes")
            
            target_host = None
            target_port = 443
            
            if request.startswith(b"CONNECT"):
                try:
                    lines = request.decode().split("\r\n")
                    connect_line = lines[0]
                    target = connect_line.split()[1]
                    if ':' in target:
                        target_host, port_str = target.split(':')
                        target_port = int(port_str)
                    else:
                        target_host = target
                        target_port = 443
                except Exception as e:
                    print("[Error] Failed to parse CONNECT request:", e)
                    sock.send(b"HTTP/1.1 400 Bad Request\r\n\r\n")
                    continue
            else:
                try:
                    for line in request.split(b"\r\n"):
                        if line.lower().startswith(b"host:"):
                            host_part = line[5:].strip()
                            if b':' in host_part:
                                target_host, port_str = host_part.split(b':')
                                target_port = int(port_str)
                            else:
                                target_host = host_part.decode()
                            break
                except Exception as e:
                    print("[Error] Failed to parse HTTP request:", e)
                    sock.send(b"HTTP/1.1 400 Bad Request\r\n\r\n")
                    continue
            
            if not target_host:
                sock.send(b"HTTP/1.1 400 Bad Request\r\n\r\n")
                continue
            
            print("[Target] Host:", target_host, "Port:", target_port)
            
            target_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            target_sock.settimeout(30)
            
            try:
                target_sock.connect((target_host, target_port))
                print("[Target] Connected successfully")
                
                if request.startswith(b"CONNECT"):
                    sock.send(b"HTTP/1.1 200 Connection Established\r\n\r\n")
                else:
                    target_sock.sendall(request)
                
                t1 = threading.Thread(target=tunnel_thread, args=(sock, target_sock))
                t2 = threading.Thread(target=tunnel_thread, args=(target_sock, sock))
                t1.start()
                t2.start()
                t1.join()
                t2.join()
                
            except ConnectionRefusedError:
                print("[Error] Cannot connect to target:", target_host, target_port)
                error_msg = "HTTP/1.1 503 Service Unavailable\r\nContent-Type: text/plain\r\n\r\nCannot connect to " + target_host + ":" + str(target_port)
                sock.send(error_msg.encode())
            except socket.timeout:
                print("[Error] Connection timeout")
                sock.send(b"HTTP/1.1 504 Gateway Timeout\r\n\r\n")
            except Exception as e:
                print("[Error] Failed to connect to target:", e)
                error_msg = "HTTP/1.1 500 Internal Server Error\r\nContent-Type: text/plain\r\n\r\n" + str(e)
                sock.send(error_msg.encode())
            finally:
                target_sock.close()
                
    except ConnectionResetError:
        print("[Device B] Connection forcibly closed")
    except Exception as e:
        print("[Error] Communication error with device B:", e)
    finally:
        sock.close()

def connect_to_b_device(b_host, b_port):
    """Connect to device B"""
    while True:
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(10)
            sock.connect((b_host, b_port))
            sock.settimeout(None)
            
            print("\n[System] Successfully connected to device B:", b_host, ":", b_port)
            
            handle_b_device(sock)
            
        except ConnectionRefusedError:
            print("[System] Device B not ready (", b_host, ":", b_port, "), retrying in 3 seconds...")
            import time
            time.sleep(3)
        except socket.timeout:
            print("[System] Connection timeout, retrying in 3 seconds...")
            import time
            time.sleep(3)
        except Exception as e:
            print("[Error] Connection error:", e)
            import time
            time.sleep(3)

def main():
    if len(sys.argv) < 3:
        print("=" * 60)
        print("           Proxy Client (Run on Computer A)")
        print("=" * 60)
        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("Device B address:", b_host, ":", b_port)
    print("=" * 60)
    
    connect_to_b_device(b_host, b_port)

if __name__ == '__main__':
    main()