kadence/desktop/kadence_desktop/blocking/manager.py

"""Turning the server's policy into what this machine actually does.

Called on every poll with the policy the backend computed. It keeps the hosts
file, the DNS routes and the process list in step with it, reports every
turned-away attempt back to the API (that is what makes the daily report fun),
and — importantly — puts everything back the moment the policy says blocking is
off, including when the app exits or crashes mid-interval.

Two mechanisms do the website blocking, because one cannot do it all:

* `/etc/hosts` answers the concrete hostnames the server expanded — exact, cheap,
  and readable with `cat`;
* the patterns a hosts file cannot express (`*.youtube.com`, `*gram.com`) are
  answered by the local DNS stub, with `systemd-resolved` sent the affected zones
  through the helper's drop-in.

The stub and the drop-in are kept for as long as the app runs, not per interval:
repointing resolved costs a restart and a dropped DNS cache, and turning the
block on and off is instantaneous — the stub answers or forwards, per query.
That is why `release()` does not touch either of them, and why `shutdown()` is
its own method: the drop-in has to be gone *before* the stub stops, or the
machine spends a moment asking a dead port.
"""

from __future__ import annotations

import logging
import time
from types import ModuleType
from typing import Callable

from kadence_desktop.blocking import processes
from kadence_desktop.blocking.hosts import HostsBlocker
from kadence_desktop.blocking.resolver import ResolvedRoute, StubResolver, split_routable

log = logging.getLogger(__name__)


