fix: validate callback hosts before generating payloads

internal/c2/beacon_host.go

@@ -1,6 +1,8 @@
 package c2
 
 import (
+	"fmt"
+	"net"
 	"strings"
 
 	"cyberstrike-ai/internal/database"
@@ -37,3 +39,34 @@ func ResolveBeaconDialHost(listener *database.C2Listener, explicitOverride strin
 	}
 	return host
 }
+
+// ValidateBeaconDialHost accepts a host only, never a URL, port or command fragment.
+func ValidateBeaconDialHost(host string) error {
+	if host == "" || len(host) > 253 {
+		return fmt.Errorf("callback host must be an IP address or hostname without a scheme, path or port")
+	}
+	if net.ParseIP(host) != nil {
+		return nil
+	}
+	name := strings.TrimSuffix(host, ".")
+	numeric := true
+	for _, c := range name {
+		if c != '.' && (c < '0' || c > '9') {
+			numeric = false
+		}
+	}
+	if numeric {
+		return fmt.Errorf("callback host contains an invalid IP address")
+	}
+	for _, label := range strings.Split(name, ".") {
+		if len(label) == 0 || len(label) > 63 || label[0] == '-' || label[len(label)-1] == '-' {
+			return fmt.Errorf("callback hostname has an invalid label")
+		}
+		for _, c := range label {
+			if !(c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' || c == '-') {
+				return fmt.Errorf("callback host must be an IP address or hostname without a scheme, path or port")
+			}
+		}
+	}
+	return nil
+}

internal/c2/beacon_host_validation_test.go

@@ -0,0 +1,22 @@
+package c2
+
+import "testing"
+
+func TestCallbackHostValidation(t *testing.T) {
+	for _, host := range []string{"127.0.0.1", "::1", "2001:db8::1", "localhost", "example.com", "host.docker.internal", "example.com."} {
+		if err := ValidateBeaconDialHost(host); err != nil {
+			t.Errorf("valid %q: %v", host, err)
+		}
+	}
+	for _, host := range []string{"", "bad host/abc", "https://example.com", "example.com:443", "999.1.1.1", "-bad.example", "bad_.example", "a..example", "x';echo test", "$(echo test)", "example.\ncom"} {
+		if err := ValidateBeaconDialHost(host); err == nil {
+			t.Errorf("invalid %q accepted", host)
+		}
+		for _, kind := range AllOnelinerKinds() {
+			result, err := GenerateOneliner(OnelinerInput{Kind: kind, Host: host, Port: 8080, HTTPBaseURL: "http://example.com:8080", ImplantToken: "test-only-token"})
+			if err == nil || result != "" {
+				t.Errorf("%s generated output for invalid host %q", kind, host)
+			}
+		}
+	}
+}

internal/c2/payload_builder.go

@@ -90,6 +90,12 @@ func (b *PayloadBuilder) BuildBeacon(in PayloadBuilderInput) (*BuildResult, erro
 		goarch = "amd64"
 	}
 
+	// 模板参数：请求 Host > 监听器 callback_host > bind 推导（见 ResolveBeaconDialHost）
+	host := ResolveBeaconDialHost(listener, in.Host, b.logger, listener.ID)
+	if err := ValidateBeaconDialHost(host); err != nil {
+		return nil, err
+	}
+
 	// 读取模板
 	tmplPath := filepath.Join(b.tmplDir, "beacon.go.tmpl")
 	tmplData, err := os.ReadFile(tmplPath)
