#!/usr/bin/env python3
"""极简 CDP 截图器（仅标准库）。
用法: cdp_shot.py <ws_port> <out.png> <width> <height> <mobile:0|1> <path> <token>
流程: 导航到 8090 -> 注入 oa.token 到 localStorage -> 重载进入登录态 -> 设备视口 -> 截图。
"""
import sys, json, socket, base64, hashlib, struct, os, urllib.request, time

def http_json(port, path):
    with urllib.request.urlopen(f"http://127.0.0.1:{port}{path}", timeout=10) as r:
        return json.loads(r.read().decode())

class WS:
    def __init__(self, url):
        # ws://127.0.0.1:port/devtools/page/XXXX
        assert url.startswith("ws://")
        hostport, _, path = url[5:].partition("/")
        host, _, port = hostport.partition(":")
        self.path = "/" + path
        self.sock = socket.create_connection((host, int(port)), timeout=15)
        key = base64.b64encode(os.urandom(16)).decode()
        req = (f"GET {self.path} HTTP/1.1\r\nHost: {hostport}\r\n"
               "Upgrade: websocket\r\nConnection: Upgrade\r\n"
               f"Sec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n\r\n")
        self.sock.sendall(req.encode())
        # read handshake response headers
        buf = b""
        while b"\r\n\r\n" not in buf:
            buf += self.sock.recv(4096)
        self.rest = buf.split(b"\r\n\r\n", 1)[1]
        self._id = 0

    def send(self, method, params=None):
        self._id += 1
        msg = json.dumps({"id": self._id, "method": method, "params": params or {}}).encode()
        header = bytearray([0x81])  # FIN + text
        mask = os.urandom(4)
        n = len(msg)
        if n < 126:
            header.append(0x80 | n)
        elif n < 65536:
            header.append(0x80 | 126); header += struct.pack(">H", n)
        else:
            header.append(0x80 | 127); header += struct.pack(">Q", n)
        header += mask
        masked = bytes(b ^ mask[i % 4] for i, b in enumerate(msg))
        self.sock.sendall(bytes(header) + masked)
        return self._id

    def _read(self, n):
        while len(self.rest) < n:
            chunk = self.sock.recv(65536)
            if not chunk:
                raise IOError("socket closed")
            self.rest += chunk
        out, self.rest = self.rest[:n], self.rest[n:]
        return out

    def recv_msg(self):
        # assemble one full (possibly fragmented) text message
        data = b""
        while True:
            b0, b1 = self._read(2)
            fin = b0 & 0x80
            opcode = b0 & 0x0F
            length = b1 & 0x7F
            if length == 126:
                length = struct.unpack(">H", self._read(2))[0]
            elif length == 127:
                length = struct.unpack(">Q", self._read(8))[0]
            payload = self._read(length) if length else b""
            if opcode == 0x9:  # ping -> pong
                continue
            if opcode == 0x8:  # close
                raise IOError("ws closed by peer")
            data += payload
            if fin:
                return json.loads(data.decode("utf-8", "replace"))

    def call(self, method, params=None, timeout=30):
        mid = self.send(method, params)
        t0 = time.time()
        while time.time() - t0 < timeout:
            m = self.recv_msg()
            if m.get("id") == mid:
                return m
        raise TimeoutError(method)

def main():
    port, out, w, h, mobile, path, token = (sys.argv[1], sys.argv[2], int(sys.argv[3]),
                                            int(sys.argv[4]), sys.argv[5] == "1",
                                            sys.argv[6], sys.argv[7])
    targets = http_json(port, "/json")
    page = next(t for t in targets if t.get("type") == "page")
    ws = WS(page["webSocketDebuggerUrl"])
    ws.call("Page.enable")
    ws.call("Runtime.enable")
    ws.call("Emulation.setDeviceMetricsOverride",
            {"width": w, "height": h, "deviceScaleFactor": 2, "mobile": mobile})
    # 1) load origin so localStorage is writable for it
    ws.call("Page.navigate", {"url": "http://127.0.0.1:8090/login"})
    time.sleep(2.5)
    # 2) inject token + go to target path
    ws.call("Runtime.evaluate", {"expression":
        f"localStorage.setItem('oa.token','{token}'); 'ok'"})
    ws.call("Page.navigate", {"url": f"http://127.0.0.1:8090{path}"})
    time.sleep(4.5)  # let SPA boot + data load
    shot = ws.call("Page.captureScreenshot", {"format": "png", "captureBeyondViewport": False})
    img = base64.b64decode(shot["result"]["data"])
    with open(out, "wb") as f:
        f.write(img)
    print(f"OK {out} {len(img)} bytes")

if __name__ == "__main__":
    main()