#!/usr/bin/env python3
"""Kadence's privileged helper: the only part that runs as root.

It does two things, both of them small: maintain a marked region inside
`/etc/hosts`, and keep one file in `/etc/systemd/resolved.conf.d/` that points
the zones Kadence routes at the app's local DNS stub. It is written to be boring
on purpose:

* it never trusts its input: every hostname is matched against a strict pattern
  and the redirect address must parse as an IP, so nothing else can be smuggled
  into the file;
* it only ever rewrites the region between its own markers, so hand-written
  entries above and below are preserved byte for byte;
* the input list is opened with `O_NOFOLLOW` and must be a regular file, which
  closes the "point the helper at a symlink" trick;
* the write is atomic (temp file in `/etc` + `rename`), so an interrupted run
  cannot leave the machine with half a hosts file;
* the first run keeps a backup at `/etc/hosts.kadence-backup`;
* the DNS route is a file it owns outright, whose every line it composes: the
  zones are validated as hostnames and the port as a port, so no other setting
  can be smuggled in, and removing the file restores plain DNS.

Installed by `install.sh` together with a polkit action, so the desktop app can
call it through `pkexec` without ever running as root itself.

Usage:
    kadence-helper apply --ip 127.0.0.1 --file /tmp/list
    kadence-helper clear
    kadence-helper resolved-on --port 5353 --domains youtube.com,t.me
    kadence-helper resolved-off
    kadence-helper status
"""

from __future__ import annotations

import argparse
import ipaddress
import os
import re
import shutil
import stat as stat_module
import subprocess
import sys
import tempfile

HOSTS_PATH = "/etc/hosts"
BACKUP_PATH = "/etc/hosts.kadence-backup"
BEGIN_MARKER = "# >>> kadence block >>>"
END_MARKER = "# <<< kadence block <<<"
MAX_DOMAINS = 5000

DROPIN_DIR = "/etc/systemd/resolved.conf.d"
DROPIN_PATH = os.path.join(DROPIN_DIR, "kadence.conf")
STUB_ADDRESS = "127.0.0.1"
MIN_STUB_PORT = 1024

HOSTNAME_RE = re.compile(r"^(?=.{1,253}$)([a-z0-9]([a-z0-9-]{0,61}[a-z0-9])?\.)+[a-z]{2,63}$")

# Blocking these would break the machine, not your concentration.
FORBIDDEN = {"localhost", "localhost.localdomain"}


def fail(message: str) -> None:
    print(f"kadence-helper: {message}", file=sys.stderr)
    raise SystemExit(2)


def read_hosts() -> list[str]:
    try:
        with open(HOSTS_PATH, "r", encoding="utf-8") as handle:
            return handle.read().splitlines()
    except OSError as exc:
        fail(f"cannot read {HOSTS_PATH}: {exc}")
        return []


def strip_region(lines: list[str]) -> list[str]:
    """Everything except our managed region (and a stale, unterminated one)."""
    output: list[str] = []
    inside = False
    for line in lines:
        stripped = line.strip()
        if stripped == BEGIN_MARKER:
            inside = True
            continue
        if stripped == END_MARKER:
            inside = False
            continue
        if not inside:
            output.append(line)
    while output and not output[-1].strip():
        output.pop()
    return output


def write_hosts(lines: list[str]) -> None:
    if not os.path.exists(BACKUP_PATH):
        try:
            shutil.copy2(HOSTS_PATH, BACKUP_PATH)
        except OSError:
            pass  # a missing backup must not stop us unblocking later

    directory = os.path.dirname(HOSTS_PATH) or "/"
    descriptor, temp_path = tempfile.mkstemp(dir=directory, prefix=".hosts.kadence.")
    try:
        with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
            handle.write("\n".join(lines) + "\n")
        os.chmod(temp_path, 0o644)
        os.chown(temp_path, 0, 0)
        os.replace(temp_path, HOSTS_PATH)
    except OSError as exc:
        os.unlink(temp_path)
        fail(f"cannot write {HOSTS_PATH}: {exc}")