@@ -97,12 +103,9 @@ func (b *PayloadBuilder) BuildBeacon(in PayloadBuilderInput) (*BuildResult, erro
 		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",
+	serverURL := fmt.Sprintf("%s://%s",
 		listenerTypeToScheme(listener.Type),
-		host,
-		listener.BindPort,
+		net.JoinHostPort(host, strconv.Itoa(listener.BindPort)),
 	)
 
 	transport := "http"

internal/c2/payload_oneliner.go

@@ -106,8 +106,8 @@ func ValidateOnelinerForListener(listener *database.C2Listener, kind OnelinerKin
 //   - 同时返回执行示例，便于 AI 在对话里直接展示给操作员。
 func GenerateOneliner(in OnelinerInput) (string, error) {
 	host := strings.TrimSpace(in.Host)
-	if host == "" {
-		return "", fmt.Errorf("host is required")
+	if err := ValidateBeaconDialHost(host); err != nil {
+		return "", err
 	}
 	switch in.Kind {
 	case OnelinerBash:

internal/handler/c2.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -748,11 +749,15 @@ func (h *C2Handler) PayloadOneliner(c *gin.Context) {
 		return
 	}
 
+	scheme := "http"
+	if listener.Type == "https_beacon" {
+		scheme = "https"
+	}
 	input := c2.OnelinerInput{
 		Kind:         kind,
 		Host:         host,
 		Port:         listener.BindPort,
-		HTTPBaseURL:  fmt.Sprintf("http://%s:%d", host, listener.BindPort),
+		HTTPBaseURL:  fmt.Sprintf("%s://%s", scheme, net.JoinHostPort(host, strconv.Itoa(listener.BindPort))),
 		ImplantToken: listener.ImplantToken,
 	}
 

internal/handler/c2_listener_validation_test.go

@@ -11,6 +11,7 @@ import (
 
 	"cyberstrike-ai/internal/c2"
 	"cyberstrike-ai/internal/database"
+	"cyberstrike-ai/internal/security"
 	"github.com/gin-gonic/gin"
 	"go.uber.org/zap"
 )
@@ -53,3 +54,43 @@ func TestC2ListenerEditRejectsInvalidFieldsWithoutSaving(t *testing.T) {
 		}
 	}
 }
+
+func TestC2OnelinerRejectsInvalidCallbackAndUsesListenerScheme(t *testing.T) {
+	db, err := database.NewDB(filepath.Join(t.TempDir(), "callback.db"), zap.NewNop())
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = db.Close() })
+	listener := &database.C2Listener{ID: "https", Name: "test", Type: "https_beacon", BindHost: "127.0.0.1", BindPort: 8443, ImplantToken: "test-only-token", CreatedAt: time.Now()}
+	if err := db.CreateC2Listener(listener); err != nil {
+		t.Fatal(err)
+	}
+	h := NewC2Handler(c2.NewManager(db, zap.NewNop(), t.TempDir()), zap.NewNop())
+	r := gin.New()
+	r.Use(func(c *gin.Context) {
+		c.Set(security.ContextSessionKey, security.Session{UserID: "test-owner", Scope: database.RBACScopeAll})
+	})
+	r.POST("/oneliner", h.PayloadOneliner)
+	for _, host := range []string{"bad host/abc", "https://example.com", "example.com:8443", "example.com", "::1"} {
+		req := httptest.NewRequest(http.MethodPost, "/oneliner", bytes.NewBufferString(fmt.Sprintf(`{"listener_id":"https","kind":"curl_beacon","host":%q}`, host)))
+		req.Header.Set("Content-Type", "application/json")
+		w := httptest.NewRecorder()
+		r.ServeHTTP(w, req)
+		valid := host == "example.com" || host == "::1"
+		if !valid && w.Code != 400 {
+			t.Fatalf("invalid callback status: %d", w.Code)
+		}
+		if valid {
+			if w.Code != 200 {
+				t.Fatalf("valid callback status: %d", w.Code)
+			}
+			expected := "https://example.com:8443"
+			if host == "::1" {
+				expected = "https://[::1]:8443"
+			}
+			if !bytes.Contains(w.Body.Bytes(), []byte(expected)) {
+				t.Fatalf("scheme or IPv6 formatting lost")
+			}
+		}
+	}
+}