"""微信窗口二维码截图能力。 封装屏幕截图(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 woc_bridge.models import BridgeError _CMD_TIMEOUT_SEC = 5.0 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)。 所有子进程执行均受 _CMD_TIMEOUT_SEC 超时保护,避免 scrot 等命令卡死 导致无限阻塞(与并发 UI 操作竞争 X11 时可能发生)。 """ proc = await asyncio.create_subprocess_exec( *args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, env=self._env(), ) try: stdout, stderr = await asyncio.wait_for( proc.communicate(), timeout=_CMD_TIMEOUT_SEC ) return proc.returncode, stdout, stderr except asyncio.TimeoutError: proc.kill() await proc.wait() raise async def _capture_full_png(self, dest_path: str) -> None: """用 `scrot ` 截取整个屏幕到 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