kadence/desktop/helper/kadence-helper

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

It does exactly one thing — maintain a marked region inside `/etc/hosts` — and
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`.

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 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

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 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")
    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()
    else:
        command_status()
    return 0


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