def flush_dns() -> None:
    """Best effort: a resolver cache would otherwise keep serving the old answer."""
    for command in (
        ["resolvectl", "flush-caches"],
        ["systemd-resolve", "--flush-caches"],
        ["nscd", "-i", "hosts"],
    ):
        if shutil.which(command[0]):
            subprocess.run(command, check=False, capture_output=True, timeout=10)


def load_domains(path: str) -> list[str]:
    """Read the list the client wrote, refusing anything surprising."""
    try:
        descriptor = os.open(path, os.O_RDONLY | os.O_NOFOLLOW)
    except OSError as exc:
        fail(f"cannot open the domain list: {exc}")
        return []
    with os.fdopen(descriptor, "r", encoding="utf-8", errors="replace") as handle:
        info = os.fstat(handle.fileno())
        if not stat_module.S_ISREG(info.st_mode):
            fail("the domain list must be a regular file")
        if info.st_size > 1_000_000:
            fail("the domain list is implausibly large")
        raw = handle.read()

    domains: list[str] = []
    for line in raw.splitlines():
        name = line.strip().lower()
        if not name or name.startswith("#"):
            continue
        if name in FORBIDDEN:
            continue
        if not HOSTNAME_RE.match(name):
            fail(f"refusing a hostname that does not look like one: {name[:80]!r}")
        if name not in domains:
            domains.append(name)
        if len(domains) > MAX_DOMAINS:
            fail(f"more than {MAX_DOMAINS} domains")
    return domains


def command_apply(ip: str, path: str) -> None:
    try:
        address = ipaddress.ip_address(ip)
    except ValueError:
        fail(f"not an IP address: {ip!r}")
        return
    domains = load_domains(path)
    lines = strip_region(read_hosts())
    if domains:
        lines.append("")
        lines.append(BEGIN_MARKER)
        lines.append("# Managed by Kadence. Edit the list in the app, not here.")
        for domain in domains:
            lines.append(f"{address}\t{domain}")
        lines.append(END_MARKER)
    write_hosts(lines)
    flush_dns()
    print(f"blocking {len(domains)} domain(s) -> {address}")


def command_clear() -> None:
    write_hosts(strip_region(read_hosts()))
    flush_dns()
    print("cleared")


def load_zones(raw: str) -> list[str]:
    """Validate the zone list the client sent. One bad name refuses the lot.

    These names go into a resolved configuration file, so they must be plain
    hostnames and nothing else — that is what makes it safe to compose the
    `Domains=` line around them.
    """
    zones: list[str] = []
    for part in raw.split(","):
        zone = part.strip().lower().strip(".")
        if not zone:
            continue
        if zone in FORBIDDEN:
            fail(f"refusing to route {zone!r}: it is a name this machine needs")
        if not HOSTNAME_RE.match(zone):
            fail(f"refusing a zone that does not look like a hostname: {zone[:80]!r}")
        if len(zone.split(".")) < 2:
            fail(f"refusing to route a whole top-level domain: {zone!r}")
        if zone not in zones:
            zones.append(zone)
        if len(zones) > MAX_DOMAINS:
            fail(f"more than {MAX_DOMAINS} zones")
    if not zones:
        fail("no zones given; use resolved-off to stop routing")
    return zones


def resolved_dropin_text(port: int, zones: list[str]) -> str:
    """The exact contents of ours, assembled from validated pieces."""
    return "\n".join(
        [
            "# Managed by Kadence. Edit the block lists in the app, not here.",
            "#",
            "# These zones are answered by the app's own DNS stub, which redirects",
            "# blocked names and forwards everything else to the real upstreams.",
            "# Deleting this file and restarting systemd-resolved restores plain",
            "# DNS; `kadence-helper resolved-off` does both.",
            "[Resolve]",
            f"DNS={STUB_ADDRESS}:{port}",
            "Domains=" + " ".join(f"~{zone}" for zone in zones),
            "",
        ]
    )


