dashboard/scripts/screenshot.py

"""Capture per-tab screenshots of the dashboard with headless Chromium.

A maintenance tool, not an application dependency: it drives Chromium over
the Chrome DevTools Protocol using the ``websocket-client`` package, so it
must be run with a Python that has it installed (the system ``python3`` has
it; the project venv does not) and a ``chromium`` (or ``chromium-browser`` /
``google-chrome``) binary on PATH. The dashboard server must already be
running, because the screenshots show its live state.

Usage:

    python3 scripts/screenshot.py [--url URL] [--out DIR] [--date YYYYMMDD]
        [--width N] [--height N] [--tabs overview,disks,...]

For every selected tab the script clicks the tab button, waits until the
tab's content has been fetched and rendered, then stores a full-page PNG at
``<out>/<Tab>_<date>.png`` (e.g. ``screenshots/Overview_20260831.png``),
matching the filenames referenced in the README.
"""

import argparse
import base64
import json
import shutil
import socket
import subprocess
import sys
import tempfile
import time
import urllib.request
from pathlib import Path
from typing import Any

import websocket

TABS = ("overview", "disks", "processes", "journal", "services", "plugins")

# JavaScript readiness conditions, one per tab: a tab is captured once its
# expression evaluates truthy. The journal waits for real lines (the error
# banner also has the jline class), and the overview extra settle below gives
# the 2 s /api/history poller time to fill the charts.
WAIT: dict[str, str] = {
    "overview": "document.querySelectorAll('#overview .card').length >= 1",
    "disks": "document.querySelectorAll('#disks .part').length >= 1",
    "processes": "document.querySelectorAll('#processes-body .table tr').length >= 1",
    "journal": "document.querySelectorAll('#journal-log .jline:not(.j-error-once)').length >= 5",
    "services": "document.querySelectorAll('#services tr.svc-row').length >= 1",
    "plugins": "document.querySelectorAll('#plugins .plugin-card').length >= 1",
}

# Extra seconds to wait after the readiness condition, per tab; the overview
# needs two history polls so its charts are not empty.
SETTLE: dict[str, float] = {
    "overview": 6.0,
}

WAIT_TIMEOUT = 20.0
POLL_INTERVAL = 0.25


class Cdp:
    """Minimal Chrome DevTools Protocol client over a WebSocket.

    Only what the capture flow needs: command/response correlation
    (events are discarded), JavaScript evaluation, and truthy-condition
    waiting.
    """

    def __init__(self, url: str) -> None:
        """Connect to a DevTools WebSocket endpoint.

        Args:
            url: The target's webSocketDebuggerUrl.

        Raises:
            websocket.WebSocketException: If the connection fails.
        """
        self._ws = websocket.create_connection(url, timeout=30)
        self._next_id = 0

    def send(self, method: str, params: dict[str, Any] | None = None) -> dict[str, Any]:
        """Send a CDP command and wait for its response.

        Args:
            method: CDP method name, e.g. "Page.navigate".
            params: Optional command parameters.

        Returns:
            The "result" object of the response ({} if absent).

        Raises:
            RuntimeError: If the response carries a protocol error.
        """
        self._next_id += 1
        message: dict[str, Any] = {"id": self._next_id, "method": method}
        if params:
            message["params"] = params
        self._ws.send(json.dumps(message))
        while True:
            data = json.loads(self._ws.recv())
            if data.get("id") != self._next_id:
                continue
            if "error" in data:
                raise RuntimeError(f"{method} failed: {data['error']}")
            return data.get("result", {})

    def evaluate(self, expression: str) -> Any:
        """Evaluate JavaScript in the page and return its value.

        Args:
            expression: JavaScript expression, evaluated by value.

        Raises:
            RuntimeError: If evaluation throws in the page.
        """
        result = self.send(
            "Runtime.evaluate",
            {"expression": expression, "returnByValue": True},
        )
        if "exceptionDetails" in result:
            detail = result["exceptionDetails"]
            raise RuntimeError(f"JS exception: {detail.get('text', detail)}")
        return result.get("result", {}).get("value")

    def wait_until(self, expression: str, timeout: float) -> None:
        """Wait until a JavaScript expression evaluates truthy.

        Args:
            expression: Expression polled every POLL_INTERVAL seconds.
            timeout: Maximum seconds to wait.

        Raises:
            TimeoutError: If the condition is not met in time.
        """
        deadline = time.monotonic() + timeout
        last: Any = None
        while time.monotonic() < deadline:
            last = self.evaluate(expression)
            if last:
                return
            time.sleep(POLL_INTERVAL)
        raise TimeoutError(
            f"condition not met within {timeout:.0f}s: {expression} (last: {last!r})"
        )

    def close(self) -> None:
        """Close the WebSocket, ignoring errors."""
        try:
            self._ws.close()
        except Exception:
            pass


def free_port() -> int:
    """Return an unused localhost TCP port.

    Returns:
        A port number that was free at call time.
    """
    with socket.socket() as sock:
        sock.bind(("127.0.0.1", 0))
        return sock.getsockname()[1]


