omarchy-screencast/daemon/screencast/util.py

"""Small helpers shared by the daemon: logging, paths, local addresses, subprocess."""
from __future__ import annotations

import asyncio
import contextlib
import ipaddress
import json
import logging
import logging.handlers
import os
import shutil
import stat
import socket
import subprocess
from pathlib import Path

log = logging.getLogger("screencast")
_firewall_tool: str | None = None


def setup_logging(verbose: bool = False) -> None:
    logging.basicConfig(
        level=logging.DEBUG if verbose else logging.INFO,
        format="%(asctime)s %(levelname).1s %(name)s: %(message)s",
        datefmt="%H:%M:%S",
    )
    # The shell owns the daemon's stderr and keeps only the last few lines in a
    # QML property, so also keep a real log on disk: when a cast misbehaves this
    # file is the only account of what the encoder and the receiver did.
    try:
        handler = logging.handlers.RotatingFileHandler(
            state_dir() / "daemon.log", maxBytes=512_000, backupCount=1
        )
        handler.setFormatter(logging.Formatter(
            "%(asctime)s %(levelname).1s %(name)s: %(message)s", "%Y-%m-%d %H:%M:%S"))
        logging.getLogger().addHandler(handler)
    except OSError as exc:
        log.warning("no log file: %s", exc)
    # These three are chatty at DEBUG and none of it is ours.
    for noisy in ("zeroconf", "pychromecast", "aiohttp.access", "dbus_fast"):
        logging.getLogger(noisy).setLevel(logging.WARNING)


def firewall_tool(refresh: bool = False) -> str:
    """Which host firewall is running, if any ("ufw", "firewalld" or "").

    A receiver fetches the stream from us, so the host firewall has to let it
    in. Omarchy installs ufw with `default deny incoming`, which drops the TV's
    SYN silently: the TV shows a spinner and we never see a client at all.
    Reading the rules needs root, so all we can tell is which firewall is up.
    """
    global _firewall_tool
    if _firewall_tool is not None and not refresh:
        return _firewall_tool
    _firewall_tool = ""
    for unit in ("ufw", "firewalld"):
        try:
            out = subprocess.run(["systemctl", "is-active", unit],
                                 capture_output=True, text=True, timeout=3)
        except Exception:  # noqa: BLE001
            continue
        if out.stdout.strip() == "active":
            _firewall_tool = unit
            break
    return _firewall_tool


def firewall_argv(tool: str, port: int, host: str = "") -> list[list[str]]:
    """The privileged command(s) that open the stream port to the local network."""
    if tool == "ufw":
        return [["ufw", "allow", "from", _subnet_of(host), "to", "any",
                 "port", str(port), "proto", "tcp", "comment", "screencast"]]
    if tool == "firewalld":
        # Scoped to the LAN like the ufw rule above: --add-port alone would open
        # this port to every network the machine ever joins, permanently.
        rule = (f'rule family="ipv{4 if ":" not in _subnet_of(host) else 6}" '
                f'source address="{_subnet_of(host)}" '
                f'port port="{port}" protocol="tcp" accept')
        return [["firewall-cmd", f"--add-rich-rule={rule}", "--permanent"],
                ["firewall-cmd", "--reload"]]
    return []


def firewall_command(tool: str, port: int, host: str = "") -> str:
    """The same thing as one line a person can paste into a terminal."""
    parts = [" ".join(["sudo", *argv]) for argv in firewall_argv(tool, port, host)]
    return " && ".join(parts)


def terminal_argv(command: list[str]) -> list[str] | None:
    """Run `command` in the user's terminal, for the times pkexec cannot ask.

    Not every desktop runs a polkit authentication agent - Omarchy does not
    install one - and without an agent pkexec has nowhere to prompt. A terminal
    window with sudo in it needs nothing installed and shows the user exactly
    what is being run.
    """
    for candidate in (os.environ.get("TERMINAL"), "xdg-terminal-exec", "ghostty",
                      "alacritty", "kitty", "foot", "wezterm", "xterm"):
        if not candidate:
            continue
        exe = shutil.which(candidate)
        if exe is None:
            continue
        name = Path(exe).name
        if name in ("xdg-terminal-exec", "foot", "kitty", "ghostty"):
            return [exe, *command]      # these take the command as plain arguments
        if name == "wezterm":
            return [exe, "start", "--", *command]
        return [exe, "-e", *command]    # alacritty, xterm and the rest
    return None


def firewall_hint(port: int, host: str = "") -> str:
    tool = firewall_tool()
    if not tool:
        return ""
    return (f"{tool} is blocking port {port} — receivers cannot fetch the stream "
            f"until it is opened")


