cmbuyer/client/src/cmbuyer_client/device/baseline.py

"""只读设备基线取证，严格限制在元数据、截图和完整节点树。"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from datetime import UTC, datetime
from hashlib import sha256
import base64
import binascii
from io import BytesIO
import json
import os
from pathlib import Path
import shutil
from typing import Any, Protocol
from uuid import uuid4
from xml.etree import ElementTree

from adbutils.errors import AdbTimeout
from PIL import Image, UnidentifiedImageError
from uiautomator2.exceptions import HTTPTimeoutError

from .adb import AdbClient, DeviceInspection


PDD_PACKAGE = "com.xunmeng.pinduoduo"
SCREENSHOT_PARAMS = [1, 80]
HIERARCHY_PARAMS = [False, 50]
_BASE64_ASCII_WHITESPACE = " \t\r\n"


class BaselineCaptureError(RuntimeError):
    """基线取证无法完整落盘时的失败，不会伪造成功产物。"""


class BaselineCaptureTimeoutError(BaselineCaptureError):
    """设备基线取证超时；底层异常文本不向 CLI 或日志泄露。"""


class UiAutomatorDevice(Protocol):
    """本任务唯一需要的 uiautomator2 只读能力。"""

    def app_info(self, package_name: str) -> dict[str, Any]:
        """读取已安装应用元数据。"""

    def jsonrpc_call(self, method: str, params: Any = None, timeout: float = 10) -> Any:
        """调用公开 uiautomator2 JSON-RPC 接口。"""


class NoReconnectUiautomatorConnector:
    """只把当前 ADB server 已列出的设备对象交给 uiautomator2。

    uiautomator2 直接接收 IP serial 时会在内部尝试 adb disconnect/connect。这里先从已列出设备中
    取对象再调用 ``u2.connect(device_object)``，避免连接阶段隐式重连已经掉线的 WiFi 通道。
    """

    def __init__(self, list_devices: Callable[[], list[Any]], connect: Callable[[Any], UiAutomatorDevice]) -> None:
        self._list_devices = list_devices
        self._connect = connect

    def __call__(self, serial: str) -> UiAutomatorDevice:
        device = next((item for item in self._list_devices() if item.serial == serial), None)
        if device is None:
            raise BaselineCaptureError("设备在连接前已从 ADB 清单消失，已拒绝自动重连。")
        return self._connect(device)


@dataclass(frozen=True)
class BaselineCaptureResult:
    """已原子发布的基线取证摘要，不包含页面正文或原始 serial。"""

    output_directory: Path
    manifest_path: Path
    screenshot_path: Path
    hierarchy_path: Path


class DeviceBaselineCapturer:
    """以先校验通道、后连接、最后原子发布的顺序采集基线。

    截图和 XML 可能包含页面敏感内容，因此仅落在调用方明确指定的本地目录；manifest 只写
    哈希、设备非敏感元数据和脱敏后的 serial 摘要，绝不嵌入 XML 或页面文本。
    """

    def __init__(
        self,
        adb_client: AdbClient,
        connector: Callable[[str], UiAutomatorDevice],
        timeout_seconds: float,
    ) -> None:
        if timeout_seconds <= 0:
            raise ValueError("timeout_seconds 必须大于 0")
        self._adb_client = adb_client
        self._connector = connector
        self._timeout_seconds = timeout_seconds

    def capture(self, serial: str, output_directory: Path) -> BaselineCaptureResult:
        """采集完整基线，任何一步失败均不发布 output_directory。"""

        inspection = self._adb_client.inspect(serial)
        target = Path(output_directory)
        if target.exists():
            raise BaselineCaptureError("输出目录已存在；为防止混入旧证据，拒绝覆盖。")
        if not target.name:
            raise BaselineCaptureError("输出目录必须是明确的新目录。")

        target.parent.mkdir(parents=True, exist_ok=True)
        staging = target.parent / f".{target.name}.staging-{uuid4().hex}"
        staging.mkdir()
        try:
            device = self._connector(serial)
            app_info = device.app_info(PDD_PACKAGE)
            version = _extract_version(app_info)

            screenshot_path = staging / "screenshot.png"
            screenshot_base64 = device.jsonrpc_call(
                "takeScreenshot",
                SCREENSHOT_PARAMS,
                timeout=self._timeout_seconds,
            )
            _save_base64_screenshot(screenshot_base64, screenshot_path)

            hierarchy = device.jsonrpc_call(
                "dumpWindowHierarchy",
                HIERARCHY_PARAMS,
                timeout=self._timeout_seconds,
            )
            _validate_hierarchy(hierarchy)
            hierarchy_path = staging / "hierarchy.xml"
            hierarchy_path.write_text(hierarchy, encoding="utf-8")

            manifest_path = staging / "manifest.json"
            manifest_path.write_text(
                json.dumps(
                    _manifest(inspection, serial, version, screenshot_path, hierarchy_path),
                    ensure_ascii=False,
                    indent=2,
                    sort_keys=True,
                )
                + "\n",
                encoding="utf-8",
            )
            os.replace(staging, target)
        except BaselineCaptureError:
            # 仅删除本次创建、名称带随机标识的暂存目录，绝不触碰调用方原有输出目录。
            if staging.exists():
                shutil.rmtree(staging)
            raise
        except (AdbTimeout, HTTPTimeoutError, TimeoutError) as error:
            if staging.exists():
                shutil.rmtree(staging)
            raise BaselineCaptureTimeoutError("设备基线取证超时，未发布任何证据产物。") from error
        except Exception as error:
            if staging.exists():
                shutil.rmtree(staging)
            # uiautomator2/adbutils 可能把 serial、路径或远端响应放入异常文本，不能直接传播到 CLI。
            raise BaselineCaptureError("设备基线取证未完成，未发布任何证据产物。") from error

        return BaselineCaptureResult(
            output_directory=target,
            manifest_path=target / "manifest.json",
            screenshot_path=target / "screenshot.png",
            hierarchy_path=target / "hierarchy.xml",
        )


def _extract_version(app_info: dict[str, Any]) -> str:
    version = app_info.get("versionName") or app_info.get("version_name")
    if not isinstance(version, str) or not version.strip():
        raise BaselineCaptureError("无法读取拼多多版本，拒绝发布不完整取证。")
    return version.strip()


def _save_base64_screenshot(value: Any, target: Path) -> None:
    """规范化常见 ASCII Base64 空白后严格解码；没有 adb screenshot fallback。"""

    if not isinstance(value, str) or not value:
        raise BaselineCaptureError("截图 RPC 未返回 base64 数据，拒绝发布不完整取证。")
    try:
        normalized = value.translate({ord(character): None for character in _BASE64_ASCII_WHITESPACE})
        raw_image = base64.b64decode(normalized.encode("ascii"), validate=True)
    except (UnicodeEncodeError, ValueError, binascii.Error) as error:
        raise BaselineCaptureError("截图 RPC Base64 语法无效，拒绝发布不完整取证。") from error
    try:
        with Image.open(BytesIO(raw_image)) as image:
            image.load()
            image.save(target, format="PNG")
    except (UnidentifiedImageError, OSError) as error:
        raise BaselineCaptureError("截图 RPC 图像数据无效，拒绝发布不完整取证。") from error


def _validate_hierarchy(value: Any) -> None:
    """确认 RPC 返回的是完整节点树，不把原始 XML 放进错误或日志。"""

    if not isinstance(value, str) or not value:
        raise BaselineCaptureError("节点树导出为空，拒绝发布不完整取证。")
    try:
        root = ElementTree.fromstring(value)
    except ElementTree.ParseError as error:
        raise BaselineCaptureError("节点树不是有效 XML，拒绝发布不完整取证。") from error
    if root.tag != "hierarchy":
        raise BaselineCaptureError("节点树根节点无效，拒绝发布不完整取证。")


def _manifest(
    inspection: DeviceInspection,
    serial: str,
    pdd_version: str,
    screenshot_path: Path,
    hierarchy_path: Path,
) -> dict[str, Any]:
    """只序列化审计摘要；页面内容留在 XML 文件，不进入日志或 manifest。"""

    return {
        "schema_version": 1,
        "captured_at": datetime.now(UTC).isoformat(),
        "channel": "wifi" if ":" in serial else "usb",
        "serial_sha256": sha256(serial.encode("utf-8")).hexdigest(),
        "device": {
            "model": inspection.model,
            "android_version": inspection.android_version,
            "pdd_package": PDD_PACKAGE,
            "pdd_version": pdd_version,
        },
        "artifacts": [
            {"path": screenshot_path.name, "sha256": _sha256_file(screenshot_path)},
            {"path": hierarchy_path.name, "sha256": _sha256_file(hierarchy_path)},
        ],
    }


def _sha256_file(path: Path) -> str:
    digest = sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()