Harden scan robustness and tests

2026-09-25 20:51:52 +08:00 · 2026-06-13 07:55:37 +08:00
parent 1595c92aed
commit 15a7670ba2
100 changed files with 4483 additions and 412 deletions
@@ -1,6 +1,11 @@
 package output

-import "testing"
+import (
+	"bufio"
+	"bytes"
+	"encoding/json"
+	"testing"
+)

 func TestSplitHostPort(t *testing.T) {
 	tests := []struct {
@@ -33,3 +38,85 @@ func TestSplitHostPort(t *testing.T) {
 		})
 	}
 }
+
+func TestNewStdoutNDJSONWriter(t *testing.T) {
+	writer := NewStdoutNDJSONWriter()
+	if writer == nil || writer.writer == nil {
+		t.Fatalf("NewStdoutNDJSONWriter = %#v, want initialized writer", writer)
+	}
+	if err := writer.Close(); err != nil {
+		t.Fatalf("Close error = %v", err)
+	}
+}
+
+func TestStdoutNDJSONWriterWriteResult(t *testing.T) {
+	var buf bytes.Buffer
+	writer := &StdoutNDJSONWriter{writer: bufio.NewWriter(&buf)}
+
+	result := &ScanResult{
+		Type:   TypeService,
+		Target: "[2001:db8::1]:8443",
+		Status: "OPEN",
+		Details: map[string]interface{}{
+			"port":          float64(9443),
+			"service":       "https",
+			"protocol":      "tcp",
+			"banner":        123,
+			"title":         "admin",
+			"url":           "https://[2001:db8::1]:8443",
+			"vulnerability": "weak credential",
+			"username":      "admin",
+			"password":      "secret",
+			"plugin":        "webtitle",
+			"version":       "1.2.3",
+			"os":            "linux",
+		},
+	}
+
+	if err := writer.WriteResult(result); err != nil {
+		t.Fatalf("WriteResult error = %v", err)
+	}
+	if err := writer.Close(); err != nil {
+		t.Fatalf("Close error = %v", err)
+	}
+
+	var rec ndjsonRecord
+	if err := json.Unmarshal(bytes.TrimSpace(buf.Bytes()), &rec); err != nil {
+		t.Fatalf("invalid ndjson output %q: %v", buf.String(), err)
+	}
+	if rec.Host != "2001:db8::1" || rec.Port != 9443 {
+		t.Fatalf("host/port = %q/%d", rec.Host, rec.Port)
+	}
+	if rec.Service != "https" || rec.Protocol != "tcp" || rec.Banner != "123" || rec.Title != "admin" {
+		t.Fatalf("flattened fields missing: %#v", rec)
+	}
+	if rec.URL != "https://[2001:db8::1]:8443" || rec.Vulnerability != "weak credential" {
+		t.Fatalf("url/vuln fields missing: %#v", rec)
+	}
+	if rec.Username != "admin" || rec.Password != "secret" || rec.Plugin != "webtitle" || rec.Version != "1.2.3" || rec.OS != "linux" {
+		t.Fatalf("credential/plugin fields missing: %#v", rec)
+	}
+}
+
+func TestStdoutNDJSONFlattenFallbacks(t *testing.T) {
+	writer := &StdoutNDJSONWriter{writer: bufio.NewWriter(&bytes.Buffer{})}
+
+	rec := writer.flatten(&ScanResult{
+		Type:   TypeHost,
+		Target: "2001:db8::1",
+		Status: "ALIVE",
+		Details: map[string]interface{}{
+			"port": int64(22),
+		},
+	})
+	if rec.Host != "2001:db8::1" || rec.Port != 22 {
+		t.Fatalf("flatten fallback = %#v", rec)
+	}
+
+	if got, ok := toInt("22"); ok || got != 0 {
+		t.Fatalf("toInt string = %d/%v, want 0/false", got, ok)
+	}
+	if got := strVal(map[string]interface{}{}, "missing"); got != "" {
+		t.Fatalf("missing strVal = %q, want empty", got)
+	}
+}
@@ -38,6 +38,19 @@ func escapeControlChars(s string) string {
 	return b.String()
 }

