Add files via upload

internal/handler/webshell_probe.go

@@ -0,0 +1,127 @@
+package handler
+
+import (
+	"bytes"
+	"io"
+	"net/http"
+	"strings"
+
+	"go.uber.org/zap"
+)
+
+// webshellOSProbeCommand 探活命令：利用 Windows cmd 与 POSIX shell 对 `%OS%` 展开差异进行判定。
+//   - Windows cmd：`%OS%` 被展开为 `Windows_NT`，回显 `:OSPROBE_Windows_NT:END`
+//   - POSIX sh/bash：`%OS%` 不是变量语法，作为字面量原样保留，回显 `:OSPROBE_%OS%:END`
+//
+// 一条命令即可得到明确的、互斥的信号，避免探活成本（相比发两次命令）。
+// 冒号包裹是为了避免部分 shell 输出多余空白/BOM 时字符串匹配失效。
+const webshellOSProbeCommand = "echo :OSPROBE_%OS%:END"
+
+// probeWebshellOSViaExec 通过一次命令执行的回显推断目标操作系统。
+//
+// 返回值：
+//   - "windows" / "linux"：识别成功
+//   - ""：无法判定（调用方应保留既有 fallback 逻辑）
+//
+// 入参 execFn 是一个"发命令并拿到回显"的闭包；让 HTTP 入口和 MCP 入口可以共用同一套探活逻辑
+// 而不必关心底层是如何发包的。
+func probeWebshellOSViaExec(execFn func(cmd string) (output string, ok bool)) string {
+	if execFn == nil {
+		return ""
+	}
+	out, ok := execFn(webshellOSProbeCommand)
+	if !ok {
+		return ""
+	}
+	return classifyWebshellOSProbeOutput(out)
+}
+
+// classifyWebshellOSProbeOutput 纯函数：根据探活命令的回显判定 OS。
+// 抽出来是为了单测可直接覆盖所有分支，无需真实 HTTP 调用。
+func classifyWebshellOSProbeOutput(out string) string {
+	if out == "" {
+		return ""
+	}
+	lower := strings.ToLower(out)
+
+	// Windows 强信号：cmd.exe 成功展开了 %OS% 变量
+	if strings.Contains(out, "Windows_NT") {
+		return "windows"
+	}
+	// 容错：部分老版本 Windows 可能 `%OS%` 展开为其他字样（极少见），再看 PATH/OS 等次级线索
+	if strings.Contains(lower, "microsoft windows") {
+		return "windows"
+	}
+
+	// Linux/Unix 强信号：`%OS%` 字面量被原样回显，说明 shell 不是 cmd.exe
+	if strings.Contains(out, "%OS%") {
+		return "linux"
+	}
+
+	// 次级线索：部分 webshell 在 Linux 上可能走了其他外壳（如 zsh/ash），
+	// 但它们对 `%OS%` 同样不展开；若命中 OSPROBE 头部却没拿到 %OS% 字面量，
+	// 说明回显被中途截断或过滤，保守返回空让上层 fallback。
+	return ""
+}
+
+// newHTTPExecFn 为 HTTP FileOp 路径构造"发命令取回显"的闭包，供探活复用。
+// 参数来自 HTTP 请求，复用 buildExecURL / buildExecBody 两个已有的命令编排器，
+// 确保探活包与实际文件操作包走完全一致的 webshell 协议（GET/POST、参数名、编码）。
+func (h *WebShellHandler) newHTTPExecFn(targetURL, password, shellType, method, cmdParam, encoding string) func(string) (string, bool) {
+	useGET := strings.ToUpper(strings.TrimSpace(method)) == "GET"
+	if strings.TrimSpace(cmdParam) == "" {
+		cmdParam = "cmd"
+	}
+	return func(cmd string) (string, bool) {
+		var (
+			httpReq *http.Request
+			err     error
+		)
+		if useGET {
+			u := h.buildExecURL(targetURL, shellType, password, cmdParam, cmd)
+			httpReq, err = http.NewRequest(http.MethodGet, u, nil)
+		} else {
+			body := h.buildExecBody(shellType, password, cmdParam, cmd)
+			httpReq, err = http.NewRequest(http.MethodPost, targetURL, bytes.NewReader(body))
+			if err == nil {
+				httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+			}
+		}
+		if err != nil {
+			return "", false
+		}
+		httpReq.Header.Set("User-Agent", "Mozilla/5.0 (compatible; CyberStrikeAI-WebShell/1.0)")
+		resp, err := h.client.Do(httpReq)
+		if err != nil {
+			return "", false
+		}
+		defer resp.Body.Close()
+		raw, _ := io.ReadAll(resp.Body)
+		return decodeWebshellOutput(raw, encoding), resp.StatusCode == http.StatusOK
+	}
+}
+
+// persistDetectedOS 把探活结果回写到连接表；失败只记日志不阻断主流程。
+// 设计上故意只触发 UPDATE，不会新建记录，因此即便 connectionID 不存在也只是悄悄放弃。
+func (h *WebShellHandler) persistDetectedOS(connectionID, detected string) {
+	connectionID = strings.TrimSpace(connectionID)
+	detected = normalizeWebshellOS(detected)
+	if connectionID == "" || detected == "" || detected == "auto" {
+		return
+	}
+	conn, err := h.db.GetWebshellConnection(connectionID)
+	if err != nil || conn == nil {
+		// 不是所有调用方都能提供有效 ID（比如临时测试），这里静默返回
+		return
+	}
+	if normalizeWebshellOS(conn.OS) != "auto" {
+		// 用户已经显式选过 OS，尊重用户选择，不自动覆盖
+		return
+	}
+	conn.OS = detected
+	if err := h.db.UpdateWebshellConnection(conn); err != nil {
+		h.logger.Warn("webshell 探活结果持久化失败", zap.String("id", connectionID), zap.String("os", detected), zap.Error(err))
+		return
+	}
+	h.logger.Info("webshell auto OS 探活成功并持久化", zap.String("id", connectionID), zap.String("os", detected))
+}