SimpleRemoter/server/2015Remote/WebService.h

#pragma once

#include <Windows.h>
#include <string>
#include <vector>
#include <map>
#include <set>
#include <atomic>
#include <chrono>
#include <mutex>
#include <thread>
#include <memory>
#include <random>
#include <ctime>

#include "common/commands.h"   // MSG64 / COMMAND_SCREEN_CONTROL：共享键盘注入 helper（BuildKeyMsg64）

// Forward declarations
class context;
class CMy2015RemoteDlg;
class CONTEXT_OBJECT;

// 构造单个键盘注入 MSG64（复用 WebService::HandleKey 的 lParam 逻辑，抽取为共享 helper，
// 供 Web 与 MCP 远程控制两处使用，避免漂移，见 docs/Mcp_RemoteControl_Design.md §6.2）。
// keyCode：虚拟键码；isDown：按下/抬起；isSysKey：true 用 WM_SYSKEYDOWN/UP 并置 Alt 上下文位。
inline MSG64 BuildKeyMsg64(int keyCode, bool isDown, bool isSysKey) {
    MSG64 msg64;
    memset(&msg64, 0, sizeof(MSG64));

    if (isSysKey) msg64.message = isDown ? WM_SYSKEYDOWN : WM_SYSKEYUP;
    else          msg64.message = isDown ? WM_KEYDOWN : WM_KEYUP;
    msg64.wParam = (uint64_t)keyCode;
    msg64.time   = GetTickCount();

    UINT scanCode = MapVirtualKey(keyCode, MAPVK_VK_TO_VSC);
    bool isExtended = (keyCode >= VK_PRIOR && keyCode <= VK_DOWN) ||
                      keyCode == VK_INSERT || keyCode == VK_DELETE ||
                      keyCode == VK_NUMLOCK || keyCode == VK_RCONTROL || keyCode == VK_RMENU ||
                      keyCode == VK_APPS;

    LPARAM lParam = 1;                       // repeat count = 1
    lParam |= (scanCode & 0xFF) << 16;
    if (isExtended) lParam |= (1 << 24);
    if (isSysKey)   lParam |= (1 << 29);     // context code for Alt
    if (!isDown)    lParam |= (3UL << 30);   // bit 30/31 for key up

    msg64.lParam = (uint64_t)lParam;
    return msg64;
}

// 构造单个鼠标注入 MSG64（复用 WebService::HandleMouse 的基座字段逻辑，抽取为共享 helper，
// 供 Web 与 MCP 远程控制两处使用，见 docs/Mcp_RemoteControl_Design.md §6.2）。
// x/y 为物理像素坐标（已由调用方从归一化坐标映射）；message/wParam 由调用方按动作设定。
inline MSG64 BuildMouseMsg64(int x, int y, UINT message, uint64_t wParam) {
    MSG64 msg64;
    memset(&msg64, 0, sizeof(MSG64));
    msg64.pt.x    = x;
    msg64.pt.y    = y;
    msg64.lParam  = MAKELPARAM(x, y);
    msg64.time    = GetTickCount();
    msg64.message = message;
    msg64.wParam  = wParam;
    return msg64;
}

// Web client state
struct WebClient {
    std::string token;
    std::string username;
    std::string role;           // "admin" | "viewer"
    uint64_t watch_device_id;   // 0 = not watching any device
    uint64_t connected_at;
    uint64_t last_activity;     // Last message/pong received (for heartbeat timeout)
    std::string client_ip;

    WebClient() : watch_device_id(0), connected_at(0), last_activity(0) {}
};

// Login attempt tracking for rate limiting
struct LoginAttempt {
    int failed_count;
    time_t locked_until;

    LoginAttempt() : failed_count(0), locked_until(0) {}
};

// Web user account
struct WebUser {
    std::string username;
    std::string password_hash;  // SHA256(password + salt)
    std::string salt;
    std::string role;           // "admin" | "viewer"
    std::vector<std::string> allowed_groups;  // Groups this user can view (empty = no access, admin = all)
};

