"""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  = 2          # "default_public_and_private_interfaces"


@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)."""
    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:
                try:
                    val = int(winreg.QueryValueEx(k, _WEBRTC_POLICY_VALUE)[0])
                    name = path.split("\\")[-1]
                    if val == _WEBRTC_BLOCK_VALUE:
                        found.append(name)
                    else:
                        missing.append(f"{name}={val} (need {_WEBRTC_BLOCK_VALUE})")
                except OSError:
                    missing.append(path.split("\\")[-1] + " key missing")
        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