WechatOnCloud/bridge/xdotool_driver.py

"""xdotool / xclip 异步驱动。

封装对微信窗口的所有 X11 自动化操作：查找窗口、检测登录态、激活窗口、
通过 Ctrl+F 搜索会话并粘贴发送文本。所有外部命令通过 asyncio 子进程执行，
不阻塞 event loop。
"""

from __future__ import annotations

import asyncio
import base64
import logging
import os
import random
import time
from typing import Optional

from models import BridgeError, LoginState


# 模块级 logger：与 server.py 同名，便于 s6 统一收日志
logger = logging.getLogger("woc-bridge")


class XdotoolDriver:
    """xdotool/xclip 异步驱动。

    所有方法均为 async，内部使用 asyncio.create_subprocess_exec 调用
    xdotool / xclip / pgrep 等命令，并通过 DISPLAY 环境变量指定 X server。
    """

    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],
        *,
        input_bytes: Optional[bytes] = None,
    ) -> tuple[int, bytes, bytes]:
        """执行一条命令并返回 (returncode, stdout, stderr)。

        Args:
            args: 命令及其参数列表，如 ["xdotool", "search", "--name", "微信"]
            input_bytes: 需要写入 stdin 的字节
        """
        proc = await asyncio.create_subprocess_exec(
            *args,
            stdin=asyncio.subprocess.PIPE if input_bytes is not None else None,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=self._env(),
        )
        stdout, stderr = await proc.communicate(input=input_bytes)
        return proc.returncode, stdout, stderr

    async def _key(self, key: str) -> None:
        """执行 xdotool key <key>。"""
        await self._run(["xdotool", "key", key])

    async def _sleep(self, seconds: float) -> None:
        """asyncio.sleep 封装，便于测试与统一调速。"""
        await asyncio.sleep(seconds)

    # ------------------------------------------------------------------
    # 窗口与进程检测
    # ------------------------------------------------------------------
    async def find_wechat_window(self) -> Optional[int]:
        """用 xdotool search --name "微信" 查找微信窗口 ID。

        Returns:
            窗口 ID（int），找不到返回 None
        """
        returncode, stdout, stderr = await self._run(
            ["xdotool", "search", "--name", "微信"]
        )
        if returncode != 0:
            return None
        text = stdout.decode(errors="ignore").strip()
        if not text:
            return None
        # 取第一个匹配的窗口 ID
        first_line = text.splitlines()[0].strip()
        try:
            return int(first_line)
        except ValueError:
            return None

    async def is_wechat_running(self) -> bool:
        """判断微信进程是否运行。

        用 pgrep -x 精确匹配进程名 wechat（避免 -f 子串匹配误中
        bridge 自身命令行里的 /config/.config/xwechat 路径）。
        若 pgrep 不可用，回落到 xdotool search（窗口存在即视为进程运行）。
        """
        # 用 -x 精确进程名匹配，避免 -f 子串误中 bridge 自身
        # 微信 4.0 Linux 进程名通常为 "wechat"
        for name in ("wechat", "WeChat"):
            returncode, stdout, _ = await self._run(["pgrep", "-x", name])
            if returncode == 0 and stdout.decode(errors="ignore").strip():
                return True
        # 回落：窗口存在即视为运行
        window_id = await self.find_wechat_window()
        return window_id is not None

    async def activate_window(self, window_id: Optional[int] = None) -> None:
        """激活微信窗口。

        Args:
            window_id: 指定窗口 ID；为 None 时自动查找
        """
        if window_id is None:
            window_id = await self.find_wechat_window()
        if window_id is None:
            raise BridgeError(
                code="WINDOW_NOT_FOUND",
                message="未找到微信窗口，无法激活",
            )
        await self._run(["xdotool", "windowactivate", "--sync", str(window_id)])

    async def _activate_window_fast(self, window_id: Optional[int] = None) -> None:
        """非阻塞激活微信窗口，避免 --sync 在 VNC 无人操作时死等。"""
        if window_id is None:
            window_id = await self.find_wechat_window()
        if window_id is None:
            raise BridgeError(
                code="WINDOW_NOT_FOUND",
                message="未找到微信窗口，无法激活",
            )
        await self._run(["xdotool", "windowactivate", str(window_id)])

    async def restart_wechat(self, timeout_sec: int = 30) -> int:
        """重启微信进程：发 SIGTERM 让微信优雅退出，等待 autostart 拉起新进程。

        微信由容器内 autostart 脚本常驻拉起（while true 循环 + sleep 2），
        故 bridge 只需 kill，autostart 会在 2 秒后自动重启。不强制 SIGKILL，
        给微信优雅退出机会（避免 DB 写入未刷盘）。

        Args:
            timeout_sec: 等待新进程出现的最大秒数，默认 30

        Returns:
            新进程的 PID

        Raises:
            BridgeError(RESTART_TIMEOUT): 超时未检测到新进程
        """
        # 1. 获取当前 PID（若有）
        returncode, stdout, _ = await self._run(["pgrep", "-x", "wechat"])
        old_pids: set[int] = set()
        if returncode == 0:
            for line in stdout.decode(errors="ignore").splitlines():
                line = line.strip()
                if line:
                    try:
                        old_pids.add(int(line))
                    except ValueError:
                        pass

        # 2. 若有旧进程，发 SIGTERM
        for pid in old_pids:
            # 用 kill（shell 内建）发 SIGTERM，不强制 SIGKILL
            await self._run(["kill", "-TERM", str(pid)])

        # 3. 等待新进程出现（autostart 会在 2 秒后拉起）
        deadline = time.monotonic() + timeout_sec
        while time.monotonic() < deadline:
            await self._sleep(1.0)
            rc, stdout, _ = await self._run(["pgrep", "-x", "wechat"])
            if rc == 0:
                for line in stdout.decode(errors="ignore").splitlines():
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        new_pid = int(line)
                    except ValueError:
                        continue
                    # 确认是新 PID（不在旧 PID 集合中）
                    if new_pid not in old_pids:
                        return new_pid

        # 4. 超时
        raise BridgeError(
            code="RESTART_TIMEOUT",
            message=f"等待微信重启超时（{timeout_sec}s）",
        )

    async def detect_login_state(self) -> str:
        """检测登录态（启发式）。

        MVP 阶段判定逻辑：
        - 微信进程未运行 → not_running
        - 进程运行但窗口未找到 → not_logged_in
        - 窗口存在 → logged_in

        Returns:
            LoginState 枚举值（字符串）
        """
        running = await self.is_wechat_running()
        if not running:
            return LoginState.NOT_RUNNING.value
        window_id = await self.find_wechat_window()
        if window_id is None:
            return LoginState.NOT_LOGGED_IN.value
        # MVP：窗口存在即视为已登录
        return LoginState.LOGGED_IN.value

    # ------------------------------------------------------------------
    # 显式启动微信（供 diagnostic autofix 使用）
    # ------------------------------------------------------------------
    # 微信可执行文件固定路径（与 docker/app-defs.sh 中 wechat 类型一致）。
    # 仅对 wechat 类型实例有意义；Telegram/Chromium/自定义应用的生命周期
    # 不应通过 bridge autofix 干预。
    _WECHAT_BIN = "/config/wechat/opt/wechat/wechat"

    async def start_wechat(self, timeout_sec: int = 10) -> Optional[int]:
        """显式启动微信进程，返回新 PID 或 None。

        用于 diagnostic autofix：pkill 后 autostart watchdog 未拉起时，
        bridge 显式启动作为兜底。避免引入 shell 依赖，直接执行二进制。

        竞态规避：
        1. 启动前再次 pgrep，已有进程则直接返回（不重复启动）
        2. 启动后轮询 pgrep，等待新 PID 出现

        Args:
            timeout_sec: 等待新进程出现的最大秒数，默认 10

        Returns:
            新进程 PID；启动失败或超时返回 None
        """
        # 1. 启动前再 pgrep，已有则跳过（避免与 autostart 竞态）
        existing_pid = await self.check_wechat_pid()
        if existing_pid is not None:
            return existing_pid

        # 2. 显式启动（disown 语义：不等待进程结束）
        #    DISPLAY/XAUTHORITY 由 _env() 透传，确保 wechat 能连接 X
        try:
            await asyncio.create_subprocess_exec(
                self._WECHAT_BIN,
                stdout=asyncio.subprocess.DEVNULL,
                stderr=asyncio.subprocess.DEVNULL,
                env=self._env(),
                # 启动后立即 detach，不让 wechat 成为 bridge 的子进程
                # （避免 bridge 退出时连带杀掉 wechat）
                start_new_session=True,
            )
        except Exception as e:
            logger.warning("start_wechat: 启动 %s 失败: %r", self._WECHAT_BIN, e)
            return None

        # 3. 等待新 PID 出现（不依赖 subprocess 返回的 pid，因为 wechat
        #    可能 fork 出独立进程，pgrep 比依赖 proc.pid 更可靠）
        deadline = time.monotonic() + timeout_sec
        while time.monotonic() < deadline:
            await self._sleep(1.0)
            new_pid = await self.check_wechat_pid()
            if new_pid is not None:
                return new_pid

        logger.warning("start_wechat: 等待 %ss 内未检测到新进程", timeout_sec)
        return None

    async def check_wechat_pid(self) -> Optional[int]:
        """返回当前运行的微信主进程 PID，无则返回 None。"""
        returncode, stdout, _ = await self._run(["pgrep", "-x", "wechat"])
        if returncode != 0:
            return None
        for line in stdout.decode(errors="ignore").splitlines():
            line = line.strip()
            if not line:
                continue
            try:
                return int(line)
            except ValueError:
                continue
        return None

    # ------------------------------------------------------------------
    # 剪贴板粘贴
    # ------------------------------------------------------------------
    async def _paste_via_xclip(self, text: str) -> None:
        """通过 xclip 写入剪贴板并触发 Ctrl+V 粘贴。

        直接把原始字节写入 xclip stdin（create_subprocess_exec 不经 shell，
        无转义问题，因此无需 base64 编码）。并行等待 stdin 写入与进程退出，
        避免大文本时管道缓冲阻塞导致的死锁。

        Args:
            text: 待粘贴文本

        Raises:
            BridgeError(SEND_FAILED): xclip 退出码非 0
        """
        xclip_proc = await asyncio.create_subprocess_exec(
            "xclip", "-selection", "clipboard",
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.DEVNULL,
            stderr=asyncio.subprocess.PIPE,
            env=self._env(),
        )
        assert xclip_proc.stdin is not None

        async def _feed() -> None:
            try:
                xclip_proc.stdin.write(text.encode("utf-8"))
                await xclip_proc.stdin.drain()
            except (BrokenPipeError, ConnectionResetError):
                # xclip 已退出，忽略
                pass
            finally:
                try:
                    xclip_proc.stdin.close()
                except Exception:
                    pass

        # 并行：喂 stdin + 等进程退出，避免管道缓冲满死锁
        _, (_, stderr) = await asyncio.gather(_feed(), xclip_proc.communicate())
        if xclip_proc.returncode != 0:
            err = stderr.decode(errors="ignore").strip() if stderr else "unknown"
            raise BridgeError(
                code="SEND_FAILED",
                message=f"xclip 写入剪贴板失败 (code={xclip_proc.returncode}): {err}",
            )
        # 触发粘贴
        await self._key("ctrl+v")

    # ------------------------------------------------------------------
    # 会话定位（Ctrl+F 搜索）
    # ------------------------------------------------------------------
    async def _open_session_by_name(
        self,
        name: str,
        timeout_sec: float = 10.0,
    ) -> None:
        """通过微信搜索框定位并进入指定联系人的会话。

        流程：
        1. 非阻塞激活微信窗口
        2. 按 Esc 关闭可能存在的搜索框/弹窗
        3. Ctrl+F 打开搜索
        4. Ctrl+A 全选旧内容，粘贴 display_name
        5. 等待搜索结果渲染
        6. 按 ↓ 选中第一个结果，回车进入会话
        7. 再按 Esc 确保退出搜索模式，焦点落在聊天输入框

        Args:
            name: 用于搜索的显示名（备注/昵称/微信号）
            timeout_sec: 整体超时（秒），默认 10 秒

        Raises:
            BridgeError(WINDOW_NOT_FOUND): 找不到微信窗口
            BridgeError(SEND_FAILED): 超时或 xdotool/xclip 操作失败
        """
        if not name:
            raise BridgeError(
                code="SEND_FAILED",
                message="display_name 不能为空，无法定位会话",
            )

        deadline = time.monotonic() + timeout_sec

        async def _step(coro: Awaitable[None], desc: str) -> None:
            """执行单个 UI 步骤并加超时保护。"""
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise BridgeError(
                    code="SEND_FAILED",
                    message=f"定位会话超时: {desc}",
                )
            try:
                await asyncio.wait_for(coro, timeout=max(1.0, remaining))
            except asyncio.TimeoutError as exc:
                raise BridgeError(
                    code="SEND_FAILED",
                    message=f"定位会话步骤超时: {desc}",
                ) from exc

        # 1. 激活窗口（非阻塞，避免 --sync 死等）
        window_id = await self.find_wechat_window()
        if window_id is None:
            raise BridgeError(
                code="WINDOW_NOT_FOUND",
                message="未找到微信窗口，无法定位会话",
            )
        await _step(self._activate_window_fast(window_id), "激活窗口")
        await _step(self._sleep(0.3), "等待窗口激活")

        # 2. 关闭可能存在的搜索框/弹窗
        await _step(self._key("Escape"), "关闭搜索框")
        await _step(self._sleep(0.2), "等待 Esc 生效")

        # 3. 打开搜索
        await _step(self._key("ctrl+f"), "打开搜索")
        await _step(self._sleep(0.4), "等待搜索框打开")

        # 4. 清空并输入搜索关键词
        await _step(self._key("ctrl+a"), "全选搜索框内容")
        await _step(self._sleep(0.1), "等待全选")
        await _step(self._paste_via_xclip(name), "粘贴搜索关键词")
        await _step(self._sleep(0.8), "等待搜索结果")

        # 5. 选中第一个结果并进入会话
        await _step(self._key("Down"), "选中搜索结果")
        await _step(self._sleep(0.3), "等待选中")
        await _step(self._key("Return"), "进入会话")
        await _step(self._sleep(0.6), "等待会话打开")

        # 6. 退出搜索模式，确保焦点在输入框
        await _step(self._key("Escape"), "退出搜索模式")
        await _step(self._sleep(0.2), "等待焦点稳定")

    # ------------------------------------------------------------------
    # 发送文本
    # ------------------------------------------------------------------
    async def send_text(
        self,
        to_wxid: str,
        content: str,
        display_name: Optional[str] = None,
    ) -> str:
        """发送文本消息。

        流程：
        1. 通过微信搜索框定位会话（使用 display_name，未提供时回退到 to_wxid）
        2. 粘贴文本内容并回车发送
        3. 返回本地生成的 channel_msg_id

        Args:
            to_wxid: 目标 wxid
            content: 文本内容
            display_name: 用于搜索定位会话的显示名（备注/昵称/微信号）

        Returns:
            本地生成的 channel_msg_id，格式 local_<unix秒>_<随机>
        """
        # 1. 定位会话
        await self._open_session_by_name(display_name if display_name else to_wxid)
        # 2. 粘贴内容并发送
        await self._paste_via_xclip(content)
        await self._sleep(0.2)
        await self._key("Return")
        # 3. 生成 local_send_id（本地 ID，非微信原生 msg_id）
        local_send_id = f"local_{int(time.time())}_{random.randint(0, 0xFFFFFF):06x}"
        return local_send_id

    # ------------------------------------------------------------------
    # 发送文件 / 图片
    # ------------------------------------------------------------------
    async def send_file(
        self,
        to_wxid: str,
        file_path: str,
        is_image: bool = False,
        display_name: Optional[str] = None,
    ) -> str:
        """发送文件/图片（MVP 简化版）。

        MVP 策略：检查文件可读 → 进入会话（复用 _open_session_by_name）
        → 发送"[图片/文件] 文件名"提示文本。真实文件传输机制留给后续完善。

        Args:
            to_wxid: 目标 wxid
            file_path: 容器内文件绝对路径
            is_image: 是否为图片
            display_name: 用于搜索定位会话的显示名

        Returns:
            本地生成的 channel_msg_id，格式 local_<unix秒>_<随机>

        Raises:
            BridgeError(INVALID_PARAMS): 文件不存在或不可读
        """
        # 文件存在性校验
        if not file_path or not os.path.isfile(file_path):
            raise BridgeError(
                code="INVALID_PARAMS",
                message=f"文件不存在: {file_path}",
            )
        if not os.access(file_path, os.R_OK):
            raise BridgeError(
                code="INVALID_PARAMS",
                message=f"文件不可读: {file_path}",
            )

        # 1. 定位会话
        await self._open_session_by_name(display_name if display_name else to_wxid)

        # 2. MVP：发送文件路径提示文本（真实文件传输留给后续完善）
        filename = os.path.basename(file_path)
        prefix = "[图片]" if is_image else "[文件]"
        hint = f"{prefix} {filename}"
        await self._paste_via_xclip(hint)
        await self._sleep(0.2)
        await self._key("Return")

        # 3. 生成 local_send_id（本地 ID，非微信原生 msg_id）
        local_send_id = f"local_{int(time.time())}_{random.randint(0, 0xFFFFFF):06x}"
        return local_send_id

    # ------------------------------------------------------------------
    # 退出登录
    # ------------------------------------------------------------------
    async def logout(self) -> None:
        """通过 UI 操作退出微信登录。

        流程：
        1. 检测登录态，若非 logged_in 直接返回（幂等）
        2. 查找并激活微信窗口（复用 window_id）
        3. 获取窗口几何，点击左下角主菜单图标
        4. 等待菜单弹出，用方向键导航到「退出登录」并回车
        5. 若出现确认对话框，按回车确认
        6. 等待 2 秒让 UI 完成切换

        注意：下方点击与方向键坐标/次数均为估算值，需在目标分辨率
        实测后调优（微信 4.0 Linux UI 路径见模块 docstring）。

        Raises:
            BridgeError(LOGOUT_FAILED): 窗口未找到或 UI 操作失败
        """
        # 1. 检测登录态（幂等）
        state = await self.detect_login_state()
        if state != LoginState.LOGGED_IN.value:
            # 未登录，无需退出
            return

        # 2. 查找窗口 ID 并激活（复用 window_id，避免二次 search）
        window_id = await self.find_wechat_window()
        if window_id is None:
            raise BridgeError(
                code="LOGOUT_FAILED",
                message="未找到微信窗口，无法退出登录",
            )
        await self.activate_window(window_id)

        # 3. 获取窗口几何（用已知的 window_id，避免命令链歧义）
        returncode, stdout, _ = await self._run(
            ["xdotool", "getwindowgeometry", "--shell", str(window_id)]
        )
        if returncode != 0:
            raise BridgeError(
                code="LOGOUT_FAILED",
                message="无法获取微信窗口几何信息",
            )
        # 解析 getwindowgeometry --shell 输出（KEY=VALUE 格式）
        geom = {}
        for line in stdout.decode(errors="ignore").splitlines():
            if "=" in line:
                k, v = line.split("=", 1)
                geom[k.strip()] = v.strip()
        try:
            win_x = int(geom.get("X", 0))
            win_y = int(geom.get("Y", 0))
            win_w = int(geom.get("WIDTH", 800))
            win_h = int(geom.get("HEIGHT", 600))
        except ValueError:
            raise BridgeError(
                code="LOGOUT_FAILED",
                message="解析窗口几何信息失败",
            )

        # 4. 点击左下角主菜单图标
        #    估算：窗口左侧栏底部，约 x=win_x+30, y=win_y+win_h-30
        #    该坐标随分辨率/缩放变化，需实测调优。
        menu_x = win_x + 30
        menu_y = win_y + win_h - 30
        rc, _, _ = await self._run(
            ["xdotool", "mousemove", "--sync", str(menu_x), str(menu_y)]
        )
        if rc != 0:
            raise BridgeError(code="LOGOUT_FAILED", message="移动鼠标到主菜单失败")
        # 用 xdotool click 1（左键单击）
        rc, _, _ = await self._run(["xdotool", "click", "1"])
        if rc != 0:
            raise BridgeError(code="LOGOUT_FAILED", message="点击主菜单失败")

        # 5. 等待菜单弹出
        await self._sleep(0.8)

        # 6. 用方向键导航到「退出登录」并回车
        #    微信菜单项顺序通常为：设置 / 切换账号 / 退出登录 / 关闭
        #    退出登录一般在第 3 项，从顶部按 ↓ 2 次到达
        #    该顺序为估算，需实测调优。
        await self._key("Down")
        await self._sleep(0.2)
        await self._key("Down")
        await self._sleep(0.2)
        await self._key("Return")

        # 7. 等待确认对话框
        await self._sleep(0.8)
        # 按回车确认（默认焦点通常在确认按钮）
        await self._key("Return")

        # 8. 等待 UI 完成切换
        await self._sleep(2.0)