feat: enhance extract system with severity, context capture, matchers, and new rules

- parsers.Extracted: add Severity and Items (ExtractItem with ctx/line/offset) - ProtonExtract: preserve severity from templates, capture N-char context around matches - Add --extract-context CLI flag (default 0, e.g. --extract-context 50) - Add word matchers pre-filtering to 14 existing templates for performance - Add 9 new HaE-inspired rules: lfi-indicator, upload-form, url-as-value, location-redirect, mac-address, sensitive-field, linkfinder, url-schemes, authorization-header - Total proton templates: 45 (was 35) - Update ColorString/TreeString to render severity and context - Comprehensive tests for severity, context, new templates, tag mappings Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-10-28 19:40:49 +00:00 · 2026-06-01 01:44:48 -07:00
1 parent 5d3ddeae56
commit 729dee3d06
@@ -1,6 +1,7 @@
 package pkg

 import (
+	"bytes"
 	"os"
 	"path/filepath"
 	"strings"
@@ -10,16 +11,17 @@ import (
 	"github.com/chainreactors/neutron/protocols"
 	"github.com/chainreactors/parsers"
 	"github.com/chainreactors/proton/protocols/file"
-	protonTmpl "github.com/chainreactors/proton/templates"
+	protonTmpl "github.com/chainreactors/proton/template"
 	yaml "sigs.k8s.io/yaml/goyaml.v3"
 )

 var (
-	ProtonScanner     *file.Scanner
-	protonTemplates   []*protonTmpl.Template
-	protonTemplateMap map[string]*protonTmpl.Template // id → template
-	protonTagMap      map[string][]string             // tag → []id
-	protonMu          sync.RWMutex
+	ProtonScanner      *file.Scanner
+	protonTemplates    []*protonTmpl.Template
+	protonTemplateMap  map[string]*protonTmpl.Template // id → template
+	protonTagMap       map[string][]string             // tag → []id
+	protonMu           sync.RWMutex
+	ExtractContextSize int // 0=disabled, >0=chars of context each side
 )

 // LoadProtonTemplates parses YAML template docs and builds a Scanner.
@@ -50,6 +52,7 @@ func LoadProtonTemplates(yamlDocs [][]byte) error {
 			tag = strings.TrimSpace(tag)
 			protonTagMap[tag] = append(protonTagMap[tag], tmpl.Id)
 		}
+
 	}

 	rebuildScanner()
@@ -146,6 +149,35 @@ func rebuildScanner() {
 	ProtonScanner = file.NewScanner(rules, nil)
 }

