SimpleRemoter/server/2015Remote/McpServer.h

#pragma once

#include <Windows.h>
#include <string>
#include <thread>
#include <atomic>
#include <chrono>
#include <mutex>
#include <condition_variable>
#include <map>
#include <vector>
#include <cstdint>

// httplib 与 Windows 头部的 min/max 宏冲突，按 file_server.h 的既有约定处理。
#undef min
#undef max
#include "httplib.h"
#ifndef max
#define max(a,b) (((a) > (b)) ? (a) : (b))
#endif
#ifndef min
#define min(a,b) (((a) < (b)) ? (a) : (b))
#endif

#pragma comment(lib, "ws2_32.lib")

class CMy2015RemoteDlg;
class context;

// MCP (Model Context Protocol) 服务端：httplib 封装 + JSON-RPC 2.0 分发 + 静态 token 校验。
// 与 CWebService 平级、互不依赖；默认禁用，经「扩展 → MCP设置」开启后监听
// （默认 127.0.0.1:6544，仅本机回环）。见 docs/Mcp_Design.md。
class CMcpServer {
public:
    static CMcpServer& Instance();

    void SetParentDlg(CMy2015RemoteDlg* pDlg) { m_parent = pDlg; }
    void SetToken(const std::string& token) { m_token = token; }

    bool Start(const std::string& bind, int port);
    void Stop();
    bool IsRunning() const { return m_running.load(); }

    // ===== P2b：无请求 id 的一次性响应 → 每 host 单飞行 =====
    // 进程/窗口列表（TOKEN_PSLIST/TOKEN_WSLIST，一次性子链接）与历史活动
    // （TOKEN_REPORT_ACTIVITY，主连接 RPC）都不带请求关联，同一 host 同一时刻
    // 只允许一个挂起请求，避免同 host 并发请求响应归属歧义。
    // 注意：单飞行只约束 MCP 侧同 host 并发，不约束与 MFC 对话框并存（各用独立子链接）。

    // 该 host 是否已有挂起请求（MessageHandle 在 TOKEN_PSLIST/TOKEN_WSLIST 分支调用）。
    bool IsPending(uint64_t device_id);

    // 拷贝响应缓冲到挂起结果并唤醒等待者（在 MessageHandle 内同步调用，数据此时仍在
    // context 缓冲中、IO 线程被阻塞，拷贝是安全的）。
    void TakeMainResponse(uint64_t device_id, const BYTE* data, ULONG len);

    // 工具线程：登记挂起（false = 该 host 已有挂起请求，设备忙）。
    bool BeginPending(uint64_t device_id, const std::string& tool);

    // 工具线程：等待响应；成功返回 true 并把缓冲写入 out，超时/失败返回 false（并清理）。
    bool WaitPending(uint64_t device_id, std::vector<BYTE>& out, int timeoutMs);

    // 工具线程：清理挂起状态（发送失败等提前退出路径）。
    void ClearPending(uint64_t device_id);

    // ===== P2c：list_files / get_screenshot 扩展 =====

    // 登记挂起（带目录路径；list_files 专用，path 为空表示只列盘）。
    bool BeginPending(uint64_t device_id, const std::string& tool, const std::string& path);

    // get_screenshot：工具线程生成 16 位 reqId（跳过 0，0 表示未设置），用于丢弃过期响应。
    uint16_t NextPreviewReqId();

    // get_screenshot：登记期望的 reqId（必须在发送 COMMAND_SCREEN_PREVIEW_REQ 前调用）。
    void SetPendingReqId(uint64_t device_id, uint16_t reqId);

    // get_screenshot：MessageHandle(TOKEN_SCREEN_PREVIEW_RSP) 在 IO 线程同步调用；
    // 挂起命中且 reqId 一致则拷贝响应并返回 true，否则 false（回落到 MFC 预览路径）。
    bool TakePreviewResponse(uint64_t device_id, uint16_t reqId, const BYTE* data, ULONG len);

