rig-pocket/server.py

#!/usr/bin/env python3
"""rig-pocket: a phone-friendly status page and push notifier for OpenRig.

Reads state only through the documented `rig` CLI (`--json`), so it needs no
access to the daemon's internals. Standard library only.
"""
import glob
import json
import os
import re
import subprocess
import threading
import time
import urllib.parse
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

HERE = Path(__file__).resolve().parent
STATIC = HERE / "static"

RIG = os.environ.get("RIG_BIN", "rig")
BIND = os.environ.get("RIG_POCKET_BIND", "127.0.0.1")
PORT = int(os.environ.get("RIG_POCKET_PORT", "7440"))
POLL = int(os.environ.get("RIG_POCKET_POLL", "30"))
REPOS = os.environ.get("RIG_POCKET_REPOS", "")  # space separated globs of git work trees
ORIGIN = os.environ.get("RIG_POCKET_ORIGIN", "")  # e.g. https://rig.example.com, checked on POST
PUBLIC_URL = os.environ.get("RIG_POCKET_URL", "")  # where a notification tap should go
THEME_CSS = os.environ.get("RIG_POCKET_THEME_CSS", "https://assets.weircon.dk/css/weirsoe-theme.css")
THEME_JS = os.environ.get("RIG_POCKET_THEME_JS", "https://assets.weircon.dk/js/weirsoe-theme.js")
NTFY_URL = os.environ.get("NTFY_URL", "")
NTFY_TOPIC = os.environ.get("NTFY_TOPIC", "")
NTFY_TOKEN = os.environ.get("NTFY_TOKEN", "")
DONE_DEST = re.compile(os.environ.get("RIG_POCKET_DONE_DEST", r"^orch-lead@"))
HUMAN_DEST = re.compile(os.environ.get("RIG_POCKET_HUMAN_DEST", r"^operator-agent@kernel$"))

PROMPT = re.compile(r"Do you want to proceed\?|^\s*❯ 1\. ", re.M)
SEAT = re.compile(r"^[A-Za-z0-9._-]+@[A-Za-z0-9._-]+$")
KEYS = {"1", "2", "3", "4"}
CLOSED = {"done", "closed", "failed", "canceled", "cancelled"}


def rig(*args, timeout=20):
    try:
        out = subprocess.run([RIG, *args], capture_output=True, text=True, timeout=timeout)
        return out.stdout
    except (OSError, subprocess.TimeoutExpired):
        return ""


def rig_json(*args):
    try:
        return json.loads(rig(*args) or "null")
    except json.JSONDecodeError:
        return None


def git(path, *args):
    try:
        return subprocess.run(["git", "-C", path, *args], capture_output=True, text=True, timeout=10).stdout
    except (OSError, subprocess.TimeoutExpired):
        return ""


def as_list(data):
    if isinstance(data, list):
        return data
    if isinstance(data, dict):
        return data.get("items") or data.get("entries") or []
    return []


def repo_paths():
    paths = []
    for pattern in REPOS.split():
        for p in sorted(glob.glob(os.path.expanduser(pattern))):
            if os.path.exists(os.path.join(p, ".git")):
                paths.append(p)
    return paths


def collect():
    rigs = []
    for r in as_list(rig_json("ps", "--json")):
        if r.get("isArchived"):
            continue
        seats = []
        for n in as_list(rig_json("ps", "--nodes", "--rig", r["name"], "--json")):
            name = n.get("canonicalSessionName", "")
            screen = rig("capture", name, "--lines", "25") if name else ""
            seats.append({
                "name": name,
                "activity": (n.get("agentActivity") or {}).get("state", "unknown"),
                "status": n.get("sessionStatus"),
                "work": n.get("assignedWorkCount", 0),
                "lastActivity": n.get("lastActivity"),
                "waiting": bool(PROMPT.search(screen)),
            })
        rigs.append({"name": r["name"], "status": r.get("status"), "uptime": r.get("uptime"),
                     "attention": r.get("attentionCount", 0), "seats": seats})

    queue = []
    for q in as_list(rig_json("queue", "list", "-A", "-a", "--limit", "60", "--json")):
        queue.append({"id": q.get("qitemId"), "state": q.get("state"), "dest": q.get("destinationSession"),
                      "source": q.get("sourceSession"), "summary": q.get("summary"), "updated": q.get("tsUpdated")})

    repos = []
    for p in repo_paths():
        commits = []
        for line in git(p, "log", "--branches", "-15", "--format=%H%x09%h%x09%cI%x09%D%x09%s").splitlines():
            full, short, when, refs, subject = (line.split("\t") + [""] * 5)[:5]
            commits.append({"hash": full, "short": short, "when": when, "refs": refs, "subject": subject})
        repos.append({"name": os.path.basename(p), "commits": commits})

    return {"updated": time.time(), "rigs": rigs, "queue": queue, "repos": repos}