// Device info for web clients
struct WebDeviceInfo {
    uint64_t id;
    std::string name;
    std::string ip;
    std::string os;
    int screen_width;
    int screen_height;
    bool online;
    bool top_down = false;          // true = Android/top-down H.264, no flip needed in browser
    std::string client_type;        // RES_CLIENT_TYPE: "APK"=Android, "EXE"=Windows, "LNX"=Linux, "MAC"=macOS
    // 当前会话的音频开关。-1=未知（客户端 BITMAPINFO 还没回来），0=关，1=开
    int audio_enabled = -1;

    // Keyframe cache for new web clients
    std::vector<uint8_t> keyframe_cache;
    std::mutex cache_mutex;

    WebDeviceInfo() : id(0), screen_width(0), screen_height(0), online(false) {}
};

// Main Web Service class
class CWebService {
public:
    static CWebService& Instance();

    // Lifecycle
    bool Start(int port = 8080);
    void Stop();
    bool IsRunning() const;

    // Set parent dialog for device list access
    void SetParentDlg(CMy2015RemoteDlg* pDlg) { m_pParentDlg = pDlg; }

    // Set admin password (use master password)
    void SetAdminPassword(const std::string& password);

    // User management
    bool CreateUser(const std::string& username, const std::string& password, const std::string& role,
                    const std::vector<std::string>& allowed_groups = {});
    bool DeleteUser(const std::string& username);
    std::vector<std::pair<std::string, std::string>> ListUsers();  // Returns [(username, role), ...]

    // Device management (called from main app)
    void MarkDeviceOnline(uint64_t device_id);
    void MarkDeviceOffline(uint64_t device_id);
    void FlushDeviceChanges();  // Called by timer to batch-notify web clients

    // H264 frame broadcasting (called from ScreenSpyDlg)
    void BroadcastFrame(uint64_t device_id, const uint8_t* data, size_t len, bool is_keyframe);
    void BroadcastH264Frame(uint64_t device_id, const uint8_t* data, size_t len);
    void CacheKeyframe(uint64_t device_id, const uint8_t* data, size_t len);

    // Resolution change notification
    void NotifyResolutionChange(uint64_t device_id, int width, int height, bool top_down = false, const std::string& client_type = "");

    // 读回 NotifyResolutionChange 缓存的物理捕获分辨率（来自 TOKEN_BITMAPINFO）。
    // 尚未收到分辨率时返回 false 并置 w/h 为 0。MCP 远程控制（remote_open）用它映射归一化坐标。
    bool GetScreenSize(uint64_t device_id, int& width, int& height);

    // Audio enable/disable notification — pushes current state to all web
    // clients watching this device and caches it for newcomers.
    void NotifyAudioState(uint64_t device_id, bool enabled);

    // Cursor change notification (called from ScreenSpyDlg)
    void BroadcastCursor(uint64_t device_id, uint8_t cursor_index);

    // Get count of web clients watching a device
    int GetWebClientCount(uint64_t device_id);

    // Start remote desktop session for web viewing
    bool StartRemoteDesktop(uint64_t device_id);

    // Stop remote desktop session
    void StopRemoteDesktop(uint64_t device_id);

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

    // Server thread
    void ServerThread(int port);

    // Signaling handlers
    void HandleLogin(void* ws_ptr, const std::string& msg, const std::string& client_ip);
    void HandleGetSalt(void* ws_ptr, const std::string& msg);
    void HandleGetDevices(void* ws_ptr, const std::string& token);
    void HandleConnect(void* ws_ptr, const std::string& token, uint64_t device_id);
    void HandleDisconnect(void* ws_ptr, const std::string& token, uint64_t requested_device_id = 0);
    void HandlePing(void* ws_ptr, const std::string& token);
    void HandleMouse(void* ws_ptr, const std::string& msg);
    void HandleKey(void* ws_ptr, const std::string& msg);
    void HandleRdpReset(void* ws_ptr, const std::string& token);
    void HandleAudioToggle(void* ws_ptr, const std::string& token);

    // Token management
    std::string GenerateToken(const std::string& username, const std::string& role);
    bool ValidateToken(const std::string& token, std::string& username, std::string& role);