    // list_files：MessageHandle(TOKEN_DRIVE_LIST) 调用。
    // 返回 true = 已取走盘列表（调用方 CancelIO）；false = 已转向下发 COMMAND_LIST_FILES、
    // 等 TOKEN_FILE_LIST（调用方不关子链接）。
    bool OnDriveList(uint64_t device_id, context* subCtx, const BYTE* buf, ULONG len);

    // ===== P3：list_registry（注册表查询，三阶段：COMMAND_REGEDIT → TOKEN_REGEDIT →
    // COMMAND_REG_FIND → TOKEN_REG_PATH + TOKEN_REG_KEY）=====

    // MessageHandle(TOKEN_REGEDIT) 调用：挂起命中且为 list_registry 则下发 COMMAND_REG_FIND
    // 并返回 true（接管子链接，不打开 MFC 对话框）；否则 false（回落 MFC）。
    bool OnRegeditReady(uint64_t device_id, context* subCtx);

    // MessageHandle(TOKEN_REG_PATH / TOKEN_REG_KEY) 调用：两包都到齐后置 done 并唤醒等待者。
    void TakeRegPath(uint64_t device_id, const BYTE* data, ULONG len);
    void TakeRegKey(uint64_t device_id, const BYTE* data, ULONG len);

    // 工具线程：等待注册表两包（子键 + 值）；成功返回 true 并把两包分别写入 out，超时返回 false。
    bool WaitPendingRegistry(uint64_t device_id, std::vector<BYTE>& subkeys,
                             std::vector<BYTE>& values, int timeoutMs);

    // ===== P3：exec_command（Windows 一次性远程命令，复刻 Web 终端链路）=====
    // 流程：工具线程 BeginTermPending → 主连接下发 COMMAND_SHELL → 客户端 shell 子连接回
    // TOKEN_TERMINAL_START / TOKEN_SHELL_START → RegisterTerminalContext 接管 → 发
    // COMMAND_NEXT(+PTY resize) 启动输出回流 → 工具线程发哨兵命令行 → OnTerminalData 收集
    // 并做哨兵检测 → WaitTerminalDone 取 stdout + exit_code → CancelIO 关子链接。
    // 单设备单终端（MVP）：BeginTermPending 命中已有会话即拒绝，避免同设备并发归属歧义。

    // 该 host 是否已有待接管的终端会话（MessageHandle TOKEN_*_START 分支调用）。
    bool IsTermPending(uint64_t device_id);

    // 终端子连接就绪：接管（登记 subCtx/isPty、清 pending、发 COMMAND_NEXT 唤醒回流）。
    void RegisterTerminalContext(uint64_t device_id, context* subCtx, bool isPty);

    // 终端子连接是否已被 MCP 接管（MessageHandle 顶部据此把 shell 输出路由到 OnTerminalData）。
    bool IsTerminalContext(context* subCtx);

    // 追加 shell 输出并做哨兵检测（命中置 done 并唤醒等待者）。
    void OnTerminalData(context* subCtx, const BYTE* data, ULONG len);

    // shell 进程退出（TOKEN_TERMINAL_CLOSE）。
    void OnTerminalClosed(context* subCtx);

    // 工具线程：登记终端挂起（false = 该 host 已有终端会话）。
    bool BeginTermPending(uint64_t device_id, const std::string& command, const std::string& nonce);

    // 工具线程：等待子连接回 START 并接管（true = 已接管，输出 subCtx/isPty）。
    bool WaitTerminalReady(uint64_t device_id, context*& subCtx, bool& isPty, int timeoutMs);

    // 工具线程：等待输出收集完成（哨兵命中 / 进程退出 / 超时）。
    // 成功 true：out=截断到哨兵前的原始字节、exitCode=0/1（进程退出无哨兵时 -1）、closed=是否进程退出。
    bool WaitTerminalDone(uint64_t device_id, std::vector<BYTE>& out, int& exitCode, bool& closed, int timeoutMs);

