WechatOnCloud/bridge/qr_capture.py

"""微信窗口二维码截图能力。

封装屏幕截图（scrot）与二维码区域裁剪（Pillow），供
/api/login/qr/start 与 /api/screenshot 使用。

所有方法为 async，内部用 asyncio 子进程执行截图命令，避免阻塞
event loop。
"""

from __future__ import annotations

import asyncio
import base64
import io
import os
import tempfile

from models import BridgeError


class QrCapture:
    """屏幕截图与二维码裁剪器。

    通过 scrot 截取整个 X11 屏幕，再用 Pillow 按比例裁剪出二维码
    所在区域（中央偏上）。
    """

    # 二维码区域裁剪比例（spec 12.3 节，MVP 阶段固定为中央偏上区域）
    _QR_LEFT = 0.25
    _QR_TOP = 0.2
    _QR_RIGHT = 0.75
    _QR_BOTTOM = 0.7

    def __init__(self, display: str = ":1") -> None:
        """保存 DISPLAY 环境变量值。

        Args:
            display: X server display 地址，如 ":1"
        """
        self.display = display

    def _env(self) -> dict:
        """返回带 DISPLAY 的环境副本。"""
        env = dict(os.environ)
        env["DISPLAY"] = self.display
        return env

    async def _run(self, args: list[str]) -> tuple[int, bytes, bytes]:
        """执行一条命令并返回 (returncode, stdout, stderr)。"""
        proc = await asyncio.create_subprocess_exec(
            *args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=self._env(),
        )
        stdout, stderr = await proc.communicate()
        return proc.returncode, stdout, stderr

    async def _capture_full_png(self, dest_path: str) -> None:
        """用 `scrot <dest>` 截取整个屏幕到 PNG 文件。

        Args:
            dest_path: 目标 PNG 文件路径
        """
        returncode, _, stderr = await self._run(
            ["scrot", dest_path]
        )
        if returncode != 0 or not os.path.exists(dest_path):
            msg = stderr.decode(errors="ignore").strip() or "unknown error"
            raise BridgeError(
                code="BRIDGE_INTERNAL_ERROR",
                message=f"截图失败: {msg}",
            )

    async def capture_qr_code(self) -> str:
        """截取并裁剪二维码区域，返回 data URL。

        流程：
        1. 截图整个屏幕到 /tmp 临时 PNG
        2. 用 Pillow 打开并按比例裁剪二维码区域
        3. 编码为 base64 data URL `data:image/png;base64,...`
        4. 删除临时文件
        5. 返回 data URL
        """
        # 临时文件
        fd, tmp_path = tempfile.mkstemp(prefix="woc_qr_", suffix=".png")
        os.close(fd)
        try:
            await self._capture_full_png(tmp_path)
            # 用 Pillow 裁剪
            data_url = await asyncio.to_thread(self._crop_qr_to_data_url, tmp_path)
            return data_url
        finally:
            if os.path.exists(tmp_path):
                try:
                    os.remove(tmp_path)
                except OSError:
                    pass

    def _crop_qr_to_data_url(self, png_path: str) -> str:
        """同步：打开 PNG，按比例裁剪二维码区域，返回 data URL。"""
        try:
            from PIL import Image  # 局部导入，避免无 Pillow 时影响模块加载
        except ImportError as e:
            raise BridgeError(
                code="BRIDGE_INTERNAL_ERROR",
                message=f"Pillow 未安装: {e}",
            )

        try:
            with Image.open(png_path) as img:
                w, h = img.size
                left = int(w * self._QR_LEFT)
                top = int(h * self._QR_TOP)
                right = int(w * self._QR_RIGHT)
                bottom = int(h * self._QR_BOTTOM)
                cropped = img.crop((left, top, right, bottom))
                buf = io.BytesIO()
                cropped.save(buf, format="PNG")
                b64 = base64.b64encode(buf.getvalue()).decode("ascii")
                return f"data:image/png;base64,{b64}"
        except BridgeError:
            raise
        except Exception as e:
            raise BridgeError(
                code="BRIDGE_INTERNAL_ERROR",
                message=f"裁剪二维码失败: {e}",
            )

    async def capture_full_screenshot(self) -> bytes:
        """截取完整屏幕，返回 PNG bytes（供 /api/screenshot 用）。"""
        fd, tmp_path = tempfile.mkstemp(prefix="woc_shot_", suffix=".png")
        os.close(fd)
        try:
            await self._capture_full_png(tmp_path)
            try:
                with open(tmp_path, "rb") as f:
                    return f.read()
            except OSError as e:
                raise BridgeError(
                    code="BRIDGE_INTERNAL_ERROR",
                    message=f"读取截图文件失败: {e}",
                )
        finally:
            if os.path.exists(tmp_path):
                try:
                    os.remove(tmp_path)
                except OSError:
                    pass