    // Client management
    void RegisterClient(void* ws_ptr, const std::string& client_ip);
    void UnregisterClient(void* ws_ptr);
    WebClient* FindClient(void* ws_ptr);

    // Rate limiting
    bool CheckRateLimit(const std::string& ip);
    void RecordFailedLogin(const std::string& ip);
    void RecordSuccessLogin(const std::string& ip);

    // JSON helpers
    std::string BuildJsonResponse(const std::string& cmd, bool ok, const std::string& msg = "");
    std::string BuildDeviceListJson(const std::string& username = "");

    // Password verification
    bool VerifyPassword(const std::string& input, const WebUser& user);
    std::string ComputeHash(const std::string& input);

    // User management helpers
    std::string GetUsersFilePath();
    void LoadUsers();
    void SaveUsers();
    void HandleCreateUser(void* ws_ptr, const std::string& msg);
    void HandleDeleteUser(void* ws_ptr, const std::string& msg);
    void HandleListUsers(void* ws_ptr, const std::string& token);
    void HandleGetGroups(void* ws_ptr, const std::string& token);

    // Send to WebSocket
    void SendText(void* ws_ptr, const std::string& text);
    void SendBinary(void* ws_ptr, const uint8_t* data, size_t len);

    // Build binary frame packet
    std::vector<uint8_t> BuildFramePacket(uint64_t device_id, bool is_keyframe,
                                          const uint8_t* data, size_t len);

private:
    // Server state
    std::thread m_ServerThread;
    std::thread m_HeartbeatThread;  // Heartbeat checker thread
    std::atomic<bool> m_bRunning;
    std::atomic<bool> m_bStopping;
    void* m_pServer;  // ws::Server*

    // Heartbeat settings
    static const int HEARTBEAT_INTERVAL_SEC = 30;   // Send ping every 30 seconds
    static const int HEARTBEAT_TIMEOUT_SEC = 90;    // Disconnect if no activity for 90 seconds

    // Heartbeat thread function
    void HeartbeatThread();

    // Parent dialog for device access
    CMy2015RemoteDlg* m_pParentDlg;

    // Web clients: ws_ptr -> WebClient
    std::map<void*, WebClient> m_Clients;
    std::mutex m_ClientsMutex;

    // Device keyframe cache: device_id -> cache
    std::map<uint64_t, std::shared_ptr<WebDeviceInfo>> m_DeviceCache;
    std::mutex m_DeviceCacheMutex;

    // Login rate limiting: ip -> LoginAttempt
    std::map<std::string, LoginAttempt> m_LoginAttempts;
    std::mutex m_LoginMutex;

    // User accounts (loaded from config)
    std::vector<WebUser> m_Users;
    std::mutex m_UsersMutex;

    // Token secret key (generated on startup)
    std::string m_SecretKey;

    // Config
    int m_nMaxClientsPerDevice;
    int m_nTokenExpireSeconds;
    bool m_bHideWebSessions;  // Whether to hide web-triggered dialogs (default: true)
    std::string m_PayloadsDir;  // Directory for file downloads (Payloads/)
    std::string m_ConfigDir;    // Directory for config files (users.json, etc.)

    // Web-triggered sessions (should be hidden)
    std::set<uint64_t> m_WebTriggeredDevices;
    std::mutex m_WebTriggeredMutex;

    // Dirty device tracking for batch notifications
    std::set<uint64_t> m_OnlineDevices;   // Devices that came online since last flush
    std::set<uint64_t> m_OfflineDevices;  // Devices that went offline since last flush
    std::mutex m_DirtyDevicesMutex;

    // Grace-period timers: delay CloseWebRemoteDesktopByClientID by 30 s after
    // the last web viewer disconnects to allow quick mobile reconnects (e.g.
    // switching to a TOTP app briefly). Cancelled by StartRemoteDesktop when a
    // viewer reconnects within the window.
    std::map<uint64_t, std::shared_ptr<std::atomic<bool>>> m_PendingCloseTimers;
    std::mutex m_PendingCloseMutex;

public:
    // Check if a device session was triggered by web (should be hidden)
    bool IsWebTriggered(uint64_t device_id);
    void ClearWebTriggered(uint64_t device_id);

