ksjsb/core/h5_kww.py

"""Nebula H5 `kww` header generator backed by the KWF WebView VM.

KWF runtime behavior confirmed from APP WebView:

- `kwf-0.0.2` installs `window.kwpsec.getData`;
- `getData()` returns the 174-char `PnGU...` value used as H5 request header
  `kww`;
- each call updates `localStorage.kwfcv1` and `localStorage.kwfv1`.

The active `PnGU...` branch is implemented in pure Python.  The Node VM bridge
is kept as a cross-check/fallback for future KWF branch changes.
"""

from __future__ import annotations

import json
import os
import queue
import subprocess
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from core.h5_kww_alg import kwf_generate_kww

DEFAULT_H5_KWW_SERVER = Path("core/h5_kww_server.mjs")


@dataclass(frozen=True)
class H5KwwResult:
    kww: str
    kwfcv1: str = ""
    kwfv1: str = ""


class H5KwwGenerator:
    """长驻 Node/KWF 进程，保持 localStorage 计数状态。"""

    def __init__(
        self,
        server: str | Path = DEFAULT_H5_KWW_SERVER,
        node_bin: str | None = None,
        timeout: int = 30,
    ) -> None:
        self.server = Path(server)
        self.node = node_bin or os.environ.get("NODE_BIN") or "node"
        self.timeout = timeout
        self._proc: subprocess.Popen[str] | None = None
        self._stdout_queue: queue.Queue[str] = queue.Queue()
        self._lock = threading.Lock()
        self._request_id = 0

    def _start(self) -> None:
        if self._proc is not None and self._proc.poll() is None:
            return
        if not self.server.exists():
            raise FileNotFoundError(f"H5 kww server not found: {self.server}")
        self._stdout_queue = queue.Queue()
        self._proc = subprocess.Popen(
            [self.node, str(self.server)],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            text=True,
            encoding="utf-8",
            bufsize=1,
        )

        def read_stdout() -> None:
            assert self._proc is not None
            assert self._proc.stdout is not None
            for line in self._proc.stdout:
                self._stdout_queue.put(line)

        thread = threading.Thread(target=read_stdout, name="H5KwwGeneratorStdout", daemon=True)
        thread.start()

    def close(self) -> None:
        proc = self._proc
        self._proc = None
        if proc is None:
            return
        try:
            if proc.stdin:
                proc.stdin.close()
        except Exception:
            pass
        try:
            if proc.stdout:
                proc.stdout.close()
        except Exception:
            pass
        try:
            proc.wait(timeout=1)
        except Exception:
            try:
                proc.kill()
            except Exception:
                pass

    def get(
        self,
        *,
        url: str = "",
        method: str = "GET",
        headers: dict[str, str] | None = None,
        body: str = "",
        cookie: str = "",
    ) -> H5KwwResult:
        with self._lock:
            self._start()
            assert self._proc is not None
            assert self._proc.stdin is not None
            self._request_id += 1
            request_id = self._request_id
            payload = json.dumps(
                {
                    "id": request_id,
                    "url": url,
                    "method": method,
                    "headers": headers or {},
                    "body": body,
                    "cookie": cookie,
                },
                ensure_ascii=False,
                separators=(",", ":"),
            )
            try:
                self._proc.stdin.write(payload + "\n")
                self._proc.stdin.flush()
            except Exception:
                self.close()
                raise

            while True:
                try:
                    line = self._stdout_queue.get(timeout=self.timeout)
                except queue.Empty as exc:
                    self.close()
                    raise TimeoutError("H5 kww server timed out") from exc
                try:
                    data = json.loads(line)
                except json.JSONDecodeError:
                    continue
                if data.get("id") not in {request_id, None}:
                    continue
                if not data.get("ok", True):
                    raise RuntimeError(f"H5 kww server failed: {data.get('error')}")
                kww = str(data.get("kww") or "")
                if not _looks_like_kww(kww):
                    raise RuntimeError(f"H5 kww server returned invalid kww: {kww!r}")
                return H5KwwResult(
                    kww=kww,
                    kwfcv1=str(data.get("kwfcv1") or ""),
                    kwfv1=str(data.get("kwfv1") or ""),
                )


class PureH5KwwGenerator:
    """Pure-Python KWF `PnGU...` generator for the current Nebula branch."""

    def __init__(
        self,
        *,
        start_collect_count: int = 1,
        now_ms: int | None = None,
        language: str = "zh-CN",
    ) -> None:
        self.collect_count = start_collect_count
        self.now_ms = now_ms
        self.language = language
        self._lock = threading.Lock()

    def close(self) -> None:
        return

    def get(
        self,
        *,
        url: str = "",
        method: str = "GET",
        headers: dict[str, str] | None = None,
        body: str = "",
        cookie: str = "",
    ) -> H5KwwResult:
        del url, method, headers, body, cookie
        with self._lock:
            current = self.collect_count
            now_ms = self.now_ms if self.now_ms is not None else int(time.time() * 1000)
            kww = kwf_generate_kww(
                collect_count=current,
                now_ms=now_ms,
                language=self.language,
            )
            self.collect_count += 1
            return H5KwwResult(kww=kww, kwfcv1=str(self.collect_count), kwfv1=kww)


def _looks_like_kww(value: str) -> bool:
    if len(value) < 120:
        return False
    return all(ch.isalnum() or ch in "+/=" for ch in value)


__all__ = [
    "DEFAULT_H5_KWW_SERVER",
    "H5KwwGenerator",
    "H5KwwResult",
    "PureH5KwwGenerator",
]