"""
Firewall kill switch — enforces that ALL traffic goes through VPN/DoH.
Uses macOS pf (packet filter) to block leaks.
"""

from __future__ import annotations

import logging
import subprocess
import tempfile
import os
from pathlib import Path

logger = logging.getLogger("school-shield.firewall")

PF_RULES_DIR = Path("/etc/pf.anchors")
SHIELD_ANCHOR = "school-shield"

# Rules: block everything EXCEPT:
# - DNS to our local DoH proxy (127.0.0.1:53)
# - HTTPS to DoH servers (Cloudflare, Quad9)
# - WireGuard traffic (UDP port 51820)
# - loopback
KILLSWITCH_RULES = f"""
# School Shield kill switch — ENFORCED MODE
# Generated by school-shield — do not edit manually

# Allow loopback
pass quick on lo0 all

# Allow DNS queries to our local DoH proxy
pass quick proto udp from any to 127.0.0.1 port 53

# Allow HTTPS to DoH servers (so DoH proxy works)
pass quick proto tcp from any to 1.1.1.1 port 443
pass quick proto tcp from any to 1.0.0.1 port 443
pass quick proto tcp from any to 9.9.9.9 port 443
pass quick proto tcp from any to 149.112.112.112 port 443

# Allow WireGuard traffic (VPN tunnel)
pass quick proto udp from any to any port 51820

# Allow established connections
pass quick proto tcp from any to any flags S/SA keep state

# Block everything else outbound
block drop out log all
"""

# Block only non-DoH DNS queries (softer mode — DNS leak prevention only)
DNS_ONLY_RULES = f"""
# School Shield — DNS leak prevention only
pass quick on lo0 all
pass quick proto udp from any to 127.0.0.1 port 53
pass quick proto udp from any to 1.1.1.1 port 53
pass quick proto udp from any to 9.9.9.9 port 53

# Block DNS to any other server
block drop out proto udp from any to any port 53
block drop out proto tcp from any to any port 53
"""


def _write_rules(rules: str) -> Path:
    """Write pf rules to a temp file."""
    rules_file = PF_RULES_DIR / SHIELD_ANCHOR
    rules_file.parent.mkdir(parents=True, exist_ok=True)
    rules_file.write_text(rules)
    rules_file.chmod(0o644)
    return rules_file


def _run_pfctl(*args: str) -> bool:
    """Run pfctl command."""
    cmd = ["sudo", "pfctl"] + list(args)
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
    if result.returncode != 0:
        logger.warning("pfctl %s failed: %s", " ".join(args), result.stderr)
        return False
    return True


def enable_killswitch(mode: str = "full") -> bool:
    """Enable the kill switch.

    Args:
        mode: 'full' (block everything except VPN+DoH) or 'dns' (block non-DoH DNS only)
    """
    rules = KILLSWITCH_RULES if mode == "full" else DNS_ONLY_RULES
    rules_file = _write_rules(rules)

    # Load rules
    if not _run_pfctl("-a", SHIELD_ANCHOR, "-f", str(rules_file)):
        return False

    # Enable pf if not already running
    _run_pfctl("-e")

    logger.info("🛡️  Kill switch ENABLED (%s mode)", mode)
    return True


def disable_killswitch() -> bool:
    """Disable the kill switch — restore normal network."""
    # Flush our anchor
    _run_pfctl("-a", SHIELD_ANCHOR, "-F", "all")

    # Remove anchor file
    rules_file = PF_RULES_DIR / SHIELD_ANCHOR
    if rules_file.exists():
        rules_file.unlink()

    logger.info("🔓 Kill switch DISABLED")
    return True


def status() -> dict:
    """Check kill switch status."""
    result = subprocess.run(
        ["sudo", "pfctl", "-a", SHIELD_ANCHOR, "-sr"],
        capture_output=True,
        text=True,
        timeout=5,
    )
    rules = result.stdout.strip()
    return {
        "enabled": bool(rules),
        "rules": rules.split("\n") if rules else [],
        "rule_count": len(rules.split("\n")) if rules else 0,
    }