    // MFC trigger management - MFC dialogs should always be visible
    void SetMfcTriggered(uint64_t device_id);
    bool IsMfcTriggered(uint64_t device_id);
    void ClearMfcTriggered(uint64_t device_id);

    // Check if a remote desktop session already exists for device
    bool HasActiveSession(uint64_t device_id);

    // Config accessors
    void SetHideWebSessions(bool hide) { m_bHideWebSessions = hide; }
    bool GetHideWebSessions() const { return m_bHideWebSessions; }

    // Real-time device updates
    void NotifyDeviceUpdate(uint64_t device_id, const std::string& rtt, const std::string& activeWindow);

    // Screen context registry (for mouse/keyboard control)
    void RegisterScreenContext(uint64_t device_id, CONTEXT_OBJECT* ctx);
    void UnregisterScreenContext(uint64_t device_id);
    CONTEXT_OBJECT* GetScreenContext(uint64_t device_id);

    // ========== Web Terminal (Phase 1: 1 user per device) ==========
    // Web 终端会话桥：把浏览器端 xterm.js ↔ 客户端 shell 子上下文连起来。
    //   设计：每台主机最多一个 Web 终端会话；如果别的浏览器请求同一台主机的终端，
    //         拒绝（UX 上后续可改成共享只读）。
    //   生命周期：term_open → COMMAND_SHELL → 客户端建子上下文 → MessageHandle
    //         看到 TOKEN_TERMINAL_START / TOKEN_SHELL_START + IsTermPending(d) →
    //         调 RegisterTerminalContext 接管，跳过 MFC dialog 打开。

    // 浏览器侧入口
    void HandleTermOpen(void* ws_ptr, const std::string& msg);
    void HandleTermInput(void* ws_ptr, const std::string& msg);
    void HandleTermResize(void* ws_ptr, const std::string& msg);
    void HandleTermClose(void* ws_ptr, const std::string& msg);

    // MessageHandle 向 WebService 询问 / 移交的钩子
    bool IsTermPending(uint64_t device_id);                              // 决定是否要拦截 dialog 打开
    void RegisterTerminalContext(uint64_t device_id, CONTEXT_OBJECT* ctx, bool isPty);
    bool IsTerminalContext(CONTEXT_OBJECT* ctx);                          // 是否是 Web 终端持有的上下文
    void OnTerminalData(CONTEXT_OBJECT* ctx, const BYTE* data, ULONG len);// 把 shell 输出泵到对应 web client
    void OnTerminalClosed(CONTEXT_OBJECT* ctx);                          // shell 子上下文断开时清理

private:
    // Screen context registry: device_id -> ScreenManager's CONTEXT_OBJECT
    std::map<uint64_t, CONTEXT_OBJECT*> m_ScreenContexts;
    std::mutex m_ScreenContextsMutex;

    // MFC triggered devices: dialogs created by MFC should always be visible
    std::set<uint64_t> m_MfcTriggeredDevices;
    std::mutex m_MfcTriggeredMutex;

    // Web 终端会话状态
    struct WebTermSession {
        void* ws_ptr;                  // browser WebSocket
        uint64_t device_id;
        CONTEXT_OBJECT* shell_ctx;     // shell 子上下文（首条消息抵达后才填）
        bool is_pty;                   // true=TOKEN_TERMINAL（现代 PTY）, false=TOKEN_SHELL（老 cmd 管道）
    };
    std::map<uint64_t, WebTermSession>      m_TermSessions;        // by device_id
    std::map<CONTEXT_OBJECT*, uint64_t>     m_TermContextToDevice; // 反查 ctx → device_id
    std::set<uint64_t>                      m_TermPending;         // 已发 COMMAND_SHELL 待响应
    std::mutex                              m_TermMutex;

    // 内部清理（已持锁版本）
    void CloseTermSessionLocked(uint64_t device_id);
};

// Global accessor
inline CWebService& WebService() { return CWebService::Instance(); }