def launch_chromium(port: int, workdir: Path) -> subprocess.Popen:
    """Start headless Chromium with remote debugging enabled.

    Args:
        port: TCP port for the DevTools endpoint.
        workdir: Temporary directory; the profile and log file go inside it.

    Returns:
        The running Popen handle.

    Raises:
        RuntimeError: If no supported browser binary is on PATH.
    """
    binary = (
        shutil.which("chromium")
        or shutil.which("chromium-browser")
        or shutil.which("google-chrome")
    )
    if binary is None:
        raise RuntimeError("no chromium/chrome binary found on PATH")
    log = open(workdir / "chromium.log", "wb")
    return subprocess.Popen(
        [
            binary,
            "--headless=new",
            "--no-first-run",
            "--no-default-browser-check",
            "--disable-gpu",
            "--hide-scrollbars",
            "--force-device-scale-factor=1",
            f"--user-data-dir={workdir / 'profile'}",
            f"--remote-debugging-port={port}",
            "--remote-allow-origins=*",
            "about:blank",
        ],
        stdout=log,
        stderr=subprocess.STDOUT,
    )


def devtools_url(port: int, timeout: float = 30.0) -> str:
    """Wait for Chromium and return the page target's DevTools WebSocket URL.

    Args:
        port: DevTools TCP port.
        timeout: Maximum seconds to wait.

    Returns:
        The webSocketDebuggerUrl of the first page target.

    Raises:
        TimeoutError: If no page target appears in time.
    """
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            with urllib.request.urlopen(f"http://127.0.0.1:{port}/json", timeout=2) as resp:
                targets = json.load(resp)
        except (OSError, ValueError):
            targets = []
        for target in targets:
            if target.get("type") == "page" and target.get("webSocketDebuggerUrl"):
                return target["webSocketDebuggerUrl"]
        time.sleep(POLL_INTERVAL)
    raise TimeoutError(f"DevTools page target not ready on port {port}")


def capture(url: str, out: Path, date: str, width: int, height: int, tabs: list[str]) -> list[Path]:
    """Capture one full-page screenshot per tab.

    Args:
        url: Dashboard base URL.
        out: Output directory for the PNG files.
        date: Date stamp embedded in the filenames (YYYYMMDD).
        width: Viewport width in CSS pixels.
        height: Viewport height in CSS pixels.
        tabs: Tab names to capture, in order.

    Returns:
        The written file paths, in capture order.

    Raises:
        OSError: If the browser or network fails.
        TimeoutError: If the browser or a tab does not become ready.
        RuntimeError: If a CDP command or page script fails.
    """
    workdir = Path(tempfile.mkdtemp(prefix="dashshot-"))
    port = free_port()
    browser = launch_chromium(port, workdir)
    try:
        cdp = Cdp(devtools_url(port))
        cdp.send("Page.enable")
        cdp.send(
            "Emulation.setDeviceMetricsOverride",
            {"width": width, "height": height, "deviceScaleFactor": 1, "mobile": False},
        )
        cdp.send("Page.navigate", {"url": url})
        cdp.wait_until("document.readyState === 'complete'", 30.0)
        written: list[Path] = []
        for tab in tabs:
            cdp.evaluate(f'document.querySelector(`.tab-btn[data-tab="{tab}"]`).click()')
            cdp.wait_until(WAIT[tab], WAIT_TIMEOUT)
            time.sleep(SETTLE.get(tab, 0.4))
            shot = cdp.send(
                "Page.captureScreenshot",
                {"format": "png", "captureBeyondViewport": True},
            )
            path = out / f"{tab.capitalize()}_{date}.png"
            path.write_bytes(base64.b64decode(shot["data"]))
            print(f"wrote {path}")
            written.append(path)
        cdp.close()
        return written
    finally:
        browser.terminate()
        try:
            browser.wait(timeout=10)
        except subprocess.TimeoutExpired:
            browser.kill()
        shutil.rmtree(workdir, ignore_errors=True)


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    """Parse command line arguments.

    Args:
        argv: Argument list (defaults to sys.argv[1:]).

    Returns:
        The parsed namespace.
    """
    parser = argparse.ArgumentParser(
        description="Capture per-tab dashboard screenshots with headless Chromium."
    )
    parser.add_argument("--url", default="http://127.0.0.1:8501", help="dashboard base URL")
    parser.add_argument(
        "--out",
        type=Path,
        default=Path(__file__).resolve().parent.parent / "screenshots",
        help="output directory for the PNG files",
    )
    parser.add_argument("--date", default=time.strftime("%Y%m%d"), help="date stamp for filenames")
    parser.add_argument("--width", type=int, default=1671, help="viewport width in CSS pixels")
    parser.add_argument("--height", type=int, default=707, help="viewport height in CSS pixels")
    parser.add_argument("--tabs", default=",".join(TABS), help="comma-separated tab names")
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    """Entry point.

    Args:
        argv: Argument list (defaults to sys.argv[1:]).

    Returns:
        Process exit code: 0 on success, 1 on runtime error, 2 on bad input.
    """
    args = parse_args(argv)
    tabs = [t.strip() for t in args.tabs.split(",") if t.strip()]
    unknown = [t for t in tabs if t not in TABS]
    if unknown:
        print(f"unknown tabs: {', '.join(unknown)}", file=sys.stderr)
        return 2
    try:
        capture(args.url, args.out, args.date, args.width, args.height, tabs)
    except (OSError, TimeoutError, RuntimeError, ValueError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())