LibrePortal/scripts/dev/lp-shot

#!/usr/bin/env python3
"""lp-shot — headless screenshot of a LibrePortal WebUI route.

    lp-shot /admin/system                              # whole route -> /tmp/webui-shot.png
    lp-shot /admin/system /tmp/x.png 12 ".sys-strip"   # one element, 12px padding, crisp

Arguments (all optional after the route):
    route       WebUI path, e.g. /apps/overview  (a full http:// URL also works)
    out         output PNG                       (default /tmp/webui-shot.png)
    pad         padding in CSS px around the element clip   (default 0)
    selector    CSS selector — capture just that element    (default: full page)

Environment:
    LP_SHOT_URL       base URL of the WebUI      (default: auto-detected, else http://localhost:3179)
    LP_SHOT_VIEWPORT  WIDTHxHEIGHT               (default 1440x900)
    LP_SHOT_SCALE     device pixel ratio         (default 2 — that's the "crisp")
    LP_SHOT_SETTLE    extra seconds after load   (default 1.5)
    LP_SHOT_CHROME    chromium binary to use     (default: first found on PATH)

The WebUI is behind a login, so every route except / needs a session. On the host
that needs no setup: lp-shot signs one itself from the jwtSecret the backend keeps
in frontend/.auth.json, exactly as /api/auth/login would, and it expires in an hour.
No password is involved — the stored one is a bcrypt hash and is never touched.

Override that when running off-host, or against another instance:

    LP_SHOT_TOKEN     an existing `libreportal_token` cookie value
    LP_SHOT_USER      username, and
    LP_SHOT_PASS      password — posted once to /api/auth/login for a cookie
    LP_SHOT_AUTH_FILE path to a .auth.json to sign from
    LP_SHOT_VERBOSE   say which .auth.json the session was signed from

Drives a headless Chromium over the DevTools protocol: navigate, wait for the SPA
to paint, then capture. Element captures are clipped by the element's real box, so
they stay 1:1 with the page rather than being cropped out of a scaled screenshot.
Page console errors are echoed to stderr — a blank shot is usually a JS error.

DEV TOOL — not part of a release. scripts/dev is `export-ignore`d in
.gitattributes, so this never lands in a user's tarball. It reads the host's
.auth.json to sign itself a session, which is fine on a maintainer's box (you
already own that file) and has no business in a shipped install.

Requires: a chromium/chrome binary, and python3-websockets. On Ubuntu:
    sudo snap install chromium && sudo apt install python3-websockets
Install with:
    sudo install -m 755 scripts/dev/lp-shot /usr/local/bin/lp-shot
"""

import base64
import hashlib
import hmac
import json
import os
import re
import shutil
import socket
import subprocess
import sys
import tempfile
import time
import urllib.parse
import urllib.request

from websockets.sync.client import connect

DEFAULT_OUT = "/tmp/webui-shot.png"


def containers_root():
    """Where app data lives on THIS install, not where it lives by default.

    The containers root is relocatable (--containers-dir), and init.sh bakes the
    resolved value into the root-owned CLI wrapper. Reading it back from there
    is the only way to find the WebUI on an install whose data sits on another
    disk; hardcoding /libreportal-containers meant lp-shot silently fell back to
    a default port and a missing .auth.json, which looks exactly like a WebUI
    that failed to boot.
    """
    if os.environ.get("LP_CONTAINERS_DIR"):
        return os.environ["LP_CONTAINERS_DIR"].rstrip("/")
    try:
        with open("/usr/local/lib/libreportal/libreportal") as fh:
            m = re.search(r'^LP_CONTAINERS_DIR="([^"]+)"', fh.read(), re.M)
        if m and "__" not in m.group(1):
            return m.group(1).rstrip("/")
    except OSError:
        pass
    return "/libreportal-containers"


COMPOSE = containers_root() + "/libreportal/docker-compose.yml"


def die(msg, code=1):
    print(f"lp-shot: {msg}", file=sys.stderr)
    sys.exit(code)


