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 <set>
#include <vector>
#include <cstdint>
#include <ctime>

// 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;

// 文件传输结果条目（download_file 输出 files[] 的一项）。
// path 为本机落盘绝对路径（ANSI，输出 JSON 时再转 UTF-8）；sha256 为小写 hex。
struct FileTransferEntry {
    std::string path;
    uint64_t    size = 0;
    std::string sha256;
};

// 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);

    // ===== P4：持久远程终端（terminal_open / terminal_exec / terminal_close）=====
    void SetTerminalEnabled(bool enabled) { m_terminalEnabled = enabled; }
    bool IsTerminalEnabled() const { return m_terminalEnabled; }

    // 登记持久终端挂起（复用单设备单终端约束；sessionId 由调用方生成）。
    bool BeginTermOpen(uint64_t device_id, const std::string& sessionId);

    // 复位单条命令字段并置 busy；输出 subCtx/isPty。false = 不存在/不匹配/未就绪/忙碌/已关。
    bool BeginTermCommand(uint64_t device_id, const std::string& sessionId,
                          const std::string& command, const std::string& nonce,
                          context*& subCtx, bool& isPty);

    // 等待本条命令完成（done || closed || 超时）。done=保持会话；closed/超时=清理会话+路由。
    bool WaitTermCommand(uint64_t device_id, std::vector<BYTE>& out, int& exitCode,
                         bool& closed, int timeoutMs);

    // 关闭持久终端（校验 sessionId；CancelIO + 清理路由/会话）。
    // 返回：0=已关闭；1=会话不存在（幂等）；2=sessionId 不匹配（调用方报错）。
    int CloseTermSession(uint64_t device_id, const std::string& sessionId);

    // idle 回收：关闭 lastActiveAt 超时且非 busy 的持久会话。返回回收数量。
    int SweepIdleTerminals(time_t idleTimeoutSec);

    // ===== P5：MCP 远程控制（remote_open / remote_close / remote_mouse / remote_keyboard）=====
    // 屏幕子连接由 WebService::StartRemoteDesktop 建立的隐藏 CScreenSpyDlg 持有（复用 Web
    // 路径，见 docs/Mcp_RemoteControl_Design.md §6.3）；会话仅记录逻辑状态 + 子连接指针。
    // 注入走该子连接（COMMAND_SCREEN_CONTROL），与终端一样单设备单会话。
    void SetRemoteControlEnabled(bool enabled) { m_remoteControlEnabled = enabled; }
    bool IsRemoteControlEnabled() const { return m_remoteControlEnabled; }

    // 该 host 的屏幕子连接是否被 MCP 远程控制会话持有（OfflineProc 反查用）。
    bool IsScreenCtrlContext(context* subCtx);

    // 登记控制会话（单设备单会话；sessionId 由调用方生成）。false = 该设备已有会话。
    bool BeginScreenCtrlOpen(uint64_t device_id, const std::string& sessionId);

    // 子连接就绪（HasActiveSession && 分辨率到达）后调用：填 subCtx/分辨率并置 started。
    // false = 会话已被并发关闭（调用方应放弃本次 open）。
    bool MarkScreenCtrlReady(uint64_t device_id, const std::string& sessionId,
                             context* subCtx, int screenW, int screenH);

    // 注入前登记（校验 sessionId/started/!busy/!closed）。输出 subCtx/screenW/H。
    // 注入期间 busy=true，保证 close/sweep/断线不与注入并发擦会话（镜像终端 busy 模式，
    // 评审发现 #6）。返回：0=ok；1=不存在/不匹配/未就绪/已断（-32002）；2=忙（-32003）。
    int BeginScreenCtrlAction(uint64_t device_id, const std::string& sessionId,
                              context*& subCtx, int& screenW, int& screenH);

    // 注入后复位 busy；若注入期间屏幕子连接已断（closed），擦会话+路由（防泄漏）。
    void EndScreenCtrlAction(uint64_t device_id, const std::string& sessionId);

    // 关闭控制会话（校验 sessionId；擦路由 + 会话）。屏幕子连接的关闭由调用方锁外执行
    // （CloseWebRemoteDesktopByClientID）。返回 0=已关；1=不存在（幂等）；2=sessionId 不匹配。
    int CloseScreenCtrlSession(uint64_t device_id, const std::string& sessionId);

    // idle 回收：关闭 lastActiveAt 超时的控制会话，并锁外关闭其隐藏对话框。返回回收数量。
    int SweepIdleScreenCtrl(time_t idleTimeoutSec);

    // 屏幕子连接断开（OfflineProc 调用）：擦除该子连接对应的会话 + 路由（幂等）。
    void OnScreenControlClosed(context* subCtx);

    // ===== P6：MCP 文件传输（download_file，V2 协议）=====
    // 下载走 CMD_DOWN_FILES_V2 → 客户端新开流式子连接回传 COMMAND_SEND_FILE_V2 /
    // COMMAND_FILE_COMPLETE_V2（dstClientID==0）。会话按 device_id 键控，单设备单传输；
    // 复用 list_files 的 COMMAND_LIST_DRIVE → TOKEN_DRIVE_LIST 开文件管理器子链接下发命令。
    void SetFileTransferEnabled(bool enabled) { m_fileTransferEnabled = enabled; }
    bool IsFileTransferEnabled() const { return m_fileTransferEnabled; }

    // 该 host 是否有进行中的文件传输会话（MessageHandle 分派守卫）。
    bool IsFileTransferPending(uint64_t device_id);

    // 登记文件传输会话（false = 该 host 已有文件会话或一次性挂起请求）。
    bool BeginFileTransferPending(uint64_t device_id, const std::string& tool,
                                  const std::string& localDir, const std::string& remotePath,
                                  bool overwrite);

    // TOKEN_DRIVE_LIST：识别 download_file，存文件管理器子链接并下发 CMD_DOWN_FILES_V2。
    // 返回 true=已接管（保持子链接，收尾由 ClearFileTransfer）；false=会话已清理（调用方 CancelIO）。
    bool OnDownloadDriveList(uint64_t device_id, context* fmSubCtx);

    // COMMAND_SEND_FILE_V2：解析 chunk，路径校验 + overwrite 判定 + 落盘（RecvFileChunkV2）。
    void OnFileChunkV2(uint64_t device_id, context* streamSubCtx, const BYTE* buf, ULONG len);

    // COMMAND_FILE_COMPLETE_V2：SHA-256 校验（HandleFileCompleteV2）+ 完成判定。
    void OnFileCompleteV2(uint64_t device_id, const BYTE* buf, ULONG len);

    // 工具线程：等待传输完成；true=成功（files 已填，error=0），false=失败/超时（error 填原因）。
    // 无论成败，files 都会返回已见文件的落盘路径（成功=完整文件，失败=供调用方删半成品）。
    // 调用方随后必须调 ClearFileTransfer 收尾（擦会话 + CancelIO 子链接）。
    bool WaitFileTransferDone(uint64_t device_id, int timeoutMs,
                              std::vector<FileTransferEntry>& files, int& skipped, int& error);

    // 收尾：擦会话 + 锁外 CancelIO 文件管理器/流式两条子链接（幂等）。
    void ClearFileTransfer(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::string       sessionId;         // 持久会话 token（一次性 exec 为空）
        bool              persistent = false;
        bool              busy = false;      // 一条命令在飞（terminal_exec 复位→WaitTermCommand 返回）
        time_t            lastActiveAt = 0;  // idle 回收用（秒）
    };
    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（顶部路由）

    // ===== P5：远程控制会话（受 m_ScreenCtrlMutex 保护；单设备单会话）=====
    struct ScreenCtrlSession {
        std::string  sessionId;             // 会话 token（每次 open 独立随机）
        context*     subCtx = nullptr;      // 屏幕子连接上下文（就绪后填；用于注入）
        bool         started = false;       // false=子连接建立中；true=已就绪
        bool         busy    = false;       // 一条注入在飞（BeginScreenCtrlAction→End 之间）
        bool         closed  = false;       // 注入期间屏幕子连接已断（由注入线程收尾）
        time_t       lastActiveAt = 0;      // idle 回收用（秒）
        int          screenW = 0;           // 物理捕获分辨率（来自 TOKEN_BITMAPINFO）
        int          screenH = 0;
    };
    std::mutex                              m_ScreenCtrlMutex;
    std::map<uint64_t, ScreenCtrlSession>   m_ScreenCtrlSessions;         // device_id → 会话
    std::map<context*, uint64_t>            m_ScreenCtrlContextToDevice;  // subCtx → device_id（OfflineProc 反查）

    // ===== P6：文件传输会话（受 m_FileXferMutex 保护；单设备单传输）=====
    struct FileTransferSession {
        std::string tool;                    // 恒为 "download_file"
        std::string localDir;                // 本机保存目录（ANSI，结尾 '\'）
        std::string remotePath;              // 客户端远程路径（ANSI，发 CMD_DOWN_FILES_V2 用）
        bool        overwrite = false;       // 是否覆盖已存在文件
        context*    fmSubCtx     = nullptr;  // 文件管理器子链接（下发 CMD_DOWN_FILES_V2）
        context*    streamSubCtx = nullptr;  // 流式子链接（首个 chunk 填，收尾 CancelIO）
        uint64_t    transferID = 0;          // 客户端回填的传输会话 ID
        uint32_t    totalFiles = 0;          // 客户端声明的文件总数（含目录项，见 F4）
        uint32_t    directoryCount = 0;      // 目录项计数（目录不发 COMPLETE 包）
        uint32_t    filesDone = 0;           // 已完成的文件数（含跳过的）
        bool        done = false;            // 全部完成或出错
        int         error = 0;               // 0=ok；8=哈希不匹配；1001=路径逃逸；1002=超时
        std::map<uint32_t, FileTransferEntry> fileEntries;   // fileIndex → {path,size,sha256}
        std::set<uint32_t>        skipIndexes;               // overwrite=false 跳过的 fileIndex
        std::vector<FileTransferEntry> skipped;              // 跳过的文件（供输出）
        time_t      startAt = 0;
    };
    std::mutex                            m_FileXferMutex;
    std::condition_variable               m_FileXferCv;
    std::map<uint64_t, FileTransferSession> m_FileXferSessions;   // device_id → 会话

    bool        m_readonly = true;
    std::string m_cmdWhitelist;
    bool        m_terminalEnabled = false;      // 持久终端开关（默认关；要求 m_readonly=false）
    bool        m_remoteControlEnabled = false; // 远程控制开关（默认关；要求 m_readonly=false）
    bool        m_fileTransferEnabled = false;  // 文件传输开关（默认关；download_file）
};

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

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