+func truncateString(s string, maxRunes int) string {
+	if maxRunes < 0 {
+		return s
+	}
+	for i := range s {
+		if maxRunes == 0 {
+			return s[:i] + "..."
+		}
+		maxRunes--
+	}
+	return s
+}
+
 func targetWithPort(target string, port interface{}) string {
 	if port == nil {
 		return target
@@ -196,10 +209,8 @@ func (w *TXTWriter) formatServiceLine(result *ScanResult) string {
 		parts = append(parts, service)
 	}
 	if banner != "" {
-		if len(banner) > 100 {
-			banner = banner[:100] + "..."
-		}
 		banner = escapeControlChars(banner)
+		banner = truncateString(banner, 100)
 		parts = append(parts, banner)
 	}
 	return strings.Join(parts, " ")
@@ -745,9 +756,7 @@ func (w *CSVWriter) formatServiceRecord(result *ScanResult) []string {
 		fingerprints = formatFingerprints(result.Details["fingerprints"])
 		if b, ok := result.Details["banner"].(string); ok {
 			banner = escapeControlChars(b)
-			if len(banner) > 100 {
-				banner = banner[:100] + "..."
-			}
+			banner = truncateString(banner, 100)
 		}
 	}
 	target := result.Target
@@ -9,6 +9,7 @@ import (
 	"sync"
 	"testing"
 	"time"
+	"unicode/utf8"
 )

 /*
@@ -82,6 +83,94 @@ func TestTargetWithPortIPv6(t *testing.T) {
 	}
 }

+func TestScanResultFormatDetailsAndDefaultManagerConfig(t *testing.T) {
+	result := &ScanResult{
+		Details: map[string]interface{}{
+			"service": "ssh",
+			"port":    22,
+			"banner":  "OpenSSH",
+		},
+	}
+	got := result.FormatDetails(";", "%s=%v")
+	want := "banner=OpenSSH;port=22;service=ssh"
+	if got != want {
+		t.Fatalf("FormatDetails = %q, want %q", got, want)
+	}
+
+	empty := (&ScanResult{}).FormatDetails(";", "%s=%v")
+	if empty != "" {
+		t.Fatalf("empty FormatDetails = %q, want empty", empty)
+	}
+
+	cfg := DefaultManagerConfig("out.json", FormatJSON)
+	if cfg.OutputPath != "out.json" || cfg.Format != FormatJSON {
+		t.Fatalf("DefaultManagerConfig = %#v", cfg)
+	}
+}
+
+func TestCSVWriterFormatRecords(t *testing.T) {
+	writer := &CSVWriter{}
+
+	host := writer.formatHostRecord(&ScanResult{Target: "192.168.1.1"})
+	if len(host) != 1 || host[0] != "192.168.1.1" {
+		t.Fatalf("host record = %#v", host)
+	}
+
+	port := writer.formatPortRecord(&ScanResult{
+		Target:  "192.168.1.1",
+		Details: map[string]interface{}{"port": 22},
+	})
+	if got, want := strings.Join(port, "|"), "192.168.1.1|22|open"; got != want {
+		t.Fatalf("port record = %q, want %q", got, want)
+	}
+
+	longBanner := strings.Repeat("界", 105)
+	service := writer.formatServiceRecord(&ScanResult{
+		Target: "2001:db8::1",
+		Details: map[string]interface{}{
+			"port":         443,
+			"name":         "https",
+			"version":      "1.2.3",
+			"title":        "hello\nworld",
+			"status":       200,
+			"server":       "nginx\r\nunit",
+			"fingerprints": []interface{}{"fp1", "", "fp2", 3},
+			"banner":       longBanner,
+		},
+	})
+	if service[0] != "[2001:db8::1]:443" || service[1] != "https" || service[2] != "1.2.3" {
+		t.Fatalf("service identity fields = %#v", service)
+	}
+	if service[3] != "hello\\nworld" || service[4] != "200" || service[5] != "nginx\\r\\nunit" {
+		t.Fatalf("service text fields = %#v", service)
+	}
+	if service[6] != "fp1,fp2" {
+		t.Fatalf("fingerprints = %q, want fp1,fp2", service[6])
+	}
+	if !utf8.ValidString(service[7]) || len([]rune(service[7])) != 103 || !strings.HasSuffix(service[7], "...") {
+		t.Fatalf("truncated banner = len %d value %q", len(service[7]), service[7])
+	}
+
+	vuln := writer.formatVulnRecord(&ScanResult{
+		Target:  "http://example.com",
+		Status:  "vulnerable",
+		Details: map[string]interface{}{"type": "poc"},
+	})
+	if got, want := strings.Join(vuln, "|"), "http://example.com|poc|vulnerable"; got != want {
+		t.Fatalf("vuln record = %q, want %q", got, want)
+	}
+
+	if got := formatFingerprints([]string{"a", "b"}); got != "a,b" {
+		t.Fatalf("string fingerprints = %q", got)
+	}
+	if got := formatFingerprints(123); got != "" {
+		t.Fatalf("unsupported fingerprints = %q, want empty", got)
+	}
+	if writer.GetFormat() != FormatCSV {
+		t.Fatalf("csv GetFormat = %q", writer.GetFormat())
+	}
+}
+
 // =============================================================================
 // TXTWriter - 基础功能测试
 // =============================================================================