def base_url():
    """Where the WebUI lives: env wins, else the live compose's published port."""
    if os.environ.get("LP_SHOT_URL"):
        return os.environ["LP_SHOT_URL"].rstrip("/")
    try:
        with open(COMPOSE) as fh:
            m = re.search(r'^\s*-\s*"(\d+):\d+"', fh.read(), re.M)
        if m:
            return f"http://localhost:{m.group(1)}"
    except OSError:
        pass
    return "http://localhost:3179"


COOKIE = "libreportal_token"


AUTH_FILES = [
    os.environ.get("LP_SHOT_AUTH_FILE"),
    containers_root() + "/libreportal/frontend/.auth.json",
    os.path.expanduser(
        "~/Documents/LibrePortal/LibrePortal/containers/libreportal/frontend/.auth.json"),
]


def jwt_hs256(payload, secret):
    """Sign a JWT the way the backend's jsonwebtoken does (HS256, compact form)."""
    def seg(raw):
        return base64.urlsafe_b64encode(raw).rstrip(b"=")
    head = seg(json.dumps({"alg": "HS256", "typ": "JWT"}, separators=(",", ":")).encode())
    body = seg(json.dumps(payload, separators=(",", ":")).encode())
    signed = head + b"." + body
    sig = hmac.new(secret.encode(), signed, hashlib.sha256).digest()
    return (signed + b"." + seg(sig)).decode()


def mint_token():
    """Sign our own session from the host's own JWT secret — no password needed.

    The backend keeps {username, passwordHash, jwtSecret} in frontend/.auth.json
    and mints session cookies as jwt.sign({sub: username}, jwtSecret). The
    password is bcrypt-hashed and unrecoverable, but the secret is right there in
    plaintext, so a tool running on the host can issue itself the same cookie the
    login endpoint would hand out. That's what makes lp-shot zero-config here.

    Deliberately short-lived: this token exists for one screenshot run, not as a
    standing credential.
    """
    for path in filter(None, AUTH_FILES):
        try:
            with open(path) as fh:
                data = json.load(fh)
        except (OSError, ValueError):
            continue
        secret, user = data.get("jwtSecret"), data.get("username")
        if not (secret and user):
            continue
        now = int(time.time())
        return jwt_hs256({"sub": user, "iat": now, "exp": now + 3600}, secret), path
    return None, None


def session_token(base):
    """A `libreportal_token` value, or None if we could not get one.

    Three ways, in order: handed to us (LP_SHOT_TOKEN), exchanged for one at the
    login endpoint using credentials in the environment, or — the usual case on
    the host itself — signed locally from .auth.json's jwtSecret. Nothing is
    cached: the cookie is injected for this run and the profile is thrown away.
    """
    tok = os.environ.get("LP_SHOT_TOKEN")
    if tok:
        return tok.strip()
    user, pw = os.environ.get("LP_SHOT_USER"), os.environ.get("LP_SHOT_PASS")
    if not (user and pw):
        tok, src = mint_token()
        if tok:
            if os.environ.get("LP_SHOT_VERBOSE"):
                print(f"lp-shot: signed a session from {src}", file=sys.stderr)
            return tok
        return None
    body = json.dumps({"username": user, "password": pw}).encode()
    req = urllib.request.Request(f"{base}/api/auth/login", data=body,
                                 headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=15) as resp:
            for header, value in resp.getheaders():
                if header.lower() == "set-cookie" and value.startswith(COOKIE + "="):
                    return value.split("=", 1)[1].split(";")[0]
    except urllib.error.HTTPError as e:
        detail = "invalid credentials" if e.code == 401 else \
                 "rate-limited, wait it out" if e.code == 429 else f"HTTP {e.code}"
        die(f"login as {user!r} failed: {detail}")
    except urllib.error.URLError as e:
        die(f"cannot reach {base}: {e.reason}")
    die("login succeeded but returned no session cookie")


def check_token(base, token):
    req = urllib.request.Request(f"{base}/api/auth/status",
                                 headers={"Cookie": f"{COOKIE}={token}"})
    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            return bool(json.loads(resp.read()).get("authenticated"))
    except (urllib.error.URLError, ValueError):
        return False


LOGIN_PROBE = "!!document.getElementById('login-form')"

