"""WebRTC leak detection and hardening checks.

WebRTC leaks expose your real IP through STUN/ICE even when a proxy is set.
This module checks every controllable surface:

  1. OS-level Chrome / Edge Group Policy  (prevents non-proxied UDP)
  2. Firefox managed profile prefs         (user.js must disable PeerConnection)
  3. STUN UDP reachability                 (if STUN is reachable, WebRTC CAN leak)
  4. Firefox process prefs                 (confirm our user.js was accepted)
"""
from __future__ import annotations

import logging
import socket
import subprocess
import winreg
from dataclasses import dataclass, field
from pathlib import Path

log = logging.getLogger(__name__)

# STUN servers commonly used by WebRTC
_STUN_HOSTS: list[tuple[str, int]] = [
    ("stun.l.google.com",    19302),
    ("stun1.l.google.com",   19302),
    ("stun.cloudflare.com",  3478),
    ("stun.stunprotocol.org", 3478),
]

# Firefox user.js prefs that block WebRTC
_REQUIRED_WEBRTC_PREFS: dict[str, str] = {
    "media.peerconnection.enabled":             "false",
    "media.peerconnection.ice.default_address_only": "true",
    "media.peerconnection.ice.no_host":         "true",
}

# Chrome / Edge HKLM policy keys
_CHROMIUM_POLICY_PATHS = (
    r"SOFTWARE\Policies\Google\Chrome",
    r"SOFTWARE\Policies\Microsoft\Edge",
    r"SOFTWARE\Policies\Chromium",
)
_WEBRTC_POLICY_VALUE    = "DefaultWebRtcIpHandlingPolicy"
_WEBRTC_BLOCK_VALUE     = 3       # disable_non_proxied_udp (was 2 = public+private only — wrong)
_WEBRTC_POLICY_STR_KEY  = "WebRtcIPHandling"
_WEBRTC_BLOCK_STR_VALUE = "disable_non_proxied_udp"


@dataclass
class WebRtcCheckResult:
    """Aggregated WebRTC leak surface scan."""
    # Policy / config checks
    chrome_edge_policy_set: bool   = False
    chrome_edge_policy_detail: str = ""
    firefox_prefs_ok: bool         = False
    firefox_prefs_detail: str      = ""
    # Network reachability
    stun_reachable: bool           = False
    stun_detail: str               = ""
    # Convenience flags
    any_leak_risk: bool            = True
    issues: list[str]              = field(default_factory=list)
    checks: list[tuple[str, bool, str]] = field(default_factory=list)  # (name, ok, detail)

    def add(self, name: str, ok: bool, detail: str) -> None:
        self.checks.append((name, ok, detail))
        if not ok:
            self.issues.append(f"{name}: {detail}")


# ─────────────────────────────────────────────────────────────────────────────
# 1. Chrome / Edge Group Policy
# ─────────────────────────────────────────────────────────────────────────────

def check_chromium_webrtc_policy() -> tuple[bool, str]:
    """Return (policy_set, detail_string).

    Accepts either the legacy DWORD DefaultWebRtcIpHandlingPolicy==3
    OR the modern REG_SZ WebRtcIPHandling=="disable_non_proxied_udp".
    """
    found: list[str] = []
    missing: list[str] = []
    for path in _CHROMIUM_POLICY_PATHS:
        try:
            with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, path, 0, winreg.KEY_QUERY_VALUE) as k:
                name = path.split("\\")[-1]
                dword_ok = False
                str_ok = False
                # Check legacy DWORD
                try:
                    val = int(winreg.QueryValueEx(k, _WEBRTC_POLICY_VALUE)[0])
                    dword_ok = val == _WEBRTC_BLOCK_VALUE
                    if not dword_ok:
                        missing.append(f"{name} DWORD={val} (need {_WEBRTC_BLOCK_VALUE})")
                except OSError:
                    pass
                # Check modern REG_SZ
                try:
                    val_str = str(winreg.QueryValueEx(k, _WEBRTC_POLICY_STR_KEY)[0])
                    str_ok = val_str == _WEBRTC_BLOCK_STR_VALUE
                    if not str_ok:
                        missing.append(f"{name} REG_SZ={val_str!r} (need {_WEBRTC_BLOCK_STR_VALUE!r})")
                except OSError:
                    pass
                if dword_ok or str_ok:
                    found.append(name)
                elif not dword_ok and not str_ok:
                    missing.append(f"{name}: no WebRTC policy keys present")
        except OSError:
            continue   # key not present at all — browser not installed or not policy-managed

    if found:
        return True, "Policy set: " + ", ".join(found)
    if missing:
        return False, "Not set: " + "; ".join(missing)
    return False, "No Chrome/Edge/Chromium policy keys found (browsers may not be installed)"