def write_dropin(text: str) -> None:
    try:
        os.makedirs(DROPIN_DIR, mode=0o755, exist_ok=True)
    except OSError as exc:
        fail(f"cannot create {DROPIN_DIR}: {exc}")
    descriptor, temp_path = tempfile.mkstemp(dir=DROPIN_DIR, prefix=".kadence.")
    try:
        with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
            handle.write(text)
        os.chmod(temp_path, 0o644)
        os.chown(temp_path, 0, 0)
        os.replace(temp_path, DROPIN_PATH)
    except OSError as exc:
        os.unlink(temp_path)
        fail(f"cannot write {DROPIN_PATH}: {exc}")


def restart_resolved() -> None:
    """Make resolved pick the file up. Its cache goes with it, which is wanted."""
    if not os.path.isdir("/run/systemd/resolve"):
        fail("this machine does not run systemd-resolved, so Kadence cannot route zones")
    if not shutil.which("systemctl"):
        fail("systemctl is not available to reload systemd-resolved")
    result = subprocess.run(
        ["systemctl", "restart", "systemd-resolved"], capture_output=True, text=True, timeout=60, check=False
    )
    if result.returncode != 0:
        detail = (result.stderr or result.stdout).strip()[:300]
        fail(f"could not restart systemd-resolved: {detail or result.returncode}")
    flush_dns()


def command_resolved_on(port: int, domains: str) -> None:
    if not MIN_STUB_PORT <= port <= 65535:
        fail(f"refusing port {port}: the stub has to bind an unprivileged port")
    zones = load_zones(domains)
    write_dropin(resolved_dropin_text(port, zones))
    restart_resolved()
    print(f"routing {len(zones)} zone(s) to {STUB_ADDRESS}:{port}")


def command_resolved_off() -> None:
    if not os.path.exists(DROPIN_PATH):
        print("no zones were routed")
        return
    try:
        os.unlink(DROPIN_PATH)
    except OSError as exc:
        fail(f"cannot remove {DROPIN_PATH}: {exc}")
    restart_resolved()
    print("stopped routing zones")


def command_status() -> None:
    inside = False
    count = 0
    for line in read_hosts():
        stripped = line.strip()
        if stripped == BEGIN_MARKER:
            inside = True
            continue
        if stripped == END_MARKER:
            inside = False
            continue
        if inside and stripped and not stripped.startswith("#"):
            count += 1
    print(f"{count} domain(s) blocked")


def main() -> int:
    parser = argparse.ArgumentParser(prog="kadence-helper", description=__doc__)
    sub = parser.add_subparsers(dest="command", required=True)
    apply_parser = sub.add_parser("apply", help="replace the managed block")
    apply_parser.add_argument("--ip", default="127.0.0.1")
    apply_parser.add_argument("--file", required=True, help="newline-separated hostnames")
    sub.add_parser("clear", help="remove the managed block")
    resolved_on = sub.add_parser("resolved-on", help="route DNS zones through the local stub")
    resolved_on.add_argument("--port", type=int, required=True, help="the stub's port")
    resolved_on.add_argument("--domains", required=True, help="comma-separated zone names")
    sub.add_parser("resolved-off", help="stop routing DNS zones through the stub")
    sub.add_parser("status", help="how many domains are blocked")

    args = parser.parse_args()
    if os.geteuid() != 0:
        fail("must run as root (it is meant to be called through pkexec)")

    if args.command == "apply":
        command_apply(args.ip, args.file)
    elif args.command == "clear":
        command_clear()
    elif args.command == "resolved-on":
        command_resolved_on(args.port, args.domains)
    elif args.command == "resolved-off":
        command_resolved_off()
    else:
        command_status()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())