+type extractHit struct {
+	value    string
+	line     int
+	offset   int
+	severity string
+}
+
+func captureContext(content []byte, lineOffset int, value string, ctxSize int) string {
+	searchEnd := len(content)
+	if nlPos := bytes.IndexByte(content[lineOffset:], '\n'); nlPos >= 0 {
+		searchEnd = lineOffset + nlPos
+	}
+	line := content[lineOffset:searchEnd]
+	idx := bytes.Index(line, []byte(value))
+	if idx == -1 {
+		return ""
+	}
+	matchStart := lineOffset + idx
+	start := matchStart - ctxSize
+	if start < 0 {
+		start = 0
+	}
+	end := matchStart + len(value) + ctxSize
+	if end > len(content) {
+		end = len(content)
+	}
+	return string(content[start:end])
+}
+
 // ProtonExtract runs proton scanner on in-memory content and returns
 // parsers.Extracted results compatible with spray's output format.
 func ProtonExtract(content []byte) parsers.Extracteds {
@@ -157,36 +189,49 @@ func ProtonExtract(content []byte) parsers.Extracteds {
 		return nil
 	}

-	resultMap := make(map[string]map[string]struct{})
+	resultMap := make(map[string][]extractHit)

 	for _, group := range scanner.Groups {
 		findings := scanner.ScanData(content, "response", group)
 		for _, f := range findings {
-			if resultMap[f.TemplateID] == nil {
-				resultMap[f.TemplateID] = make(map[string]struct{})
-			}
 			for _, e := range f.Extracts {
-				resultMap[f.TemplateID][e.Value] = struct{}{}
-			}
-			for _, events := range f.Matches {
-				for _, e := range events {
-					resultMap[f.TemplateID][e.Value] = struct{}{}
-				}
+				resultMap[f.TemplateID] = append(resultMap[f.TemplateID], extractHit{
+					value: e.Value, line: e.Line, offset: e.Offset, severity: f.Severity,
+				})
 			}
+			// f.Matches contains matcher hits (gate only), not extractor results — skip
 			if f.Result != nil {
 				for _, val := range f.Result.OutputExtracts {
-					resultMap[f.TemplateID][val] = struct{}{}
+					resultMap[f.TemplateID] = append(resultMap[f.TemplateID], extractHit{
+						value: val, severity: f.Severity,
+					})
 				}
 			}
 		}
 	}

 	var extracteds parsers.Extracteds
-	for name, vals := range resultMap {
-		displayName := strings.TrimPrefix(name, "spray-")
-		extracted := &parsers.Extracted{Name: displayName}
-		for v := range vals {
-			extracted.ExtractResult = append(extracted.ExtractResult, v)
+	for templateID, hits := range resultMap {
+		displayName := strings.TrimPrefix(templateID, "spray-")
+		extracted := &parsers.Extracted{
+			Name:     displayName,
+			Severity: hits[0].severity,
+		}
+		seen := make(map[string]struct{})
+		for _, h := range hits {
+			if _, ok := seen[h.value]; ok {
+				continue
+			}
+			seen[h.value] = struct{}{}
+			extracted.ExtractResult = append(extracted.ExtractResult, h.value)
+			if ExtractContextSize > 0 {
+				extracted.Items = append(extracted.Items, parsers.ExtractItem{
+					Value:  h.value,
+					Line:   h.line,
+					Offset: h.offset,
+					Ctx:    captureContext(content, h.offset, h.value, ExtractContextSize),
+				})
+			}
 		}
 		if len(extracted.ExtractResult) > 0 {
 			extracteds = append(extracteds, extracted)
@@ -222,3 +267,19 @@ func ProtonExtractorNames() []string {
 	}
 	return names
 }
+
+// IsProtonExtractor checks if a name matches a proton template ID or tag.
+func IsProtonExtractor(name string) bool {
+	protonMu.RLock()
+	defer protonMu.RUnlock()
+	if _, ok := protonTemplateMap["spray-"+name]; ok {
+		return true
+	}
+	if _, ok := protonTemplateMap[name]; ok {
+		return true
+	}
+	if _, ok := protonTagMap[name]; ok {
+		return true
+	}
+	return false
+}
@@ -0,0 +1,495 @@
+package pkg
+
+import (
+	"strings"
+	"testing"
+
+	yaml "sigs.k8s.io/yaml/goyaml.v3"
+)
+
+func mustLoadTemplates(t *testing.T) {
+	t.Helper()
+	data := LoadEmbeddedConfig("proton_rules")
+	if len(data) == 0 {
+		t.Fatal("proton_rules embedded data is empty")
+	}
+
+	var templates []interface{}
+	if err := yaml.Unmarshal(data, &templates); err != nil {
+		t.Fatalf("unmarshal proton_rules: %v", err)
+	}
+	docs := make([][]byte, 0, len(templates))
+	for _, tmpl := range templates {
+		doc, err := yaml.Marshal(tmpl)
+		if err != nil {
+			continue
+		}
+		docs = append(docs, doc)
+	}
+
+	if err := LoadProtonTemplates(docs); err != nil {
+		t.Fatalf("LoadProtonTemplates: %v", err)
+	}
+}
+
+func TestProtonTemplatesLoad(t *testing.T) {
+	mustLoadTemplates(t)
+
+	names := ProtonExtractorNames()
+	if len(names) == 0 {
+		t.Fatal("no proton templates loaded")
+	}
+
+	want := []string{
+		"spray-ipv4-address",
+		"spray-jwt-token",
+		"spray-js-file",
+		"spray-email-address",
+		"spray-aws-access-key",
+		"spray-password-leak",
+	}
+	nameSet := make(map[string]bool)
+	for _, n := range names {
+		nameSet[n] = true
+	}
+	for _, w := range want {
+		if !nameSet[w] {
+			t.Errorf("missing template: %s", w)
+		}
+	}
+}
+
+func TestProtonTagMapping(t *testing.T) {
+	mustLoadTemplates(t)
+
+	tests := []struct {
+		tag       string
+		wantMatch string
+	}{
+		{"ip", "spray-ipv4-address"},
+		{"jwt", "spray-jwt-token"},
+		{"js", "spray-js-file"},
+		{"mail", "spray-email-address"},
+		{"phone", "spray-phone-number"},
+		{"idcard", "spray-id-card"},
+		{"password", "spray-password-leak"},
+		{"username", "spray-username-leak"},
+		{"inter-ip", "spray-internal-ip"},
+		{"oss", "spray-oss-key"},
+		{"s3", "spray-s3-bucket"},
+		{"aws-ak", "spray-aws-access-key"},
+		{"rsa-key", "spray-rsa-private-key"},
+		{"jdbc", "spray-jdbc-connection"},
+		{"swagger", "spray-swagger-endpoint"},
+		{"lfi", "spray-lfi-indicator"},
+		{"upload", "spray-upload-form"},
+		{"ssrf", "spray-url-as-value"},
+		{"redirect", "spray-location-redirect"},
+		{"mac", "spray-mac-address"},
+		{"sensitive-field", "spray-sensitive-field"},
+		{"linkfinder", "spray-linkfinder"},
+		{"url-schemes", "spray-url-schemes"},
+		{"auth-header", "spray-auth-header"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.tag, func(t *testing.T) {
+			if !IsProtonExtractor(tt.tag) {
+				t.Fatalf("IsProtonExtractor(%q) = false, want true", tt.tag)
+			}
+			protonMu.RLock()
+			ids := protonTagMap[tt.tag]
+			protonMu.RUnlock()
+			found := false
+			for _, id := range ids {
+				if id == tt.wantMatch {
+					found = true
+					break
+				}
+			}
+			if !found {
+				t.Errorf("tag %q -> got %v, want to include %s", tt.tag, ids, tt.wantMatch)
+			}
+		})
+	}
+}
+
+func TestProtonCategoryTags(t *testing.T) {
+	mustLoadTemplates(t)
+
+	categories := []struct {
+		tag      string
+		minCount int
+	}{
+		{"cloud", 7},
+		{"token", 9},
+		{"credential", 9},
+		{"pentest", 20},
+		{"info", 10},
+		{"crawl", 6},
+	}
+
+	for _, cat := range categories {
+		t.Run(cat.tag, func(t *testing.T) {
+			if !IsProtonExtractor(cat.tag) {
+				t.Fatalf("IsProtonExtractor(%q) = false", cat.tag)
+			}
+			protonMu.RLock()
+			ids := protonTagMap[cat.tag]
+			protonMu.RUnlock()
+			if len(ids) < cat.minCount {
+				t.Errorf("tag %q has %d templates, want >= %d", cat.tag, len(ids), cat.minCount)
+			}
+		})
+	}
+}
+
+func TestEnableExtractorsByTag(t *testing.T) {
+	mustLoadTemplates(t)
+
+	EnableExtractors([]string{"ip"})
+
+	protonMu.RLock()
+	scanner := ProtonScanner
+	protonMu.RUnlock()
+
+	if scanner == nil || len(scanner.Groups) == 0 {
+		t.Fatal("scanner is empty after EnableExtractors([ip])")
+	}
+
+	EnableExtractors(nil)
+}
+
+func TestEnableExtractorsByID(t *testing.T) {
+	mustLoadTemplates(t)
+
+	EnableExtractors([]string{"spray-jwt-token"})
+
+	protonMu.RLock()
+	scanner := ProtonScanner
+	protonMu.RUnlock()
+
+	if scanner == nil || len(scanner.Groups) == 0 {
+		t.Fatal("scanner is empty after EnableExtractors([spray-jwt-token])")
+	}
+
+	EnableExtractors(nil)
+}
+
+func TestEnableExtractorsShortName(t *testing.T) {
+	mustLoadTemplates(t)
+
+	EnableExtractors([]string{"ipv4-address"})
+
+	protonMu.RLock()
+	scanner := ProtonScanner
+	protonMu.RUnlock()
+
+	if scanner == nil || len(scanner.Groups) == 0 {
+		t.Fatal("scanner is empty after EnableExtractors([ipv4-address]) — spray- prefix resolution failed")
+	}
+
+	EnableExtractors(nil)
+}
+
+func TestProtonExtractIP(t *testing.T) {
+	mustLoadTemplates(t)
+
+	EnableExtractors([]string{"ip"})
+	defer EnableExtractors(nil)
+
+	body := []byte(`{"origin": "192.168.1.100", "host": "example.com"}`)
+	results := ProtonExtract(body)
+
+	if len(results) == 0 {
+		t.Fatal("ProtonExtract returned no results for body containing IP")
+	}
+
+	found := false
+	for _, r := range results {
+		for _, v := range r.ExtractResult {
+			if strings.Contains(v, "192.168.1.100") {
+				found = true
+			}
+		}
+	}
+	if !found {
+		t.Errorf("expected to extract 192.168.1.100, got %v", results)
+	}
+}
+
+func TestProtonExtractJWT(t *testing.T) {
+	mustLoadTemplates(t)
+
+	EnableExtractors([]string{"jwt"})
+	defer EnableExtractors(nil)
+
+	jwt := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.dozjgNryP4J3jVmNHl0w5N_XgL0n3I9PlFUP0THsR8U"
+	body := []byte(`<html><body>token: ` + jwt + `</body></html>`)
+	results := ProtonExtract(body)
+
+	if len(results) == 0 {
+		t.Fatal("ProtonExtract returned no results for body containing JWT")
+	}
+
+	found := false
+	for _, r := range results {
+		for _, v := range r.ExtractResult {
+			if strings.Contains(v, "eyJ") {
+				found = true
+			}
+		}
+	}
+	if !found {
+		t.Errorf("expected to extract JWT, got %v", results)
+	}
+}
+
+func TestProtonExtractEmail(t *testing.T) {
+	mustLoadTemplates(t)
+
+	EnableExtractors([]string{"mail"})
+	defer EnableExtractors(nil)
+
+	body := []byte(`<html><body>Contact us at [email protected] for support</body></html>`)
+	results := ProtonExtract(body)
+
+	if len(results) == 0 {
+		t.Fatal("ProtonExtract returned no results for body containing email")
+	}
+
+	found := false
+	for _, r := range results {
+		for _, v := range r.ExtractResult {
+			if strings.Contains(v, "[email protected]") {
+				found = true
+			}
+		}
+	}
+	if !found {
+		t.Errorf("expected to extract [email protected], got %v", results)
+	}
+}
+
+func TestProtonExtractPassword(t *testing.T) {
+	mustLoadTemplates(t)
+
+	EnableExtractors([]string{"password"})
+	defer EnableExtractors(nil)
+
+	body := []byte(`config.password = "s3cretP@ss"`)
+	results := ProtonExtract(body)
+
+	if len(results) == 0 {
+		t.Fatal("ProtonExtract returned no results for body containing password leak")
+	}
+}
+
+func TestProtonExtractNoMatch(t *testing.T) {
+	mustLoadTemplates(t)
+
+	EnableExtractors([]string{"ip"})
+	defer EnableExtractors(nil)
+
+	body := []byte(`<html><body>Hello World, nothing interesting here</body></html>`)
+	results := ProtonExtract(body)
+
+	if len(results) != 0 {
+		t.Errorf("expected no extracts for clean body, got %v", results)
+	}
+}
+
+func TestProtonExtractMultipleTags(t *testing.T) {
+	mustLoadTemplates(t)
+
+	EnableExtractors([]string{"ip", "mail"})
+	defer EnableExtractors(nil)
+
+	body := []byte(`server=192.168.0.1 [email protected]`)
+	results := ProtonExtract(body)
+
+	nameSet := make(map[string]bool)
+	for _, r := range results {
+		nameSet[r.Name] = true
+	}
+
+	if !nameSet["ipv4-address"] {
+		t.Error("missing ipv4-address extraction")
+	}
+	if !nameSet["email-address"] {
+		t.Error("missing email-address extraction")
+	}
+}
+
+func TestAddCustomExtractor(t *testing.T) {
+	mustLoadTemplates(t)
+
+	AddCustomExtractor("version", `version[:\s]+(\d+\.\d+\.\d+)`)
+	defer EnableExtractors(nil)
+
+	if !IsProtonExtractor("custom-version") {
+		t.Error("custom extractor not found by ID")
+	}
+
+	body := []byte(`<html>version: 2.1.0</html>`)
+	EnableExtractors([]string{"custom-version"})
+	results := ProtonExtract(body)
+
+	if len(results) == 0 {
+		t.Fatal("custom extractor returned no results")
+	}
+}
+
+func TestProtonExtractSeverity(t *testing.T) {
+	mustLoadTemplates(t)
+
+	EnableExtractors([]string{"ip"})
+	defer EnableExtractors(nil)
+
+	body := []byte(`{"origin": "192.168.1.100"}`)
+	results := ProtonExtract(body)
+
+	if len(results) == 0 {
+		t.Fatal("no results")
+	}
+	for _, r := range results {
+		if r.Severity == "" {
+			t.Errorf("Severity is empty for %s", r.Name)
+		}
+	}
+}
+
+func TestProtonExtractContext(t *testing.T) {
+	mustLoadTemplates(t)
+
+	oldCtx := ExtractContextSize
+	ExtractContextSize = 20
+	defer func() { ExtractContextSize = oldCtx }()
+
+	EnableExtractors([]string{"ip"})
+	defer EnableExtractors(nil)
+
+	body := []byte(`server_addr = "192.168.1.100" # internal`)
+	results := ProtonExtract(body)
+
+	if len(results) == 0 {
+		t.Fatal("no results")
+	}
+	for _, r := range results {
+		if len(r.Items) == 0 {
+			t.Fatal("Items is empty when ExtractContextSize > 0")
+		}
+		item := r.Items[0]
+		if item.Value == "" {
+			t.Error("item.Value is empty")
+		}
+		if item.Ctx == "" {
+			t.Error("item.Ctx is empty — context capture failed")
+		}
+		if !strings.Contains(item.Ctx, "192.168.1.100") {
+			t.Errorf("context should contain the value, got: %s", item.Ctx)
+		}
+	}
+}
+
+func TestProtonExtractContextDisabled(t *testing.T) {
+	mustLoadTemplates(t)
+
+	oldCtx := ExtractContextSize
+	ExtractContextSize = 0
+	defer func() { ExtractContextSize = oldCtx }()
+
+	EnableExtractors([]string{"ip"})
+	defer EnableExtractors(nil)
+
+	body := []byte(`{"origin": "192.168.1.100"}`)
+	results := ProtonExtract(body)
+
+	if len(results) == 0 {
+		t.Fatal("no results")
+	}
+	for _, r := range results {
+		if len(r.Items) != 0 {
+			t.Error("Items should be empty when ExtractContextSize=0")
+		}
+	}
+}
+
+func TestProtonExtractLFI(t *testing.T) {
+	mustLoadTemplates(t)
+
+	EnableExtractors([]string{"lfi"})
+	defer EnableExtractors(nil)
+
+	body := []byte("root:x:0:0:root:/root:/bin/bash\ndaemon:x:1:1:daemon:/usr/sbin:/usr/sbin/nologin\n")
+	results := ProtonExtract(body)
+
+	if len(results) == 0 {
+		t.Fatal("LFI indicator not detected")
+	}
+}
+
+func TestProtonExtractAuthHeader(t *testing.T) {
+	mustLoadTemplates(t)
+
+	EnableExtractors([]string{"auth-header"})
+	defer EnableExtractors(nil)
+
+	body := []byte("HTTP/1.1 200 OK\r\nAuthorization: Bearer eyJhbGciOiJSUzI1NiJ9.test\r\n\r\n")
+	results := ProtonExtract(body)
+
+	if len(results) == 0 {
+		t.Fatal("Authorization header not extracted")
+	}
+}
+
+func TestProtonExtractSensitiveField(t *testing.T) {
+	mustLoadTemplates(t)
+
+	EnableExtractors([]string{"sensitive-field"})
+	defer EnableExtractors(nil)
+
+	body := []byte(`{"api_key": "sk-1234567890abcdef", "name": "test"}`)
+	results := ProtonExtract(body)
+
+	if len(results) == 0 {
+		t.Fatal("sensitive field not extracted")
+	}
+}
+
+func TestProtonExtractLocationRedirect(t *testing.T) {
+	mustLoadTemplates(t)
+
+	EnableExtractors([]string{"redirect"})
+	defer EnableExtractors(nil)
+
+	body := []byte("HTTP/1.1 302 Found\r\nLocation: https://example.com/login\r\n\r\n")
+	results := ProtonExtract(body)
+
+	if len(results) == 0 {
+		t.Fatal("Location redirect not extracted")
+	}
+}
+
+func TestIsProtonExtractor(t *testing.T) {
+	mustLoadTemplates(t)
+
+	tests := []struct {
+		name string
+		want bool
+	}{
+		{"ip", true},
+		{"cloud", true},
+		{"spray-jwt-token", true},
+		{"jwt-token", true},
+		{"nonexistent-xyz", false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := IsProtonExtractor(tt.name); got != tt.want {
+				t.Errorf("IsProtonExtractor(%q) = %v, want %v", tt.name, got, tt.want)
+			}
+		})
+	}
+}
@@ -21,6 +21,9 @@ var sprayDictData []byte
 //go:embed data/extract.bin
 var extractData []byte

+//go:embed data/proton_rules.bin
+var protonRulesData []byte
+
 //go:embed data/port.bin
 var portData []byte

@@ -34,6 +37,8 @@ func loadEmbeddedConfig(typ string) []byte {
 		return encode.MustDeflateDeCompress(sprayDictData)
 	}else if typ == "extract" {
 		return encode.MustDeflateDeCompress(extractData)
+	}else if typ == "proton_rules" {
+		return encode.MustDeflateDeCompress(protonRulesData)
 	}else if typ == "port" {
 		return encode.MustDeflateDeCompress(portData)
 	}