From ed47ee202e41aaa790ba22806783aa20ddcb477c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=85=AC=E6=98=8E?=
 <83812544+Ed1s0nZ@users.noreply.github.com>
Date: Thu, 6 Aug 2026 17:14:06 +0800
Subject: [PATCH] Add files via upload

---
 internal/c2/listener_http.go      | 24 ++++-----
 internal/c2/listener_http_test.go | 85 +++++++++++++++++++++++++++++++
 internal/c2/manager.go            | 67 ++++++++++++++++++++++--
 internal/c2/manager_blob_test.go  | 75 +++++++++++++++++++++++++++
 internal/c2/tcp_beacon_server.go  |  6 ++-
 5 files changed, 238 insertions(+), 19 deletions(-)
 create mode 100644 internal/c2/manager_blob_test.go

diff --git a/internal/c2/listener_http.go b/internal/c2/listener_http.go
index 22fef328..a4fc62e4 100644
--- a/internal/c2/listener_http.go
+++ b/internal/c2/listener_http.go
@@ -291,16 +291,13 @@ func (l *HTTPBeaconListener) handleResult(w http.ResponseWriter, r *http.Request
 	}
 	var report TaskResultReport
 	plaintext, decErr := DecryptAESGCM(l.rec.EncryptionKey, string(body))
-	if decErr == nil {
-		if err := json.Unmarshal(plaintext, &report); err != nil {
-			l.disguisedReject(w)
-			return
-		}
-	} else {
-		if err := json.Unmarshal(body, &report); err != nil {
-			l.disguisedReject(w)
-			return
-		}
+	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)
@@ -341,12 +338,15 @@ func (l *HTTPBeaconListener) handleUpload(w http.ResponseWriter, r *http.Request
 		l.disguisedReject(w)
 		return
 	}
-	dir := filepath.Join(l.manager.StorageDir(), "uploads")
+	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
 	}
-	dst := filepath.Join(dir, taskID+".bin")
 	if err := os.WriteFile(dst, plaintext, 0o644); err != nil {
 		http.Error(w, "save failed", http.StatusInternalServerError)
 		return
diff --git a/internal/c2/listener_http_test.go b/internal/c2/listener_http_test.go
index 8db0e34f..58f1d9d1 100644
--- a/internal/c2/listener_http_test.go
+++ b/internal/c2/listener_http_test.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"net/http/httptest"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -227,3 +228,87 @@ func TestHTTPBeaconListener_HandleFileServe(t *testing.T) {
 		})
 	}
 }
+
+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/manager.go b/internal/c2/manager.go
index d620c87b..6b0c2adf 100644
--- a/internal/c2/manager.go
+++ b/internal/c2/manager.go
@@ -42,6 +42,11 @@ type Manager struct {
 // 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 {
@@ -736,18 +741,24 @@ func (m *Manager) IngestTaskResult(report TaskResultReport) error {
 }
 
 func (m *Manager) saveResultBlob(taskID, b64Content, suffix string) (string, error) {
-	suffix = strings.TrimSpace(suffix)
-	if suffix == "" {
-		suffix = ".bin"
+	taskID = strings.TrimSpace(taskID)
+	if taskID == "" || taskID == "." || taskID == ".." ||
+		strings.ContainsAny(taskID, `/\`) {
+		return "", fmt.Errorf("invalid task_id")
 	}
-	if !strings.HasPrefix(suffix, ".") {
-		suffix = "." + suffix
+
+	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
@@ -758,6 +769,52 @@ func (m *Manager) saveResultBlob(taskID, b64Content, suffix string) (string, 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
+}
+
 // ----------------------------------------------------------------------------
 // 事件总线辅助
 // ----------------------------------------------------------------------------
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/tcp_beacon_server.go b/internal/c2/tcp_beacon_server.go
index 4eeda05e..9c044196 100644
--- a/internal/c2/tcp_beacon_server.go
+++ b/internal/c2/tcp_beacon_server.go
@@ -209,11 +209,13 @@ func (l *TCPReverseListener) handleTCPBeaconSession(conn net.Conn, br *bufio.Rea
 			if err != nil {
 				return
 			}
-			dir := filepath.Join(l.manager.StorageDir(), "uploads")
+			dir, dst, err := uploadPathForTask(l.manager.StorageDir(), up.TaskID)
+			if err != nil {
+				return
+			}
 			if err := os.MkdirAll(dir, 0o755); err != nil {
 				return
 			}
-			dst := filepath.Join(dir, up.TaskID+".bin")
 			if err := os.WriteFile(dst, plainFile, 0o644); err != nil {
 				return
 			}