hospital-front-agent/agent/dispatcher.py

# -*- coding: utf-8 -*-
"""下行命令分发器。

B 端通过 WebSocket 下发 JSON 命令：{"type":"CMD_TYPE","payload":{...}}

已支持命令：
- PING            探活（立即回 pong）
- GET_HEALTH      查询主程序健康
- GET_VERSION     查询主程序版本
- GET_LOG_FILES   日志文件列表
- LOG_TAIL        日志尾部（payload: file, lines）
- GET_LOG_DOWNLOAD 日志全量下载（payload: file）
- RESTART         重启主程序（优雅停机后由 ProcessGuard 拉起）
- GET_CONFIG      读取医院配置文件（payload: hospital）
- SAVE_CONFIG     写回医院配置文件（payload: hospital, content）
- UPGRADE         升级主程序
- GET_APIS        查询主程序 /order 接口清单
- HTTP_PROXY      代理调用前置机任意接口（穿透白名单）

所有命令执行结果通过 AgentClient 上行（CMD_RESULT）。
"""

import logging
import threading
import time

log = logging.getLogger(__name__)


class CommandDispatcher:
    """命令分发：解析 type → 执行 → 回执 CMD_RESULT（透传 cmdId）。"""

    def __init__(self, config, main_app_client):
        self.config = config
        self.main_app_client = main_app_client

        self._client = None                  # 由 main 注入（避免构造环）
        self._upgrade_command = None         # 升级命令（懒加载）
        self._upgrade_lock = threading.Lock()

    def set_client(self, client):
        """注入 AgentClient（main 组装阶段调用，避免构造环）。"""
        self._client = client

    def _get_upgrade_command(self):
        """懒加载升级命令。"""
        with self._upgrade_lock:
            if self._upgrade_command is None:
                from agent.upgrade import UpgradeCommand
                self._upgrade_command = UpgradeCommand(self.config, self.main_app_client, self._client)
            return self._upgrade_command

    # ================= 分发 =================

    def dispatch(self, msg_type, payload):
        """分发并执行命令。payload 为 dict 或 None。"""
        log.info("收到命令：%s", msg_type)
        payload = payload if isinstance(payload, dict) else None

        if msg_type == "PING":
            # 命令协议：B 端下发的 payload 内可带 cmdId，回执时透传以便并发匹配
            result = {"pong": True, "ts": int(time.time())}

        elif msg_type == "GET_HEALTH":
            result = self.main_app_client.health()

        elif msg_type == "GET_VERSION":
            result = self.main_app_client.version()

        elif msg_type == "GET_LOG_FILES":
            result = {"ok": True, "data": self.main_app_client.log_files()}

        elif msg_type == "LOG_TAIL":
            file = (payload or {}).get("file") or "info.log"
            lines = (payload or {}).get("lines") or 200
            result = {"ok": True, "data": self.main_app_client.log_tail(file, lines)}

        elif msg_type == "RESTART":
            # 优先走 NSSM 服务重启（主程序假死时 HTTP 不可达也能重启成功）；
            # 未配置服务名时回退环回 HTTP 触发主程序优雅重启
            if self.config.main_service_name:
                from agent.process_guard import restart_main_service
                ok, reason = restart_main_service(self.config.nssm_path, self.config.main_service_name)
                result = {
                    "ok": ok,
                    "msg": "已通过 NSSM 重启主程序服务" if ok else f"NSSM 重启失败：{reason}",
                }
            else:
                ok, reason = self.main_app_client.restart(1000)
                result = {
                    "ok": ok,
                    "msg": "已触发优雅重启，等待 ProcessGuard 拉起" if ok else f"重启请求失败：{reason}",
                }

        elif msg_type == "GET_CONFIG":
            hospital = (payload or {}).get("hospital") or ""
            content = self.main_app_client.config_file(hospital)
            result = {"ok": content is not None, "data": content}

        elif msg_type == "SAVE_CONFIG":
            hospital = (payload or {}).get("hospital") or ""
            content = (payload or {}).get("content") or ""
            result = {"ok": self.main_app_client.save_config_file(hospital, content)}

        elif msg_type == "UPGRADE":
            # 异步执行（升级耗时长，避免阻塞 WS 消息循环）
            version = (payload or {}).get("version") or "unknown"
            download_url = (payload or {}).get("downloadUrl")
            sha256 = (payload or {}).get("sha256")
            if not download_url or not download_url.strip():
                result = {"ok": False, "msg": "缺少 downloadUrl"}
            else:
                threading.Thread(
                    target=self._run_upgrade,
                    args=(version, download_url, sha256),
                    name="upgrade",
                    daemon=True,
                ).start()
                result = {"ok": True, "msg": "升级已开始（异步执行，结果将通过 UPGRADE_RESULT 上报）"}

        elif msg_type == "GET_APIS":
            # 返回主程序 /order 接口清单（供 B 端接口测试工具加载）
            result = {"ok": True, "data": self.main_app_client.get_apis()}

        elif msg_type == "GET_LOG_DOWNLOAD":
            # 日志全量下载（供 B 端浏览器下载，Agent 自动带 admin token 调主程序）
            file = (payload or {}).get("file") or "info.log"
            content = self.main_app_client.log_download(file)
            result = {"ok": content is not None, "data": content}

        elif msg_type == "HTTP_PROXY":
            # 代理调用前置机任意接口（穿透白名单，本机环回）
            p = payload or {}
            result = self.main_app_client.http_proxy(
                method=p.get("method") or "GET",
                path=p.get("path") or "/",
                query=p.get("query"),
                body=p.get("body"),
                extra_headers=p.get("headers"),
            )

        else:
            result = {"ok": False, "msg": f"未知命令：{msg_type}"}

        envelope = {"cmd": msg_type}
        # 透传 B 端下发的 cmdId（并发命令匹配用）
        if payload is not None and payload.get("cmdId") is not None:
            envelope["cmdId"] = payload.get("cmdId")
        envelope["result"] = result
        self._client.send("CMD_RESULT", envelope)

    def _run_upgrade(self, version, download_url, sha256):
        """升级线程体：执行并上报 UPGRADE_RESULT。"""
        r = self._get_upgrade_command().upgrade(version, download_url, sha256)
        self._client.send("UPGRADE_RESULT", r)