def notify(title, body, priority=3, tags=""):
    if not (NTFY_URL and NTFY_TOPIC):
        return
    headers = {"Title": title.encode("utf-8"), "Priority": str(priority)}
    if tags:
        headers["Tags"] = tags
    if PUBLIC_URL:
        headers["Click"] = PUBLIC_URL
    if NTFY_TOKEN:
        headers["Authorization"] = f"Bearer {NTFY_TOKEN}"
    req = urllib.request.Request(f"{NTFY_URL.rstrip('/')}/{NTFY_TOPIC}", data=body.encode("utf-8"), headers=headers)
    try:
        urllib.request.urlopen(req, timeout=10).close()
    except OSError:
        pass


class Watcher:
    """Turns successive snapshots into notifications. The first snapshot only seeds."""

    def __init__(self):
        self.seen_commits = None
        self.waiting_polls = {}
        self.notified_waiting = set()
        self.queue_state = None

    def events(self, snap):
        out = []
        commits = {c["hash"]: (r["name"], c) for r in snap["repos"] for c in r["commits"]}
        queue = {q["id"]: q for q in snap["queue"]}
        if self.seen_commits is None:
            self.seen_commits = set(commits)
            self.queue_state = {k: q["state"] for k, q in queue.items()}
        else:
            new = [commits[h] for h in commits if h not in self.seen_commits]
            self.seen_commits |= set(commits)
            by_repo = {}
            for repo, c in new:
                by_repo.setdefault(repo, []).append(c["subject"])
            for repo, subjects in by_repo.items():
                out.append((f"{repo}: {len(subjects)} ny commit" + ("s" if len(subjects) > 1 else ""),
                            "\n".join(subjects), 2, "package"))
            for k, q in queue.items():
                before = self.queue_state.get(k)
                if before is None and HUMAN_DEST.search(q["dest"] or "") and q["state"] not in CLOSED:
                    out.append(("Teamet spørger dig", q["summary"] or q["source"] or k, 4, "raising_hand"))
                if before and before not in CLOSED and q["state"] in CLOSED and DONE_DEST.search(q["dest"] or ""):
                    out.append((f"Opgave {q['state']}: {q['dest']}", q["summary"] or k, 4, "white_check_mark"))
            self.queue_state = {k: q["state"] for k, q in queue.items()}

        # a prompt must survive two polls before it pings, so prompts answered quickly stay quiet
        for r in snap["rigs"]:
            for s in r["seats"]:
                n = s["name"]
                self.waiting_polls[n] = self.waiting_polls.get(n, 0) + 1 if s["waiting"] else 0
                if self.waiting_polls[n] >= 2 and n not in self.notified_waiting:
                    self.notified_waiting.add(n)
                    out.append((f"{n} venter på svar", "Åbn for at se skærmen og svare.", 4, "warning"))
                if not s["waiting"]:
                    self.notified_waiting.discard(n)
        return out


STATE = {"snap": None}
LOCK = threading.Lock()


def poll_loop():
    watcher = Watcher()
    while True:
        try:
            snap = collect()
            with LOCK:
                STATE["snap"] = snap
            for title, body, prio, tags in watcher.events(snap):
                notify(title, body, prio, tags)
        except Exception as e:  # keep polling whatever one round did
            print("poll failed:", e, flush=True)
        time.sleep(POLL)