class BlockerManager:
    def __init__(
        self,
        *,
        helper_path: str,
        report: Callable[[str, str, str, str | None], None],
        notify: Callable[[str, str], None] | None = None,
        stub: StubResolver | None = None,
        route: ResolvedRoute | None = None,
        procs: ModuleType | None = None,
        report_cooldown: float = 20.0,
        kill_grace: float = 5.0,
    ) -> None:
        self._hosts = HostsBlocker(helper_path)
        # Injectable so tests can drive the ordering without binding ports or
        # asking a real systemd-resolved for anything.
        self._stub = stub if stub is not None else StubResolver()
        self._route = route if route is not None else ResolvedRoute(helper_path)
        # Injectable too, so the app-blocking tests do not have to own a victim
        # process — and because `/proc` is not something a test should need.
        self._procs = procs if procs is not None else processes
        self._report = report
        self._notify = notify
        self._active = False
        self._applied: list[str] = []
        #: The app list and kill/warn flag from the last policy, so the scan
        #: timer has something to work from between polls.
        self._apps: list[str] = []
        self._kill_apps = True
        #: What we have already told you about, keyed by (pid, start time) and
        #: stamped with when. It throttles the *report*, never the block — see
        #: [_enforce_apps].
        self._reported: dict[tuple[int, int], float] = {}
        #: When each process was asked politely to quit, so the next pass can
        #: insist. Keyed the same way.
        self._asked: dict[tuple[int, int], float] = {}
        self.report_cooldown = report_cooldown
        self.kill_grace = kill_grace
        self._redirect_ip = "127.0.0.1"
        self.enforce_hosts = True
        self.enforce_apps = True
        self.last_error: str | None = None
        self.route_error: str | None = None
        #: Patterns no machine can enforce at DNS level (`*gram*`); the UI says so.
        self.unroutable: list[str] = []
        #: Zones currently handed to `systemd-resolved`.
        self.routed_zones: list[str] = []
        self.killed_today = 0

        if self._route.applied is not None:
            # A drop-in survived a crash: the zones point at this stub, so it has
            # to be listening *now*, before the first policy even arrives.
            self._stub.start()

    # ---- state -------------------------------------------------------------

    @property
    def active(self) -> bool:
        return self._active

    @property
    def helper_installed(self) -> bool:
        return self._hosts.helper_installed

    @property
    def blocked_domains(self) -> list[str]:
        """What is really in `/etc/hosts` right now — not what the policy wants."""
        return list(self._applied)

    @property
    def error(self) -> str | None:
        """The first thing that went wrong, hosts or DNS, for the UI to show."""
        return self.last_error or self.route_error

    # ---- the one entry point ----------------------------------------------

    def sync(self, policy: dict) -> None:
        """Bring the machine in line with `policy` (from /api/blocker/policy)."""
        should_block = bool(policy.get("active")) and bool(policy.get("enabled"))
        domains = [str(item) for item in policy.get("domains", [])]
        apps = [str(item) for item in policy.get("apps", [])]
        self._redirect_ip = str(policy.get("redirect_ip") or "127.0.0.1")
        patterns = [str(item) for item in policy.get("unexpandable", [])]
        allow = [str(item) for item in policy.get("allowed", [])]

        routable, unroutable, zones = split_routable(patterns)
        self.unroutable = unroutable
        if not self.enforce_hosts:
            routable, zones = [], []
        self._sync_dns(routable, allow, zones, active=should_block)

        if should_block:
            if self.enforce_hosts and domains != self._applied:
                if self._hosts.apply(domains, self._redirect_ip):
                    self._applied = domains
                    self.last_error = None
                    log.info("blocking %s domain(s)", len(domains))
                else:
                    self.last_error = self._hosts.last_error
            self._apps = apps
            self._kill_apps = bool(policy.get("kill_apps", True))
            self._active = True
            self.enforce_apps_now()
        else:
            self.release()

    def _sync_dns(
        self,
        patterns: list[str],
        allow: list[str],
        zones: list[str],
        *,
        active: bool,
    ) -> None:
        """Put the stub and the resolved route in step, in a safe order."""
        if not zones:
            if self._route.applied is not None:
                self.route_error = None if self._route.off() else self._route.last_error
            else:
                self.route_error = None
            # Nothing points at the stub any more, so it can go.
            if self._stub.running:
                self._stub.stop()
            self.routed_zones = []
            return

        # The stub first. Writing the drop-in before anything listens on its port
        # would break resolution for the very zones we are trying to protect.
        if not self._stub.running and not self._stub.start():
            self.route_error = self._stub.error
            if self._route.applied is not None:
                self._route.off()
            self.routed_zones = []
            return
        if self._route.on(zones, self._stub.port or 0):
            self.route_error = None
        else:
            self.route_error = self._route.last_error
        self._stub.set_rules(patterns, allow, self._redirect_ip, active=active)
        self.routed_zones = zones

    def release(self) -> None:
        """Unblock everything. Called when the interval ends, and on shutdown."""
        if self._applied or self._hosts_dirty():
            if self._hosts.clear():
                self._applied = []
                self.last_error = None
            else:
                self.last_error = self._hosts.last_error
        self._active = False
        self._reported.clear()
        self._asked.clear()
        # The stub keeps running and keeps its route: the zones are still pointed
        # at it, so it has to answer — it just forwards everything again.
        self._stub.set_rules([], [], self._redirect_ip, active=False)

    def shutdown(self) -> None:
        """Give the machine's DNS back, then everything else. Called on exit."""
        if self._route.applied is not None and not self._route.off():
            self.route_error = self._route.last_error
        if self._stub.running:
            self._stub.stop()
        self.routed_zones = []
        self.release()

    def _hosts_dirty(self) -> bool:
        from kadence_desktop.blocking.hosts import current_block

        return bool(current_block())

    # ---- apps --------------------------------------------------------------

    def enforce_apps_now(self) -> None:
        """One app-blocking pass, from the timer that owns the cadence.

        Separate from [sync] because the two halves want different clocks: the
        hosts file only changes when the policy does, while a process can be
        launched at any moment and the server says how often to look
        (`blocker.app_poll_seconds`, arriving as `poll_seconds`).
        """
        if not (self._active and self.enforce_apps and self._apps):
            return
        self._enforce_apps(self._apps, kill=self._kill_apps)

    def _enforce_apps(self, patterns: list[str], *, kill: bool) -> None:
        """Close (or report) every blocked process running right now.

        The important word is *every*, on *every* pass. An earlier version kept a
        set of pids it had dealt with and skipped them forever, which conflated
        two different jobs: not spamming the daily report with one row per poll,
        and actually blocking the app. The second launch of a blocked app — or
        the first one whose SIGTERM was refused — then went unchallenged for the
        rest of the interval, which is the whole bug this docstring exists for.

        So the enforcement runs again each pass and only the *report* is
        throttled, per process per [report_cooldown]. "Per process" is
        (pid, start time), not pid: a pid alone is reused, and a relaunch
        landing on a recycled number must not inherit the old entry.

        The kill escalates across passes rather than inside one: SIGTERM now,
        SIGKILL once [kill_grace] has gone by and the process is still there.
        That is gentler than the old two-second busy-wait and it never blocks
        the thread drawing the countdown.
        """
        self._forget_dead()
        now = time.monotonic()
        for app in self._procs.scan(patterns):
            identity = app.identity
            if not kill:
                self._announce(app, "warned")
                continue

            asked_at = self._asked.get(identity)
            if asked_at is None:
                delivered = self._procs.terminate(app)
                self._asked[identity] = now
                if delivered:
                    self.killed_today += 1
                self._announce(app, "killed" if delivered else "warned")
            elif now - asked_at >= self.kill_grace:
                # It sat through the polite request. Insist, and keep insisting
                # one grace period at a time if something refuses the signal.
                self._procs.terminate(app, force=True)
                self._asked[identity] = now
                self._announce(app, "killed")

    def _announce(self, app: "processes.RunningApp", action: str) -> None:
        """Report and notify, at most once per process per cooldown."""
        if not self._claim_report(app.identity):
            return
        details = f"{app.name} (pid {app.pid})"
        self._report("app", app.matched, action, details)
        if not self._notify:
            return
        if action == "killed":
            self._notify("Closed a distraction", f"{app.name} is blocked during focus.")
        else:
            self._notify("That app is blocked", f"{app.name} is on your block list.")

    def _claim_report(self, identity: tuple[int, int]) -> bool:
        """Whether this process is due a report, and stamp it if so."""
        now = time.monotonic()
        last = self._reported.get(identity)
        if last is not None and now - last < self.report_cooldown:
            return False
        self._reported[identity] = now
        return True

    def _forget_dead(self) -> None:
        """Drop processes that are gone, so an all-day interval cannot leak."""
        for book in (self._reported, self._asked):
            for pid, started in list(book):
                if not self._procs.alive(pid, started):
                    del book[(pid, started)]