@mengruo/dsh-vision-toolkit 0.0.1

package/vendor/agent-vision-toolkit/skills/vision-tools/scripts/html_shot.py

@@ -0,0 +1,420 @@
+#!/usr/bin/env python3
+"""Screenshot an HTML file with headless Chrome/Chromium/Edge.
+
+The default path intentionally remains the small Chrome CLI wrapper used by
+the original upstream tool. ``--full-page`` uses Chrome DevTools Protocol
+instead: the requested viewport remains the layout viewport while the page's
+CSS scroll height is measured and captured in one pass.
+"""
+
+import argparse
+import base64
+import hashlib
+import json
+import os
+import shutil
+import socket
+import struct
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+from urllib.error import URLError
+from urllib.parse import urlparse
+from urllib.request import urlopen
+
+
+CHROME_CANDIDATES = (
+    # macOS
+    "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome",
+    "/Applications/Chromium.app/Contents/MacOS/Chromium",
+    "/Applications/Microsoft Edge.app/Contents/MacOS/Microsoft Edge",
+    # Linux
+    "google-chrome", "google-chrome-stable", "chromium", "chromium-browser",
+    "microsoft-edge", "brave-browser",
+)
+
+
+def find_chrome() -> str | None:
+    for candidate in CHROME_CANDIDATES:
+        path = shutil.which(candidate)
+        if path:
+            return path
+        if candidate.startswith("/") and Path(candidate).is_file():
+            return candidate
+    for env in ("PROGRAMFILES", "PROGRAMFILES(X86)", "LOCALAPPDATA"):
+        base = os.environ.get(env)
+        if not base:
+            continue
+        for sub in ("Google/Chrome/Application/chrome.exe",
+                    "Microsoft/Edge/Application/msedge.exe"):
+            path = Path(base) / sub
+            if path.is_file():
+                return str(path)
+    return None
+
+
+def default_output(source: str) -> str:
+    if source.startswith(("http://", "https://", "file://", "data:")):
+        stem = Path(urlparse(source).path).stem
+        return f"{stem or 'page'}.png"
+    return f"{Path(source).stem}.png"
+
+
+def reserve_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
+        listener.bind(("127.0.0.1", 0))
+        return int(listener.getsockname()[1])
+
+
+def read_http_json(port: int, path: str) -> object:
+    with urlopen(f"http://127.0.0.1:{port}{path}", timeout=1.0) as response:
+        return json.loads(response.read().decode("utf-8"))
+
+
+class DevToolsSocket:
+    """Minimal RFC 6455 client for the Chrome DevTools Protocol."""
+
+    def __init__(self, websocket_url: str):
+        parsed = urlparse(websocket_url)
+        if parsed.scheme != "ws" or parsed.hostname is None or parsed.port is None:
+            raise RuntimeError("invalid DevTools websocket URL")
+        self.socket = socket.create_connection((parsed.hostname, parsed.port), timeout=5.0)
+        self.socket.settimeout(5.0)
+        self.buffer = bytearray()
+        key = base64.b64encode(os.urandom(16)).decode("ascii")
+        request = (
+            f"GET {parsed.path or '/'} HTTP/1.1\r\n"
+            f"Host: {parsed.hostname}:{parsed.port}\r\n"
+            "Upgrade: websocket\r\n"
+            "Connection: Upgrade\r\n"
+            f"Sec-WebSocket-Key: {key}\r\n"
+            "Sec-WebSocket-Version: 13\r\n\r\n"
+        ).encode("ascii")
+        self.socket.sendall(request)
+        response = self._read_until(b"\r\n\r\n")
+        if not response.startswith(b"HTTP/1.1 101"):
+            raise RuntimeError("Chrome did not accept the DevTools websocket")
+        expected_accept = base64.b64encode(hashlib.sha1(
+            f"{key}258EAFA5-E914-47DA-95CA-C5AB0DC85B11".encode("ascii")
+        ).digest()).decode("ascii")
+        headers = response.decode("latin-1").split("\r\n")
+        accept = next((line.split(":", 1)[1].strip() for line in headers
+                       if line.lower().startswith("sec-websocket-accept:")), None)
+        if accept != expected_accept:
+            raise RuntimeError("Chrome returned an invalid DevTools websocket handshake")
+        self.next_id = 0
+
+    def _read_until(self, marker: bytes) -> bytes:
+        data = bytearray()
+        while marker not in data:
+            chunk = self.socket.recv(4096)
+            if not chunk:
+                raise RuntimeError("DevTools websocket closed during handshake")
+            data.extend(chunk)
+        boundary = data.index(marker) + len(marker)
+        self.buffer.extend(data[boundary:])
+        return bytes(data[:boundary])
+
+    def _send_frame(self, opcode: int, payload: bytes) -> None:
+        mask = os.urandom(4)
+        masked = bytes(value ^ mask[index % 4] for index, value in enumerate(payload))
+        length = len(masked)
+        if length < 126:
+            header = struct.pack("!BB", 0x80 | opcode, 0x80 | length)
+        elif length < 65536:
+            header = struct.pack("!BBH", 0x80 | opcode, 0x80 | 126, length)
+        else:
+            header = struct.pack("!BBQ", 0x80 | opcode, 0x80 | 127, length)
+        self.socket.sendall(header + mask + masked)
+
+    def _read_frame(self) -> tuple[bool, int, bytes]:
+        header = self._read_exact(2)
+        first, second = header
+        final = (first & 0x80) != 0
+        opcode = first & 0x0F
+        length = second & 0x7F
+        masked = (second & 0x80) != 0
+        if length == 126:
+            length = struct.unpack("!H", self._read_exact(2))[0]
+        elif length == 127:
+            length = struct.unpack("!Q", self._read_exact(8))[0]
+        mask = self._read_exact(4) if masked else None
+        payload = self._read_exact(length)
+        if mask is not None:
+            payload = bytes(value ^ mask[index % 4] for index, value in enumerate(payload))
+        return final, opcode, payload
+
+    def _read_message(self) -> tuple[int, bytes]:
+        fragments: list[bytes] = []
+        message_opcode: int | None = None
+        while True:
+            final, opcode, payload = self._read_frame()
+            if opcode == 0x9:
+                self._send_frame(0xA, payload)
+                continue
+            if opcode == 0x8:
+                raise RuntimeError("DevTools websocket closed")
+            if opcode in (0x1, 0x2):
+                fragments = [payload]
+                message_opcode = opcode
+            elif opcode == 0x0 and message_opcode is not None:
+                fragments.append(payload)
+            else:
+                continue
+            if final and message_opcode is not None:
+                return message_opcode, b"".join(fragments)
+
+    def _read_exact(self, length: int) -> bytes:
+        data = bytearray()
+        if self.buffer:
+            take = min(length, len(self.buffer))
+            data.extend(self.buffer[:take])
+            del self.buffer[:take]
+        while len(data) < length:
+            chunk = self.socket.recv(length - len(data))
+            if not chunk:
+                raise RuntimeError("DevTools websocket closed unexpectedly")
+            data.extend(chunk)
+        return bytes(data)
+
+    def command(self, method: str, params: dict | None = None) -> dict:
+        self.next_id += 1
+        request_id = self.next_id
+        request = {"id": request_id, "method": method}
+        if params is not None:
+            request["params"] = params
+        self._send_frame(0x1, json.dumps(request, separators=(",", ":")).encode("utf-8"))
+        while True:
+            opcode, payload = self._read_message()
+            if opcode != 0x1:
+                continue
+            message = json.loads(payload.decode("utf-8"))
+            if message.get("id") != request_id:
+                continue
+            if "error" in message:
+                raise RuntimeError(f"{method}: {message['error'].get('message', 'CDP command failed')}")
+            return message.get("result", {})
+
+    def close(self) -> None:
+        try:
+            self._send_frame(0x8, b"")
+        except OSError:
+            pass
+        self.socket.close()
+
+
+def evaluate(client: DevToolsSocket, expression: str) -> object:
+    result = client.command("Runtime.evaluate", {
+        "expression": expression,
+        "returnByValue": True,
+        "awaitPromise": True,
+    })
+    if "exceptionDetails" in result:
+        raise RuntimeError(result["exceptionDetails"].get("text", "Runtime.evaluate failed"))
+    remote = result.get("result", {})
+    if remote.get("subtype") == "error":
+        raise RuntimeError(remote.get("description", "Runtime.evaluate failed"))
+    return remote.get("value")
+
+
+def wait_for_target(port: int, deadline: float) -> dict:
+    last_error: Exception | None = None
+    while time.monotonic() < deadline:
+        try:
+            targets = read_http_json(port, "/json/list")
+            if isinstance(targets, list):
+                for target in targets:
+                    if isinstance(target, dict) and target.get("type") == "page" and target.get("webSocketDebuggerUrl"):
+                        return target
+        except (OSError, URLError, ValueError) as error:
+            last_error = error
+        time.sleep(0.05)
+    detail = f": {last_error}" if last_error else ""
+    raise RuntimeError(f"timed out waiting for Chrome DevTools{detail}")
+
+
+def wait_for_ready(client: DevToolsSocket, deadline: float) -> None:
+    while time.monotonic() < deadline:
+        if evaluate(client, "document.readyState") == "complete":
+            return
+        time.sleep(0.05)
+    raise RuntimeError("timed out waiting for the HTML document")
+
+
+def pause_for_frame(seconds: float) -> None:
+    time.sleep(max(0.02, seconds))
+
+
+def capture_full_page(chrome: str, source: str, output: Path, width: int, height: int,
+                      scale: int, wait_ms: int, max_pixels: int | None) -> int:
+    port = reserve_port()
+    guarded_profile = os.environ.get("DSH_VISION_CHROME_PROFILE")
+    profile = guarded_profile or tempfile.mkdtemp(prefix="dsh-vision-chrome-")
+    process: subprocess.Popen | None = None
+    client: DevToolsSocket | None = None
+    try:
+        command = [
+            chrome, "--headless=new", "--disable-gpu", "--hide-scrollbars",
+            "--no-first-run", "--no-default-browser-check",
+            "--blink-settings=imagesLazyLoadingEnabled=false",
+            f"--window-size={width},{height}",
+            f"--remote-debugging-port={port}",
+            "--remote-debugging-address=127.0.0.1",
+        ]
+        if guarded_profile is None:
+            command.append(f"--user-data-dir={profile}")
+        if scale != 1:
+            command.append(f"--force-device-scale-factor={scale}")
+        command.append(source)
+        process = subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+        target = wait_for_target(port, time.monotonic() + 15.0)
+        client = DevToolsSocket(str(target["webSocketDebuggerUrl"]))
+        client.command("Page.enable")
+        client.command("Runtime.enable")
+        client.command("Emulation.setDeviceMetricsOverride", {
+            "width": width,
+            "height": height,
+            "deviceScaleFactor": scale,
+            "mobile": False,
+            "screenWidth": width,
+            "screenHeight": height,
+        })
+        client.command("Page.navigate", {"url": source})
+        wait_for_ready(client, time.monotonic() + 15.0)
+        if wait_ms > 0:
+            pause_for_frame(wait_ms / 1000.0)
+
+        # Keep the requested viewport fixed. Scrolling, rather than growing
+        # the viewport, wakes IO/reveal handlers without changing vh/svh.
+        evaluate(client, """
+            (() => {
+              document.documentElement.style.scrollBehavior = 'auto';
+              if (document.body) document.body.style.scrollBehavior = 'auto';
+              window.scrollTo(0, 0);
+              return true;
+            })()
+        """)
+        page_height = int(evaluate(client, """
+            Math.max(document.documentElement.scrollHeight,
+                     document.body ? document.body.scrollHeight : 0)
+        """) or 0)
+        for _ in range(2):
+            measured_before_sweep = page_height
+            for position in range(0, page_height, max(1, height)):
+                evaluate(client, f"window.scrollTo(0, {position}); true")
+                pause_for_frame(0.12)
+            evaluate(client, f"window.scrollTo(0, {page_height}); true")
+            pause_for_frame(0.12)
+            page_height = int(evaluate(client, """
+                Math.max(document.documentElement.scrollHeight,
+                         document.body ? document.body.scrollHeight : 0)
+            """) or 0)
+            if page_height <= measured_before_sweep:
+                break
+        evaluate(client, """
+            Promise.race([
+              Promise.all(Array.from(document.images)
+                .filter(image => !image.complete)
+                .map(image => new Promise(resolve => {
+                  image.addEventListener('load', resolve, {once: true});
+                  image.addEventListener('error', resolve, {once: true});
+                }))),
+              new Promise(resolve => setTimeout(resolve, 3000))
+            ]).then(() => true)
+        """)
+        evaluate(client, "window.scrollTo(0, 0); true")
+        pause_for_frame(0.5)
+        page_height = int(evaluate(client, """
+            Math.max(document.documentElement.scrollHeight,
+                     document.body ? document.body.scrollHeight : 0)
+        """) or 0)
+        if page_height <= 0:
+            raise RuntimeError("document has no measurable scroll height")
+        if max_pixels is not None and width * page_height * scale * scale > max_pixels:
+            raise RuntimeError(f"full-page screenshot would exceed maxImagePixels ({width * page_height * scale * scale} > {max_pixels})")
+        result = client.command("Page.captureScreenshot", {
+            "format": "png",
+            "fromSurface": True,
+            "captureBeyondViewport": True,
+            "clip": {"x": 0, "y": 0, "width": width, "height": page_height, "scale": 1},
+        })
+        data = result.get("data")
+        if not isinstance(data, str):
+            raise RuntimeError("Chrome did not return screenshot data")
+        output.write_bytes(base64.b64decode(data))
+        return page_height
+    finally:
+        if client is not None:
+            client.close()
+        if process is not None and process.poll() is None:
+            process.terminate()
+            try:
+                process.wait(timeout=2)
+            except subprocess.TimeoutExpired:
+                process.kill()
+                process.wait(timeout=2)
+        if guarded_profile is None:
+            shutil.rmtree(profile, ignore_errors=True)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        prog="html_shot",
+        description="Screenshot an HTML file (or URL) to PNG with headless Chrome/Chromium/Edge",
+    )
+    parser.add_argument("source", help="path to a local .html file, or an http(s):// URL")
+    parser.add_argument("-o", "--output", help="output PNG path (default: <source-stem>.png in the current directory)")
+    parser.add_argument("--width", type=int, default=1280, help="viewport width in CSS pixels (default: 1280)")
+    parser.add_argument("--height", type=int, default=800, help="viewport height in CSS pixels (default: 800)")
+    parser.add_argument("--scale", type=int, default=1,
+                        help="device scale factor: 2 makes the PNG 2x for small text (default: 1)")
+    parser.add_argument("--wait-ms", type=int, default=0,
+                        help="virtual time budget in ms before capturing (default: 0, capture immediately)")
+    parser.add_argument("--full-page", action="store_true",
+                        help="capture the full CSS scroll height while keeping the requested viewport")
+    parser.add_argument("--max-pixels", type=int,
+                        help="reject full-page captures larger than this many output pixels")
+    args = parser.parse_args()
+
+    chrome = find_chrome()
+    if not chrome:
+        parser.exit(1, "html_shot: no Chrome/Chromium/Edge found; install one and retry\n")
+
+    source = args.source
+    if not source.startswith(("http://", "https://", "file://", "data:")):
+        path = Path(source).expanduser()
+        if not path.is_file():
+            parser.exit(1, f"html_shot: file not found: {path}\n")
+        source = path.resolve().as_uri()
+
+    output = Path(args.output).expanduser().resolve() if args.output else Path(default_output(source)).resolve()
+    if args.full_page:
+        try:
+            page_height = capture_full_page(chrome, source, output, args.width, args.height, args.scale, args.wait_ms, args.max_pixels)
+        except Exception as error:
+            parser.exit(1, f"html_shot: full-page capture failed: {error}\n")
+        print(f"wrote {output} ({args.width * args.scale}x{page_height * args.scale}; pageHeight={page_height})")
+        return
+
+    command = [
+        chrome, "--headless=new", "--disable-gpu", "--hide-scrollbars",
+        "--no-first-run", "--no-default-browser-check",
+        f"--window-size={args.width},{args.height}",
+        f"--screenshot={output}",
+    ]
+    if args.scale != 1:
+        command.append(f"--force-device-scale-factor={args.scale}")
+    if args.wait_ms > 0:
+        command.append(f"--virtual-time-budget={args.wait_ms}")
+    command.append(source)
+
+    result = subprocess.run(command, text=True, capture_output=True)
+    if result.returncode != 0 or not output.is_file():
+        message = result.stderr.strip() or result.stdout.strip() or f"chrome exited with code {result.returncode}"
+        parser.exit(1, f"html_shot: capture failed: {message}\n")
+    print(f"wrote {output} ({args.width * args.scale}x{args.height * args.scale})")
+
+
+if __name__ == "__main__":
+    main()