########################################################################################################################
# IMPORTS

import contextlib
import logging
import select
import socket
import struct
import threading

import requests
import socks

########################################################################################################################
# SETUP

logger = logging.getLogger(__name__)

########################################################################################################################
# CLASS


class Socks5Bridge:
    """
    Local SOCKS5 proxy (no auth) that forwards TCP through an authenticated SOCKS5 proxy
    and intercepts WebRTC STUN via UDP ASSOCIATE.

    Playwright's Firefox does not support SOCKS5 proxy authentication, but it does support
    unauthenticated SOCKS5. This bridge:

    - TCP CONNECT  → forwarded to upstream SOCKS5 with credentials via PySocks
    - UDP ASSOCIATE → local relay that intercepts STUN Binding Requests and replies
                      with the proxy's exit IP, preventing any real STUN UDP from
                      leaving the machine and revealing the host's real IP.

    Firefox sees ``socks5://127.0.0.1:{port}`` with no credentials required.

    Usage:
        bridge = Socks5Bridge(host, port, username, password)
        local_port = bridge.start()          # resolves exit_ip and begins listening
        # use socks5://127.0.0.1:{local_port} as proxy for Playwright
        bridge.stop()
    """

    # SOCKS5
    _VER = 0x05
    _NO_AUTH = 0x00
    _CMD_CONNECT = 0x01
    _CMD_UDP_ASSOCIATE = 0x03
    _ATYP_IPV4 = 0x01
    _ATYP_DOMAIN = 0x03
    _ATYP_IPV6 = 0x04
    _REP_SUCCESS = 0x00
    _REP_FAILURE = 0x01
    _CMD_NOT_SUPPORTED = 0x07

    # STUN (RFC 5389)
    _STUN_MAGIC = 0x2112A442
    _STUN_BINDING_REQUEST = 0x0001
    _STUN_BINDING_RESPONSE = 0x0101

    def __init__(self, host: str, port: int, username: str | None = None, password: str | None = None):
        self.host = host
        self.port = int(port)
        self.username = username
        self.password = password
        self.exit_ip: str | None = None
        self._server: socket.socket | None = None
        self._local_port: int | None = None
        self._running = False
        self._thread: threading.Thread | None = None

    # ------------------------------------------------------------------------------------------------------------------
    # Public API

    @classmethod
    def resolve_exit_ip(cls, host: str, port: int, username: str | None = None, password: str | None = None) -> str:
        """Fetch the public exit IP of the SOCKS5 proxy directly (no bridge needed)."""
        auth = f"{username}:{password}@" if username and password else ""
        proxy_url = f"socks5://{auth}{host}:{port}"
        proxies = {"http": proxy_url, "https": proxy_url}
        r = requests.get("https://wtfismyip.com/json", proxies=proxies, timeout=10)
        return r.json()["YourFuckingIPAddress"]

    def start(self) -> int:
        """Resolve exit IP, start the bridge, return the local port."""
        if self.exit_ip is None:
            self.exit_ip = self.resolve_exit_ip(self.host, self.port, self.username, self.password)
            logger.debug(f"SOCKS5 exit IP: {self.exit_ip}")

        self._server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._server.bind(("127.0.0.1", 0))
        self._server.settimeout(1)
        self._local_port = self._server.getsockname()[1]
        self._server.listen(50)
        self._running = True
        self._thread = threading.Thread(target=self._accept_loop, daemon=True, name="socks5-bridge")
        self._thread.start()
        logger.info(f"SOCKS5 bridge listening on 127.0.0.1:{self._local_port} → {self.host}:{self.port}")
        return self._local_port

    def stop(self) -> None:
        self._running = False
        if self._server:
            with contextlib.suppress(Exception):
                self._server.close()
        if self._thread:
            self._thread.join(timeout=2)

    # ------------------------------------------------------------------------------------------------------------------
    # TCP accept loop

    def _accept_loop(self) -> None:
        while self._running:
            try:
                conn, _ = self._server.accept()
                threading.Thread(target=self._handle, args=(conn,), daemon=True).start()
            except socket.timeout:
                continue
            except Exception:
                break

    # ------------------------------------------------------------------------------------------------------------------
    # SOCKS5 handshake dispatcher

    def _handle(self, conn: socket.socket) -> None:
        try:
            # Greeting
            header = self._recvall(conn, 2)
            if len(header) < 2 or header[0] != self._VER:
                logger.warning(f"SOCKS5 bridge: bad greeting {header!r}")
                return
            self._recvall(conn, header[1])  # discard offered methods
            conn.sendall(bytes([self._VER, self._NO_AUTH]))

            # Request header (4 bytes: VER CMD RSV ATYP) — wait after auth negotiation
            req = self._recvall(conn, 4)
            if len(req) < 4 or req[0] != self._VER:
                logger.warning(f"SOCKS5 bridge: bad request header {req!r}")
                conn.sendall(self._reply(self._REP_FAILURE))
                return

            cmd, atyp = req[1], req[3]
            logger.debug(f"SOCKS5 bridge: cmd={cmd} atyp={atyp}")

            if cmd == self._CMD_CONNECT:
                self._handle_connect(conn, atyp)
            elif cmd == self._CMD_UDP_ASSOCIATE:
                self._handle_udp_associate(conn, atyp)
            else:
                logger.warning(f"SOCKS5 bridge: unsupported cmd={cmd}")
                conn.sendall(self._reply(self._CMD_NOT_SUPPORTED))
        except Exception as ex:
            logger.warning(f"SOCKS5 bridge handle error: {ex}")
        finally:
            with contextlib.suppress(Exception):
                conn.close()

    # ------------------------------------------------------------------------------------------------------------------
    # TCP CONNECT → upstream SOCKS5

    def _handle_connect(self, conn: socket.socket, atyp: int) -> None:
        dest_host, dest_port = self._read_addr_port(conn, atyp)
        logger.debug(f"SOCKS5 bridge CONNECT → {dest_host}:{dest_port}")
        remote = None
        try:
            remote = socks.socksocket()
            remote.set_proxy(socks.SOCKS5, self.host, self.port, rdns=True,
                             username=self.username, password=self.password)
            remote.settimeout(30)
            remote.connect((dest_host, dest_port))
            conn.sendall(self._reply(self._REP_SUCCESS))
            self._relay(conn, remote)
        except Exception as ex:
            logger.warning(f"SOCKS5 bridge CONNECT error ({dest_host}:{dest_port}): {ex}")
            with contextlib.suppress(Exception):
                conn.sendall(self._reply(self._REP_FAILURE))
        finally:
            for s in (remote,):
                if s:
                    with contextlib.suppress(Exception):
                        s.close()

    # ------------------------------------------------------------------------------------------------------------------
    # UDP ASSOCIATE → local relay with STUN interception

    def _handle_udp_associate(self, conn: socket.socket, atyp: int) -> None:
        self._read_addr_port(conn, atyp)  # discard client's address hint

        udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udp_sock.bind(("127.0.0.1", 0))
        udp_port = udp_sock.getsockname()[1]

        # Reply with UDP relay address
        conn.sendall(self._reply_bind("127.0.0.1", udp_port))

        stop = threading.Event()
        threading.Thread(target=self._udp_relay_loop, args=(udp_sock, stop), daemon=True).start()

        # Keep TCP control connection alive until client closes it
        try:
            conn.settimeout(None)
            while conn.recv(1):
                pass
        except Exception:
            pass
        finally:
            stop.set()
            with contextlib.suppress(Exception):
                udp_sock.close()

    def _udp_relay_loop(self, udp_sock: socket.socket, stop: threading.Event) -> None:
        """
        Receives SOCKS5-encapsulated UDP datagrams from Firefox.

        STUN Binding Requests are intercepted and answered locally with the proxy
        exit IP — no real STUN UDP ever leaves the machine.
        """
        udp_sock.settimeout(1)
        while not stop.is_set():
            try:
                data, client_addr = udp_sock.recvfrom(65536)
            except socket.timeout:
                continue
            except Exception:
                break

            try:
                dest_host, dest_port, payload, hdr = self._parse_udp_packet(data)
                if payload is None:
                    continue

                if self._is_stun_binding_request(payload):
                    stun_resp = self._build_stun_response(payload, self.exit_ip)
                    # Wrap in SOCKS5 UDP header (source = original STUN server so Firefox accepts it)
                    udp_sock.sendto(hdr + stun_resp, client_addr)
                    logger.debug(f"STUN intercepted → spoofed with {self.exit_ip}")
                # Non-STUN UDP: drop (no real UDP upstream proxy)
            except Exception as ex:
                logger.debug(f"UDP relay error: {ex}")

    # ------------------------------------------------------------------------------------------------------------------
    # Helpers: SOCKS5

    @staticmethod
    def _recvall(conn: socket.socket, n: int) -> bytes:
        """Read exactly n bytes from a socket."""
        buf = b""
        while len(buf) < n:
            chunk = conn.recv(n - len(buf))
            if not chunk:
                raise ConnectionError("Connection closed before receiving expected bytes")
            buf += chunk
        return buf

    def _read_addr_port(self, conn: socket.socket, atyp: int) -> tuple[str, int]:
        if atyp == self._ATYP_IPV4:
            host = socket.inet_ntoa(self._recvall(conn, 4))
        elif atyp == self._ATYP_DOMAIN:
            length = self._recvall(conn, 1)[0]
            host = self._recvall(conn, length).decode()
        elif atyp == self._ATYP_IPV6:
            host = socket.inet_ntop(socket.AF_INET6, self._recvall(conn, 16))
        else:
            raise ValueError(f"Unknown ATYP: {atyp}")
        port = struct.unpack("!H", self._recvall(conn, 2))[0]
        return host, port

    @staticmethod
    def _reply(rep: int) -> bytes:
        """Standard reply with BND 0.0.0.0:0."""
        return struct.pack("!BBBBIH", 0x05, rep, 0x00, 0x01, 0, 0)

    @staticmethod
    def _reply_bind(host: str, port: int) -> bytes:
        """Reply carrying the actual bind address (for UDP ASSOCIATE)."""
        return (struct.pack("!BBBB", 0x05, 0x00, 0x00, 0x01)
                + socket.inet_aton(host)
                + struct.pack("!H", port))

    @staticmethod
    def _relay(client: socket.socket, remote: socket.socket) -> None:
        while True:
            try:
                r, _, _ = select.select([client, remote], [], [], 30)
            except Exception:
                break
            if not r:
                break
            for s in r:
                try:
                    data = s.recv(8192)
                    if not data:
                        return
                    (remote if s is client else client).sendall(data)
                except Exception:
                    return

    # ------------------------------------------------------------------------------------------------------------------
    # Helpers: SOCKS5 UDP encapsulation

    @staticmethod
    def _parse_udp_packet(data: bytes) -> tuple:
        """Parse SOCKS5 UDP datagram. Returns (dest_host, dest_port, payload, reply_header)."""
        if len(data) < 4:
            return None, None, None, None
        frag = data[2]
        if frag != 0:
            return None, None, None, None  # fragmentation not supported
        atyp = data[3]
        if atyp == 0x01:  # IPv4
            if len(data) < 10:
                return None, None, None, None
            dest_host = socket.inet_ntoa(data[4:8])
            dest_port = struct.unpack("!H", data[8:10])[0]
            payload = data[10:]
            hdr = b"\x00\x00\x00\x01" + data[4:10]
        elif atyp == 0x03:  # domain
            if len(data) < 5:
                return None, None, None, None
            dlen = data[4]
            if len(data) < 5 + dlen + 2:
                return None, None, None, None
            dest_host = data[5:5 + dlen].decode()
            dest_port = struct.unpack("!H", data[5 + dlen:7 + dlen])[0]
            payload = data[7 + dlen:]
            hdr = b"\x00\x00\x00\x03" + bytes([dlen]) + data[5:5 + dlen] + data[5 + dlen:7 + dlen]
        else:
            return None, None, None, None
        return dest_host, dest_port, payload, hdr

    # ------------------------------------------------------------------------------------------------------------------
    # Helpers: STUN (RFC 5389)

    @staticmethod
    def _is_stun_binding_request(data: bytes) -> bool:
        if len(data) < 20:
            return False
        msg_type = struct.unpack("!H", data[0:2])[0]
        magic = struct.unpack("!I", data[4:8])[0]
        return msg_type == 0x0001 and magic == 0x2112A442

    @classmethod
    def _build_stun_response(cls, request: bytes, mapped_ip: str) -> bytes:
        """Build a STUN Binding Success Response with XOR-MAPPED-ADDRESS = mapped_ip."""
        transaction_id = request[8:20]
        magic = cls._STUN_MAGIC

        # XOR-MAPPED-ADDRESS attribute (0x0020)
        port = 0  # we don't know the source port Firefox used, use 0
        x_port = port ^ (magic >> 16)
        x_addr = struct.unpack("!I", socket.inet_aton(mapped_ip))[0] ^ magic
        attr_val = struct.pack("!BBHI", 0x00, 0x01, x_port, x_addr)
        attr = struct.pack("!HH", 0x0020, len(attr_val)) + attr_val

        # STUN header
        header = struct.pack("!HHI", cls._STUN_BINDING_RESPONSE, len(attr), magic) + transaction_id
        return header + attr
