本次提交将woc-bridge项目重构为模块化包结构,按职责拆分多个子域: 1. 新增models层定义所有Pydantic数据模型与统一错误体系 2. 拆分db/ui/messaging/routes等业务域模块 3. 实现基础API路由:状态查询、截图、登录、媒体获取等 4. 重构tools脚本的模块导入路径 5. 补充版本号与能力清单定义 6. 完善全局配置与依赖管理 整体完成项目从单文件脚本到可维护的包结构迁移,为后续功能开发打下基础。
150 lines
4.9 KiB
Python
150 lines
4.9 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
|
||
|
||
|
||
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
|