39227-vm/tools/local_screenshot_bridge.py

#!/usr/bin/env python3
import argparse
import base64
import json
import os
import re
import subprocess
import sys
from datetime import datetime, timezone
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path


def _slug(s: str, max_len: int = 80) -> str:
    s = (s or "").strip().lower()
    s = re.sub(r"[^a-z0-9]+", "-", s)
    s = s.strip("-")
    if not s:
        return "screenshot"
    return s[:max_len]


class Handler(BaseHTTPRequestHandler):
    server_version = "LocalScreenshotBridge/0.1"

    def _send_json(self, status: int, payload: dict):
        body = json.dumps(payload, ensure_ascii=True).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        # Chrome extension fetch() to localhost will preflight; allow it.
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):  # noqa: N802
        if self.path not in ("/", "/health"):
            self._send_json(404, {"ok": False, "error": "not_found"})
            return
        self._send_json(
            200,
            {
                "ok": True,
                "service": "local_screenshot_bridge",
                "out_dir": str(self.server.out_dir),  # type: ignore[attr-defined]
                "has_run_cmd": bool(getattr(self.server, "run_cmd", None)),  # type: ignore[attr-defined]
            },
        )

    def do_OPTIONS(self):  # noqa: N802
        self.send_response(204)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.end_headers()

    def do_POST(self):  # noqa: N802
        if self.path != "/screenshot":
            self._send_json(404, {"ok": False, "error": "not_found"})
            return

        try:
            length = int(self.headers.get("Content-Length", "0"))
        except ValueError:
            self._send_json(400, {"ok": False, "error": "bad_content_length"})
            return

        raw = self.rfile.read(length)
        try:
            req = json.loads(raw.decode("utf-8"))
        except Exception:
            self._send_json(400, {"ok": False, "error": "bad_json"})
            return

        data_url = req.get("data_url") or ""
        title = req.get("title") or ""
        page_url = req.get("url") or ""
        client_ts = req.get("ts") or ""

        m = re.match(r"^data:image/png;base64,(.*)$", data_url)
        if not m:
            self._send_json(400, {"ok": False, "error": "expected_png_data_url"})
            return

        try:
            png_bytes = base64.b64decode(m.group(1), validate=True)
        except Exception:
            self._send_json(400, {"ok": False, "error": "bad_base64"})
            return

        now = datetime.now(timezone.utc)
        stamp = now.strftime("%Y%m%dT%H%M%SZ")
        base = f"{stamp}-{_slug(title)}"

        out_dir: Path = self.server.out_dir  # type: ignore[attr-defined]
        out_dir.mkdir(parents=True, exist_ok=True)
        png_path = out_dir / f"{base}.png"
        meta_path = out_dir / f"{base}.json"

        try:
            png_path.write_bytes(png_bytes)
            meta_path.write_text(
                json.dumps(
                    {
                        "title": title,
                        "url": page_url,
                        "client_ts": client_ts,
                        "saved_utc": now.isoformat(),
                        "png_path": str(png_path),
                    },
                    indent=2,
                    ensure_ascii=True,
                )
                + "\n",
                encoding="utf-8",
            )
        except Exception as e:
            self._send_json(500, {"ok": False, "error": "write_failed", "detail": str(e)})
            return

        run = getattr(self.server, "run_cmd", None)  # type: ignore[attr-defined]
        ran = None
        if run:
            try:
                proc = subprocess.run(
                    run + [str(png_path), str(meta_path)],
                    cwd=str(self.server.project_root),  # type: ignore[attr-defined]
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    text=True,
                )
                ran = {
                    "cmd": run,
                    "exit_code": proc.returncode,
                    "stdout": proc.stdout[-4000:],
                    "stderr": proc.stderr[-4000:],
                }
            except Exception as e:
                ran = {"cmd": run, "error": str(e)}

        self._send_json(
            200,
            {
                "ok": True,
                "png_path": str(png_path),
                "meta_path": str(meta_path),
                "ran": ran,
            },
        )


def main(argv: list[str]) -> int:
    p = argparse.ArgumentParser(description="Receive screenshots from a Chrome extension and save into this project.")
    p.add_argument("--port", type=int, default=8765)
    p.add_argument("--bind", default="127.0.0.1", help="Bind address (default: 127.0.0.1)")
    p.add_argument("--out-dir", default="screenshots", help="Output directory relative to project root")
    p.add_argument(
        "--run",
        nargs="+",
        default=None,
        help="Optional command to run after saving. Screenshot paths are appended as args: PNG then JSON.",
    )
    args = p.parse_args(argv)

    project_root = Path(__file__).resolve().parents[1]
    out_dir = (project_root / args.out_dir).resolve()

    httpd = HTTPServer((args.bind, args.port), Handler)
    httpd.project_root = project_root  # type: ignore[attr-defined]
    httpd.out_dir = out_dir  # type: ignore[attr-defined]
    httpd.run_cmd = args.run  # type: ignore[attr-defined]

    print(f"Listening on http://{args.bind}:{args.port}/screenshot", file=sys.stderr)
    print(f"Saving screenshots to {out_dir}", file=sys.stderr)
    if args.run:
        print(f"Will run: {' '.join(args.run)} <png_path> <meta_path>", file=sys.stderr)
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))