def known_seats():
    with LOCK:
        snap = STATE["snap"] or {"rigs": []}
    return {s["name"] for r in snap["rigs"] for s in r["seats"]}


TYPES = {".html": "text/html; charset=utf-8", ".js": "text/javascript; charset=utf-8",
         ".webmanifest": "application/manifest+json", ".svg": "image/svg+xml", ".png": "image/png",
         ".json": "application/json"}


def csp():
    ext = " ".join(sorted({u.split("/", 3)[0] + "//" + u.split("/", 3)[2] for u in (THEME_CSS, THEME_JS) if u.startswith("http")}))
    return (f"default-src 'self'; style-src 'self' {ext}; script-src 'self' {ext}; img-src 'self' data:; "
            "connect-src 'self'; manifest-src 'self'; worker-src 'self'; frame-ancestors 'none'; base-uri 'none'")


class Handler(BaseHTTPRequestHandler):
    server_version = "rig-pocket"

    def log_message(self, *args):
        pass

    def send(self, code, body, ctype="application/json", cache="no-store"):
        data = body if isinstance(body, bytes) else body.encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(data)))
        self.send_header("Cache-Control", cache)
        self.send_header("Content-Security-Policy", csp())
        self.send_header("X-Content-Type-Options", "nosniff")
        self.send_header("Referrer-Policy", "no-referrer")
        self.end_headers()
        self.wfile.write(data)

    def do_GET(self):
        path = self.path.split("?", 1)[0]
        if path == "/api/status":
            with LOCK:
                snap = STATE["snap"]
            return self.send(200, json.dumps(snap or {"updated": 0, "rigs": [], "queue": [], "repos": []}))
        if path.startswith("/api/capture/"):
            seat = urllib.parse.unquote(path.rsplit("/", 1)[1])
            if not SEAT.match(seat) or seat not in known_seats():
                return self.send(404, '{"error":"unknown seat"}')
            return self.send(200, json.dumps({"seat": seat, "screen": rig("capture", seat, "--lines", "60")}))
        name = "index.html" if path in ("/", "") else path.lstrip("/")
        f = (STATIC / name).resolve()
        if STATIC not in f.parents or not f.is_file():
            return self.send(404, "not found", "text/plain")
        data = f.read_bytes()
        if f.name == "index.html":
            data = data.replace(b"{{THEME_CSS}}", THEME_CSS.encode()).replace(b"{{THEME_JS}}", THEME_JS.encode())
        cache = "no-cache" if f.suffix in (".html", ".webmanifest") or f.name == "sw.js" else "max-age=300"
        self.send(200, data, TYPES.get(f.suffix, "application/octet-stream"), cache)

    def do_POST(self):
        if self.path != "/api/send":
            return self.send(404, '{"error":"not found"}')
        # a custom header forces a CORS preflight, so other sites cannot post here
        if self.headers.get("X-Rig-Pocket") != "1" or (ORIGIN and self.headers.get("Origin") != ORIGIN):
            return self.send(403, '{"error":"forbidden"}')
        try:
            body = json.loads(self.rfile.read(min(int(self.headers.get("Content-Length", 0)), 8192)))
        except (ValueError, json.JSONDecodeError):
            return self.send(400, '{"error":"bad json"}')
        seat, text, key = body.get("seat", ""), (body.get("text") or "").strip(), body.get("key")
        if not SEAT.match(seat) or seat not in known_seats():
            return self.send(404, '{"error":"unknown seat"}')
        if key in KEYS:
            out = rig("send", seat, key, "--dangerously-interact", "--reason", "rig-pocket: answered from phone")
        elif text and len(text) <= 2000:
            out = rig("send", seat, f"From the person (via rig-pocket on their phone): {text}")
        else:
            return self.send(400, '{"error":"nothing to send"}')
        self.send(200, json.dumps({"ok": True, "result": out.strip()[:500]}))


def main():
    threading.Thread(target=poll_loop, daemon=True).start()
    ThreadingHTTPServer((BIND, PORT), Handler).serve_forever()


if __name__ == "__main__":
    main()