    // 工具线程：清理终端挂起（发送失败等提前退出路径）。
    void ClearTermPending(uint64_t device_id);

    // 安全配置（启动时由 CMy2015RemoteDlg 读 THIS_CFG 后设置）。
    void SetReadonly(bool readonly) { m_readonly = readonly; }
    void SetCmdWhitelist(const std::string& whitelist) { m_cmdWhitelist = whitelist; }
    bool IsReadonly() const { return m_readonly; }
    const std::string& GetCmdWhitelist() const { return m_cmdWhitelist; }

private:
    CMcpServer();
    ~CMcpServer();
    CMcpServer(const CMcpServer&) = delete;
    CMcpServer& operator=(const CMcpServer&) = delete;

    // POST /mcp 处理器：token 校验 + JSON-RPC 分发。
    void HandleMcp(const httplib::Request& req, httplib::Response& res);

    httplib::Server    m_server;
    std::thread        m_thread;
    std::atomic<bool>  m_running{false};
    std::string        m_token;
    CMy2015RemoteDlg*  m_parent = nullptr;

    // 挂起请求注册表（受 m_PendingMutex 保护，键 = device_id）。
    struct PendingRequest {
        std::string       tool;
        std::string       path;              // list_files：目录路径；list_registry：rootToken(1B)+相对子键路径
        uint16_t          expectedReqId = 0; // get_screenshot：期望的预览 reqId（0 = 未设置）
        std::vector<BYTE> data;
        std::vector<BYTE> regPath;           // list_registry：TOKEN_REG_PATH 子键包（含 token 字节）
        std::vector<BYTE> regKey;            // list_registry：TOKEN_REG_KEY 值包（含 token 字节）
        bool              pathDone = false;  // list_registry：已收到 TOKEN_REG_PATH
        bool              keyDone  = false;  // list_registry：已收到 TOKEN_REG_KEY
        bool              done = false;
    };
    std::mutex                   m_PendingMutex;
    std::condition_variable      m_PendingCv;
    std::map<uint64_t, PendingRequest> m_Pending;

    // get_screenshot 的 reqId 发生器（16 位，跳过 0）。与 MFC 预览的 m_PreviewReqId 各自独立计数。
    std::atomic<uint16_t>        m_PreviewReqId{1};

    // ===== P3：exec_command 终端会话（受 m_TermMutex 保护）=====
    struct TermSession {
        bool              started = false;   // false=已发 COMMAND_SHELL 待 START；true=已接管
        context*          subCtx  = nullptr; // shell 子连接上下文
        bool              isPty   = false;   // true=ConPTY(UTF-8)；false=老 cmd 管道(GBK)
        std::string       command;           // 原始 UTF-8 命令（审计用）
        std::string       nonce;             // 哨兵随机串
        std::vector<BYTE> data;              // 收集的原始 shell 输出
        size_t            sentPos = 0;       // 哨兵在 data 中的位置
        int               exitCode = -1;
        bool              done = false;      // 哨兵命中
        bool              closed = false;    // TOKEN_TERMINAL_CLOSE（进程退出）
    };
    std::mutex                       m_TermMutex;
    std::condition_variable          m_TermCv;
    std::map<uint64_t, TermSession>  m_TermSessions;         // device_id → 会话
    std::map<context*, uint64_t>     m_TermContextToDevice;  // subCtx → device_id（顶部路由）

    bool        m_readonly = true;
    std::string m_cmdWhitelist;
};

// 全局访问器（仿 WebService()，见 WebService.h 末尾）
inline CMcpServer& McpServer() { return CMcpServer::Instance(); }

// 生成 32 hex（128-bit）随机 token；供运行时兜底与「MCP设置」对话框预填复用。
std::string GenerateRandomToken();