def _subnet_of(host: str) -> str:
    """The /24 the receiver lives on, so the opened port stays on the LAN."""
    # `host` comes from a discovery announcement, which anyone on the wire can
    # send: a device claiming a public address would otherwise pick the subnet
    # in a rule we are about to install as root. Only a private address gets to
    # choose; otherwise ask the kernel which address the default route leaves
    # from, which is the LAN the receivers actually live on.
    try:
        if host and not ipaddress.ip_address(host).is_private:
            host = ""
    except ValueError:
        host = ""
    try:
        addr = ipaddress.ip_address(local_ip_for(host or "1.1.1.1"))
    except ValueError:
        return "192.168.0.0/16"
    if not addr.is_private:
        return "192.168.0.0/16"
    if isinstance(addr, ipaddress.IPv4Address):
        return str(ipaddress.ip_network(f"{addr}/24", strict=False))
    return str(ipaddress.ip_network(f"{addr}/64", strict=False))


def runtime_dir() -> Path:
    """Where the control socket, the preview frame and HLS segments live.

    Without XDG_RUNTIME_DIR this falls back to /tmp, which is shared and
    predictable: another local user could pre-create the directory, or plant a
    symlink where the preview frame is about to be written. So refuse anything
    we do not own, and keep it to ourselves.
    """
    base = os.environ.get("XDG_RUNTIME_DIR") or "/tmp"
    p = Path(base) / "screencast"
    p.mkdir(parents=True, exist_ok=True, mode=0o700)
    st = p.lstat()
    if not stat.S_ISDIR(st.st_mode) or st.st_uid != os.getuid():
        raise RuntimeError(f"{p} is not ours — refusing to use it")
    if st.st_mode & 0o077:
        p.chmod(0o700)
    return p


def config_dir() -> Path:
    base = os.environ.get("XDG_CONFIG_HOME") or (Path.home() / ".config")
    p = Path(base) / "screencast"
    p.mkdir(parents=True, exist_ok=True)
    return p


def state_dir() -> Path:
    """Logs and, for AirPlay, the receiver pairing credentials — so 0700."""
    base = os.environ.get("XDG_STATE_HOME") or (Path.home() / ".local/state")
    p = Path(base) / "screencast"
    p.mkdir(parents=True, exist_ok=True, mode=0o700)
    if p.stat().st_mode & 0o077:
        p.chmod(0o700)
    return p


def read_json(path: Path, default):
    try:
        return json.loads(path.read_text())
    except Exception:
        return default


def write_json(path: Path, data) -> None:
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(data, indent=2))
    tmp.replace(path)


def local_ip_for(host: str) -> str:
    """The address a receiver at `host` would see us as.

    The machine can sit on several networks at once (several NICs, a VPN), so
    the stream URL must carry the address that routes back to that particular
    device: let the kernel pick it.
    """
    fam = socket.AF_INET
    try:
        if isinstance(ipaddress.ip_address(host), ipaddress.IPv6Address):
            fam = socket.AF_INET6
    except ValueError:
        pass
    s = socket.socket(fam, socket.SOCK_DGRAM)
    try:
        s.connect((host, 9))
        ip = s.getsockname()[0]
    except OSError:
        ip = "127.0.0.1"
    finally:
        s.close()
    return ip


def url_host(ip: str) -> str:
    return f"[{ip}]" if ":" in ip else ip


async def run_cmd(*args: str, timeout: float = 5.0) -> tuple[int, str, str]:
    proc = await asyncio.create_subprocess_exec(
        *args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
    )
    try:
        out, err = await asyncio.wait_for(proc.communicate(), timeout)
    except asyncio.TimeoutError:
        proc.kill()
        return 124, "", "timeout"
    return proc.returncode or 0, out.decode(errors="replace"), err.decode(errors="replace")


def have(binary: str) -> bool:
    from shutil import which

    return which(binary) is not None


async def port_open(host: str, port: int, timeout: float = 2.0) -> bool:
    """Is anything listening there? Distinguishes "asleep" from "not running"."""
    try:
        reader, writer = await asyncio.wait_for(asyncio.open_connection(host, port), timeout)
    except Exception:  # noqa: BLE001
        return False
    writer.close()
    with contextlib.suppress(Exception):
        await writer.wait_closed()
    return True


def gst_has(element: str) -> bool:
    """Is this GStreamer element installed? (cached; a plugin set never changes mid-run)"""
    if element in _gst_cache:
        return _gst_cache[element]
    ok = False
    try:
        ok = subprocess.run(
            ["gst-inspect-1.0", "--exists", element], timeout=10
        ).returncode == 0
    except Exception:
        ok = False
    _gst_cache[element] = ok
    return ok


_gst_cache: dict[str, bool] = {}