Kris
ed2ee086de
主要变更： 1. 新增UIActionScheduler实现优先级UI任务调度，支持CRITICAL/HIGH/NORMAL/LOW四级优先级 2. 替换原有SendQueue为UIActionScheduler，统一所有UI操作的调度逻辑 3. 为截图、登录、重启等API添加调度器封装，支持超时、限流与fast-fail机制 4. 新增Prometheus监控指标，统计UI任务执行、等待耗时、超时、限流与快速失败情况 5. 为QrCapture添加命令执行超时保护，避免X11操作卡死 6. 扩展StatusResponse与状态接口，暴露调度器指标 7. 新增完整的单元测试覆盖调度器功能 8. 为看门狗注入调度器，实现kill前队列清空与快速失败窗口
163 lines
5.3 KiB
Python
163 lines
5.3 KiB
Python
"""微信窗口二维码截图能力。
|
||
|
||
封装屏幕截图（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 <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
|