AUTH_HELP = (
    "the WebUI needs a session, and none could be signed locally.\n"
    "  - on the host: check .auth.json is readable (tried: %s)\n"
    "  - elsewhere:   export LP_SHOT_TOKEN=... (the libreportal_token cookie)\n"
    "                 or LP_SHOT_USER=admin LP_SHOT_PASS=...\n"
    "See `lp-shot --help`." % ", ".join(p for p in AUTH_FILES if p)
)


def find_chrome():
    env = os.environ.get("LP_SHOT_CHROME")
    if env:
        return env
    for name in ("chromium", "chromium-browser", "google-chrome", "google-chrome-stable", "chrome"):
        path = shutil.which(name)
        if path:
            return path
    die("no chromium found — install one (`sudo snap install chromium`) or set LP_SHOT_CHROME")


def profile_dir(chrome):
    """A throwaway user-data-dir the browser can actually write to.

    The snap build is confined: it cannot see /tmp or dot-directories in $HOME,
    so park the profile under its own SNAP_USER_COMMON. Everything else gets a
    normal temp dir.

    One profile per run, never a shared path: chromium refuses to start on a
    directory another instance still holds, and a snap-confined browser can't be
    killed from outside (AppArmor drops the signal even for root), so a single
    wedged run would otherwise brick the tool until someone rebooted.
    """
    if "/snap/" in os.path.realpath(chrome) or chrome.startswith("/snap/"):
        home = os.path.expanduser("~/snap/chromium/common")
        os.makedirs(home, exist_ok=True)
        sweep_stale(home)
        return tempfile.mkdtemp(prefix="lp-shot-", dir=home)
    return tempfile.mkdtemp(prefix="lp-shot-")


def sweep_stale(home):
    """Drop lp-shot profiles left behind by runs that died mid-flight."""
    cutoff = time.time() - 3600
    try:
        for name in os.listdir(home):
            path = os.path.join(home, name)
            if name.startswith("lp-shot-") and os.path.isdir(path) \
                    and os.path.getmtime(path) < cutoff:
                shutil.rmtree(path, ignore_errors=True)
    except OSError:
        pass


