kadence/desktop/kadence_desktop/blocking/hosts.py

"""Website blocking through `/etc/hosts`.

Why the hosts file rather than a proxy or an extension: it needs no browser
support, survives a browser restart, covers every app on the machine, and is
trivially auditable — you can read exactly what is blocked with `cat`.

Writing it needs root, which a desktop app must not have, so all writes go
through the small privileged helper (`kadence-helper`) under polkit. The helper
only ever rewrites the region between its own markers and validates every
hostname, so the worst a compromised client can do is block a site.
"""

from __future__ import annotations

import logging
import os
import re
import shutil
import subprocess
import tempfile

log = logging.getLogger(__name__)

HOSTS_PATH = "/etc/hosts"
BEGIN_MARKER = "# >>> kadence block >>>"
END_MARKER = "# <<< kadence block <<<"

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


def valid_hostname(name: str) -> bool:
    return bool(_HOSTNAME_RE.match(name.strip().lower()))


def current_block() -> list[str]:
    """The hostnames currently inside our managed region of `/etc/hosts`."""
    try:
        with open(HOSTS_PATH, "r", encoding="utf-8") as handle:
            lines = handle.read().splitlines()
    except OSError:
        return []

    inside = False
    hosts: list[str] = []
    for line in lines:
        if line.strip() == BEGIN_MARKER:
            inside = True
            continue
        if line.strip() == END_MARKER:
            inside = False
            continue
        if inside:
            parts = line.split()
            if len(parts) >= 2 and not line.lstrip().startswith("#"):
                hosts.extend(parts[1:])
    return hosts


class HostsBlocker:
    """Applies and clears the managed block, as rarely as possible.

    Every change costs a polkit prompt in the worst case, so the desired state is
    diffed against what is already in the file and a no-op is really a no-op.
    """

    def __init__(self, helper_path: str = "/usr/local/bin/kadence-helper") -> None:
        self.helper_path = helper_path
        self._last_error: str | None = None

    # ---- availability ------------------------------------------------------

    @property
    def helper_installed(self) -> bool:
        return os.path.isfile(self.helper_path) and os.access(self.helper_path, os.X_OK)

    @property
    def last_error(self) -> str | None:
        return self._last_error

    def _runner(self) -> list[str] | None:
        """How to become root: polkit if we can, passwordless sudo otherwise."""
        if not self.helper_installed:
            self._last_error = (
                f"The privileged helper is not installed at {self.helper_path}. "
                "Run desktop/helper/install.sh once."
            )
            return None
        if shutil.which("pkexec") and os.environ.get("DISPLAY", os.environ.get("WAYLAND_DISPLAY")):
            return ["pkexec", self.helper_path]
        if shutil.which("sudo"):
            return ["sudo", "-n", self.helper_path]
        self._last_error = "Neither pkexec nor sudo is available to run the helper."
        return None

    # ---- operations --------------------------------------------------------

    def apply(self, domains: list[str], redirect_ip: str = "127.0.0.1") -> bool:
        """Make the managed region exactly `domains`. True when it already was."""
        wanted = sorted({d.strip().lower() for d in domains if valid_hostname(d)})
        if sorted(set(current_block())) == wanted:
            return True
        if not wanted:
            return self.clear()

        runner = self._runner()
        if runner is None:
            return False

        with tempfile.NamedTemporaryFile("w", suffix=".kadence", delete=False, encoding="utf-8") as handle:
            handle.write("\n".join(wanted))
            list_path = handle.name
        os.chmod(list_path, 0o644)  # the helper reads it as root from /tmp
        try:
            return self._run(runner + ["apply", "--ip", redirect_ip, "--file", list_path])
        finally:
            try:
                os.unlink(list_path)
            except OSError:
                pass

    def clear(self) -> bool:
        """Remove the managed region. Safe to call when it is not there."""
        if not current_block():
            return True
        runner = self._runner()
        if runner is None:
            return False
        return self._run(runner + ["clear"])

    def _run(self, command: list[str]) -> bool:
        try:
            result = subprocess.run(command, capture_output=True, text=True, timeout=60, check=False)
        except (OSError, subprocess.SubprocessError) as exc:
            self._last_error = f"Could not run the helper: {exc}"
            log.warning(self._last_error)
            return False
        if result.returncode != 0:
            detail = (result.stderr or result.stdout).strip()[:300]
            # 126/127 is polkit's "you said no" — worth saying plainly.
            if result.returncode in (126, 127):
                detail = detail or "the authorisation prompt was dismissed"
            self._last_error = f"Helper failed: {detail}"
            log.warning(self._last_error)
            return False
        self._last_error = None
        return True