# ─────────────────────────────────────────────────────────────────────────────
# 2. Firefox managed profile prefs
# ─────────────────────────────────────────────────────────────────────────────

def check_firefox_webrtc_prefs(profile_dir: Path) -> tuple[bool, str]:
    """Verify our written user.js contains the required WebRTC-disabling prefs."""
    user_js = profile_dir / "user.js"
    if not user_js.is_file():
        return False, f"user.js not found at {profile_dir} — profile not yet initialized"

    try:
        content = user_js.read_text(encoding="utf-8", errors="replace")
    except Exception as e:
        return False, f"Cannot read user.js: {e}"

    missing: list[str] = []
    for pref, expected_val in _REQUIRED_WEBRTC_PREFS.items():
        # Look for:  user_pref("pref.name", value);
        needle = f'"{pref}", {expected_val}'
        if needle not in content:
            missing.append(pref)

    if missing:
        return False, "Missing in user.js: " + ", ".join(missing)
    return True, "All WebRTC prefs present in managed profile"


# ─────────────────────────────────────────────────────────────────────────────
# 3. STUN server UDP reachability
# ─────────────────────────────────────────────────────────────────────────────

def _probe_stun_udp(host: str, port: int, timeout: float = 2.5) -> bool:
    """
    Send a minimal STUN Binding Request over UDP and wait for a response.
    Returns True if we get *any* response (STUN or ICMP unreachable).
    """
    # Minimal STUN Binding Request (RFC 5389)
    msg = (
        b"\x00\x01"  # Message Type: Binding Request
        b"\x00\x00"  # Message Length: 0
        b"\x21\x12\xa4\x42"  # Magic Cookie
        + b"\x00" * 12       # Transaction ID (12 bytes)
    )
    try:
        ip = socket.gethostbyname(host)
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.settimeout(timeout)
        s.sendto(msg, (ip, port))
        data, _ = s.recvfrom(512)
        s.close()
        return len(data) > 0
    except socket.timeout:
        return False
    except Exception:
        return False


def check_stun_reachability(timeout: float = 3.0) -> tuple[bool, str]:
    """
    Returns (reachable, detail).
    reachable = True means STUN servers ARE reachable → WebRTC COULD leak.
    """
    reachable: list[str] = []
    for host, port in _STUN_HOSTS[:3]:
        try:
            if _probe_stun_udp(host, port, timeout):
                reachable.append(f"{host}:{port}")
        except Exception:
            continue

    if reachable:
        return True, "STUN reachable: " + ", ".join(reachable) + " — WebRTC UDP is NOT blocked"
    return False, "STUN servers unreachable — WebRTC UDP appears blocked (good)"


# ─────────────────────────────────────────────────────────────────────────────
# Composite scan
# ─────────────────────────────────────────────────────────────────────────────

def run_webrtc_check(profile_dir: Path | None = None) -> WebRtcCheckResult:
    """
    Run all WebRTC leak surface checks.

    profile_dir: path to the managed Firefox profile (to verify user.js).
    """
    res = WebRtcCheckResult()

    # 1. Chrome/Edge policy
    policy_ok, policy_detail = check_chromium_webrtc_policy()
    res.chrome_edge_policy_set   = policy_ok
    res.chrome_edge_policy_detail = policy_detail
    res.add("Chrome/Edge WebRTC policy (HKLM)", policy_ok, policy_detail)

    # 2. Firefox profile prefs
    if profile_dir is not None:
        fx_ok, fx_detail = check_firefox_webrtc_prefs(profile_dir)
        res.firefox_prefs_ok     = fx_ok
        res.firefox_prefs_detail = fx_detail
        res.add("Firefox managed profile WebRTC prefs", fx_ok, fx_detail)
    else:
        res.add("Firefox managed profile WebRTC prefs", False,
                "Profile not specified — not checked")

    # 3. STUN UDP reachability
    stun_reachable, stun_detail = check_stun_reachability()
    res.stun_reachable = stun_reachable
    res.stun_detail    = stun_detail
    # STUN reachable = risk; NOT reachable = good (firewall is blocking it)
    res.add("STUN UDP reachability", not stun_reachable, stun_detail)

    # Overall verdict
    res.any_leak_risk = bool(res.issues)
    return res