def launch(chrome, width, height):
    prof = profile_dir(chrome)
    port_file = os.path.join(prof, "DevToolsActivePort")
    try:
        os.remove(port_file)
    except OSError:
        pass
    proc = subprocess.Popen(
        [
            chrome,
            "--headless=new",
            "--remote-debugging-port=0",   # port 0 -> the real one lands in DevToolsActivePort
            f"--user-data-dir={prof}",
            f"--window-size={width},{height}",
            "--disable-gpu",
            "--hide-scrollbars",
            "--no-first-run",
            "--no-default-browser-check",
            "--disable-extensions",
            "--disable-background-networking",
            "--disable-features=Translate,MediaRouter",
            "about:blank",
        ],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    for _ in range(300):  # up to 30s — a cold snap start is slow
        if proc.poll() is not None:
            die("chromium exited before it opened a debugging port")
        try:
            with open(port_file) as fh:
                lines = fh.read().split("\n")
            if lines and lines[0].strip():
                return proc, prof, int(lines[0].strip())
        except OSError:
            pass
        time.sleep(0.1)
    proc.kill()
    die("timed out waiting for chromium's debugging port")


def page_socket(port):
    for _ in range(100):
        try:
            raw = urllib.request.urlopen(f"http://127.0.0.1:{port}/json/list", timeout=2).read()
            for t in json.loads(raw):
                if t.get("type") == "page" and t.get("webSocketDebuggerUrl"):
                    return t["webSocketDebuggerUrl"]
        except (urllib.error.URLError, socket.timeout, ValueError):
            pass
        time.sleep(0.1)
    die("chromium never exposed a page target")


class CDP:
    def __init__(self, ws_url):
        # max_size=None: a full-page screenshot is far past the 1MB default frame cap.
        self.ws = connect(ws_url, max_size=None, open_timeout=20)
        self.n = 0
        self.events = []

    def send(self, method, timeout=60, **params):
        self.n += 1
        self.ws.send(json.dumps({"id": self.n, "method": method, "params": params}))
        deadline = time.time() + timeout
        while time.time() < deadline:
            msg = json.loads(self.ws.recv(timeout=max(1, deadline - time.time())))
            if msg.get("id") == self.n:
                if "error" in msg:
                    die(f"{method}: {msg['error'].get('message')}")
                return msg.get("result", {})
            if "method" in msg:
                self.events.append(msg)
        die(f"{method}: timed out")

    def drain(self):
        """Collect any events that arrived while we weren't listening."""
        while True:
            try:
                msg = json.loads(self.ws.recv(timeout=0.05))
            except Exception:
                return
            if "method" in msg:
                self.events.append(msg)

    def eval(self, expr):
        r = self.send("Runtime.evaluate", expression=expr, returnByValue=True, awaitPromise=True)
        if r.get("exceptionDetails"):
            return None
        return r.get("result", {}).get("value")

    def close(self):
        try:
            self.ws.close()
        except Exception:
            pass


def main():
    # Print a session cookie and exit. A screenshot is enough for "does it
    # render", but not for "does this wizard step work" — that needs clicking,
    # which means driving a real browser, which needs the same session lp-shot
    # already knows how to mint. Without this the only way in is typing the
    # admin password into the login form.
    #
    #   lp-shot --token            -> the raw cookie VALUE
    #   lp-shot --cookie-js        -> a document.cookie assignment to paste/eval
    if len(sys.argv) > 1 and sys.argv[1] in ("--token", "--cookie-js"):
        token, src = mint_token()
        if os.environ.get("LP_SHOT_VERBOSE"):
            print(f"signed from {src}", file=sys.stderr)
        if sys.argv[1] == "--token":
            print(token)
        else:
            print(f'document.cookie = "{COOKIE}={token}; path=/"')
        sys.exit(0)

    if len(sys.argv) < 2 or sys.argv[1] in ("-h", "--help"):
        print(__doc__.strip())
        sys.exit(0 if len(sys.argv) > 1 else 2)

    route = sys.argv[1]
    out = os.path.abspath(sys.argv[2]) if len(sys.argv) > 2 else DEFAULT_OUT
    pad = float(sys.argv[3]) if len(sys.argv) > 3 else 0.0
    selector = sys.argv[4] if len(sys.argv) > 4 else None

    base = base_url()
    url = route if re.match(r"^https?://", route) else base + "/" + route.lstrip("/")
    vp = os.environ.get("LP_SHOT_VIEWPORT", "1440x900")
    try:
        width, height = (int(x) for x in vp.lower().split("x"))
    except ValueError:
        die(f"bad LP_SHOT_VIEWPORT {vp!r} — want WIDTHxHEIGHT")
    scale = float(os.environ.get("LP_SHOT_SCALE", "2"))
    settle = float(os.environ.get("LP_SHOT_SETTLE", "1.5"))

    token = session_token(base)
    if token and not check_token(base, token):
        die("that session is not valid (expired token, or wrong LP_SHOT_URL)")

    chrome = find_chrome()
    proc, prof, port = launch(chrome, width, height)
    cdp = None
    try:
        cdp = CDP(page_socket(port))
        cdp.send("Page.enable")
        cdp.send("Runtime.enable")
        cdp.send("Log.enable")
        cdp.send("Network.enable")
        if token:
            host = re.sub(r"^https?://", "", base).split(":")[0].split("/")[0]
            cdp.send("Network.setCookie", name=COOKIE, value=token, domain=host,
                     path="/", httpOnly=True, sameSite="Strict")
        # Pin the metrics rather than trusting --window-size: headless sizes the
        # window including chrome-less padding, and we want an exact CSS viewport.
        cdp.send("Emulation.setDeviceMetricsOverride",
                 width=width, height=height, deviceScaleFactor=scale, mobile=False)

        cdp.send("Page.navigate", url=url, timeout=60)
        deadline = time.time() + 30
        while time.time() < deadline:
            if any(e["method"] == "Page.loadEventFired" for e in cdp.events):
                break
            cdp.drain()
            time.sleep(0.1)

        # The WebUI is an SPA: the document loads long before the route paints.
        # Wait for the thing we're actually capturing (or for a non-empty body),
        # then settle for late renders/fonts.
        # The SPA shows a full-screen boot loader (#libreportal-loading-screen,
        # removed from the DOM once it finishes) over an otherwise-ready page.
        # Waiting only for "body has text" happily captures that splash at 9%,
        # so every probe is gated on the loader being gone first.
        want = f"!!document.querySelector({json.dumps(selector)})" if selector else \
               "!!document.body && document.body.innerText.trim().length > 0"
        probe = f"(!document.getElementById('libreportal-loading-screen')) && ({want})"
        # An unauthenticated run lands on the sign-in box: without a check it
        # would either quietly screenshot that, or sit here until the selector
        # times out — and neither says what's actually wrong. But "login form on
        # screen" isn't itself the failure (you may be shooting the login page,
        # or an element inside it), so it only counts once what was ASKED for
        # has failed to show: an unsatisfied probe on a route that isn't the
        # login route, or a selector that never turned up.
        root_route = urllib.parse.urlparse(url).path.rstrip("/") in ("", "/")
        deadline = time.time() + 45   # cold SPA boot is slow
        while time.time() < deadline:
            if cdp.eval(probe):
                break
            if not selector and not root_route and cdp.eval(LOGIN_PROBE):
                die(AUTH_HELP)
            time.sleep(0.2)
        else:
            if selector and cdp.eval(LOGIN_PROBE):
                die(AUTH_HELP)
            if selector:
                die(f"selector {selector!r} never appeared on {url}")
        if not selector and not root_route and cdp.eval(LOGIN_PROBE):
            die(AUTH_HELP)
        time.sleep(settle)

        for e in cdp.events:
            if e["method"] == "Log.entryAdded" and e["params"]["entry"].get("level") == "error":
                print(f"  page error: {e['params']['entry'].get('text')}", file=sys.stderr)

        if selector:
            box = cdp.eval(
                "(() => { const el = document.querySelector(%s); if (!el) return null;"
                " el.scrollIntoView({block:'center', inline:'center', behavior:'instant'});"
                " const r = el.getBoundingClientRect();"
                " return {x: r.left + window.scrollX, y: r.top + window.scrollY,"
                "         w: r.width, h: r.height}; })()" % json.dumps(selector)
            )
            if not box:
                die(f"selector {selector!r} matched nothing on {url}")
            if box["w"] < 1 or box["h"] < 1:
                die(f"selector {selector!r} matched a zero-size element (hidden?)")
            clip = {
                "x": max(0.0, box["x"] - pad),
                "y": max(0.0, box["y"] - pad),
                "width": box["w"] + pad * 2,
                "height": box["h"] + pad * 2,
                "scale": scale,
            }
        else:
            m = cdp.send("Page.getLayoutMetrics")
            size = m.get("cssContentSize") or m.get("contentSize")
            clip = {
                "x": 0.0, "y": 0.0,
                "width": float(size["width"]),
                "height": min(float(size["height"]), 20000.0),
                "scale": scale,
            }

        shot = cdp.send("Page.captureScreenshot", format="png", clip=clip,
                        captureBeyondViewport=True, fromSurface=True, timeout=90)
        data = base64.b64decode(shot["data"])
        os.makedirs(os.path.dirname(out) or ".", exist_ok=True)
        with open(out, "wb") as fh:
            fh.write(data)
        print(f"{out}  ({int(clip['width'])}x{int(clip['height'])} css @{scale:g}x, "
              f"{len(data) // 1024} KB)  {url}")
    finally:
        # Shut the browser down through the protocol, not with a signal: the snap
        # launcher is setuid-root, so the PID we hold is root-owned and os.kill
        # comes back EPERM. Browser.close is the only teardown that actually
        # works there; the signals are just a backstop for an unconfined chrome.
        if cdp:
            try:
                cdp.send("Browser.close", timeout=5)
            except SystemExit:
                pass
            cdp.close()
        try:
            proc.wait(timeout=10)
        except subprocess.TimeoutExpired:
            for stop in (proc.terminate, proc.kill):
                try:
                    stop()
                    proc.wait(timeout=5)
                    break
                except (PermissionError, subprocess.TimeoutExpired):
                    continue
        shutil.rmtree(prof, ignore_errors=True)


if __name__ == "__main__":
    main()