auto_control/mcp_server/platform_client.py

"""平台（auto_control）HTTP 客户端：登录会话 + CSRF + API 封装。

会话失效自动重登；CSRF token 每次登录后获取，POST 必带。
只读接口（GET）与写接口（POST）分离封装，写操作由 MCP 门控层决定是否调用。
"""
import httpx
import logging

_log = logging.getLogger("mcp.platform")

_CSRF_HEADER = "X-CSRF-Token"


class PlatformError(Exception):
    def __init__(self, code, message):
        super().__init__(message)
        self.code = code
        self.message = message


class PlatformClient:
    def __init__(self, url, user, password, timeout=30.0):
        self.url = url.rstrip("/")
        self.user = user
        self.password = password
        self._client = httpx.Client(base_url=self.url, timeout=timeout,
                                    follow_redirects=True)
        self._csrf = ""
        self._login()

    # ---------- 会话 ----------
    def _login(self):
        """登录平台，拿会话 cookie + CSRF token。"""
        r = self._client.post("/login", data={
            "username": self.user, "password": self.password})
        if r.status_code != 302 and r.status_code != 200:
            raise PlatformError("platform_unavailable",
                                f"平台登录失败 HTTP {r.status_code}: {r.text[:120]}")
        csrf = self._client.get("/api/csrf")
        if csrf.status_code == 200:
            self._csrf = (csrf.json() or {}).get("token", "")
        _log.info("平台登录成功: %s", self.user)

    def _ensure_session(self):
        """会话过期（401/403/302 到登录页）时重登。"""
        r = self._client.get("/api/status")
        if r.status_code in (401, 403) or r.url.path.endswith("/login"):
            _log.warning("会话失效，重新登录")
            self._login()
            return True
        return False

    # ---------- 基础请求 ----------
    def _get(self, path, **params):
        self._ensure_session()
        return self._client.get(path, params=params)

    def _post(self, path, json=None):
        self._ensure_session()
        headers = {_CSRF_HEADER: self._csrf} if self._csrf else {}
        return self._client.post(path, json=json or {}, headers=headers)

    # ---------- 平台 API 封装（M0） ----------
    def list_devices(self):
        """设备列表与状态（GET /api/status）。"""
        r = self._get("/api/status")
        if r.status_code != 200:
            raise PlatformError("platform_unavailable",
                                f"/api/status HTTP {r.status_code}")
        j = r.json() or {}
        return j.get("devices", [])

    def screenshot(self, serial):
        """截图（GET /api/screen/thumb），返回 (jpeg_bytes, screen_state)。"""
        r = self._get("/api/screen/thumb", serial=serial)
        if r.status_code == 503:
            raise PlatformError("device_offline", r.text[:120])
        if r.status_code != 200:
            raise PlatformError("platform_unavailable",
                                f"截图 HTTP {r.status_code}")
        return r.content, r.headers.get("X-Screen-State", "unknown")

    def snapshot(self, serial):
        """一次取齐截图+元素树（GET /api/uiauto/snapshot）。

        平台侧用**同一个 u2 连接**背靠背 dump+screenshot，并做双截图校验
        （unstable 表示"抓取期间界面在变化"）。返回平台 JSON：
        {ok, image(data-url), width, height, elements, unstable, screen_state, cost_ms}
        """
        r = self._get("/api/uiauto/snapshot", serial=serial, _t=0)
        if r.status_code in (502, 503):
            raise PlatformError("device_offline", r.text[:160])
        if r.status_code != 200:
            raise PlatformError("platform_unavailable",
                                f"快照 HTTP {r.status_code}")
        j = r.json() or {}
        if not j.get("ok"):
            raise PlatformError("device_offline", str(j.get("error", "快照失败"))[:160])
        return j

    def screen_size(self, serial):
        """屏幕原生分辨率（GET /api/screen/size），返回 (w, h)。"""
        r = self._get("/api/screen/size", serial=serial)
        if r.status_code == 503:
            raise PlatformError("device_offline", r.text[:120])
        if r.status_code != 200:
            raise PlatformError("platform_unavailable",
                                f"分辨率 HTTP {r.status_code}")
        j = r.json() or {}
        if not j.get("ok"):
            raise PlatformError("device_offline", str(j.get("error", "取分辨率失败"))[:120])
        return int(j["width"]), int(j["height"])

    def tap(self, serial, x, y, snap=False):
        """点击（POST /api/screen/tap）。

        snap=True：点落在可点击元素内则吸附到元素中心（AI 粗略坐标也能点准）。
        返回平台 JSON（含 snapped/x/y/label）。
        """
        r = self._post("/api/screen/tap", json={"serial": serial,
                                                "x": int(x), "y": int(y),
                                                "snap": 1 if snap else 0})
        return self._check_op(r, "tap")

    def tap_text(self, serial, text):
        """按屏幕文字点击（平台解析：UI 树子串匹配 → OCR 兜底）。

        返回 {ok, found, method, matched, x, y}——found=false 是业务结果
        （屏幕无该文字），非设备错误；设备离线/不可达仍抛 PlatformError。
        """
        r = self._post("/api/screen/tap_text",
                       json={"serial": serial, "text": str(text)})
        if r.status_code == 503:
            raise PlatformError("device_offline", r.text[:120])
        if r.status_code != 200:
            raise PlatformError("platform_unavailable",
                                f"tap_text HTTP {r.status_code}: {r.text[:120]}")
        return r.json() or {}

    def swipe(self, serial, x1, y1, x2, y2, duration=0.2):
        """滑动（POST /api/screen/swipe）。"""
        r = self._post("/api/screen/swipe", json={
            "serial": serial, "x1": int(x1), "y1": int(y1),
            "x2": int(x2), "y2": int(y2),
            "duration": float(duration)})
        return self._check_op(r, "swipe")

    def ui_elements(self, serial):
        """UI 元素树（GET /api/uiauto/elements，uiautodev 服务）。"""
        r = self._get("/api/uiauto/elements", serial=serial)
        if r.status_code == 503:
            raise PlatformError("device_offline", r.text[:120])
        if r.status_code != 200:
            raise PlatformError("platform_unavailable",
                                f"元素树 HTTP {r.status_code}")
        j = r.json() or {}
        if not j.get("ok"):
            raise PlatformError("device_offline", str(j.get("error", "取元素失败"))[:120])
        return j.get("elements", [])

    def wake(self, serial):
        """亮屏并解锁（POST /api/device/screen_all mode=on）。"""
        r = self._post("/api/device/screen_all",
                       json={"mode": "on", "serials": [serial]})
        return self._check_op(r, "wake")

    def sleep(self, serial):
        """熄屏（POST /api/device/screen_all mode=off）。"""
        r = self._post("/api/device/screen_all",
                       json={"mode": "off", "serials": [serial]})
        return self._check_op(r, "sleep")

    def list_tasks(self):
        """任务计划列表（GET /api/jobs）。"""
        r = self._get("/api/jobs")
        if r.status_code != 200:
            raise PlatformError("platform_unavailable",
                                f"/api/jobs HTTP {r.status_code}")
        j = r.json() or {}
        tasks = []
        for t in j.get("jobs") or []:
            tasks.append({"id": t.get("id"), "name": t.get("name"),
                          "task_type": t.get("task_type"),
                          "enabled": t.get("enabled"),
                          "schedule": (t.get("schedule") or {}).get("mode", "")})
        return tasks

    def press_key(self, serial, key):
        """按键（POST /api/screen/key）。"""
        r = self._post("/api/screen/key",
                       json={"serial": serial, "key": key})
        return self._check_op(r, "key")

    @staticmethod
    def _check_op(r, name):
        if r.status_code == 503:
            raise PlatformError("device_offline", r.text[:120])
        if r.status_code != 200:
            raise PlatformError("platform_unavailable",
                                f"{name} HTTP {r.status_code}: {r.text[:120]}")
        j = r.json() or {}
        if not j.get("ok"):
            raise PlatformError("device_offline", str(j.get("error", "操作失败"))[:120])
        return j