github.com/chainreactors/fingers@v1.2.1/fingerprinthub/ac_index.go (about)

     1  package fingerprinthub
     2  
     3  import (
     4  	"strings"
     5  
     6  	"github.com/chainreactors/neutron/operators"
     7  	"github.com/chainreactors/neutron/templates"
     8  	"github.com/chainreactors/utils/ahocorasick"
     9  )
    10  
    11  type TemplateKeywordIndex struct {
    12  	dual     *ahocorasick.DualKeywordIndex
    13  	fastPath map[int]bool
    14  }
    15  
    16  type MatchResult struct {
    17  	Matched    map[int]bool
    18  	NeedsCheck map[int]bool
    19  }
    20  
    21  func NewTemplateKeywordIndex(tmpls []*templates.Template) *TemplateKeywordIndex {
    22  	builder := ahocorasick.NewDualKeywordIndexBuilder().SetOverlapping(true)
    23  	fastPath := make(map[int]bool)
    24  
    25  	for ti, tmpl := range tmpls {
    26  		requests := tmpl.GetRequests()
    27  		if len(requests) == 0 {
    28  			continue
    29  		}
    30  
    31  		hasKeyword := false
    32  		forceNonKeyword := false
    33  		isFastPath := true
    34  
    35  		for _, req := range requests {
    36  			if req.Matchers == nil {
    37  				continue
    38  			}
    39  
    40  			isAnd := req.MatchersCondition == "and"
    41  			if isAnd {
    42  				isFastPath = false
    43  			}
    44  			hasNonKeywordMatcher := false
    45  
    46  			for _, matcher := range req.Matchers {
    47  				if matcher.GetType() != operators.WordsMatcher {
    48  					isFastPath = false
    49  				} else if matcher.Condition == "and" && len(matcher.Words) > 1 {
    50  					isFastPath = false
    51  				}
    52  
    53  				kws := extractMatcherKeywords(matcher)
    54  				if len(kws) == 0 {
    55  					hasNonKeywordMatcher = true
    56  					if isAnd {
    57  						break
    58  					}
    59  					continue
    60  				}
    61  
    62  				hasKeyword = true
    63  				part := resolveMatcherPart(matcher)
    64  
    65  				for _, kw := range kws {
    66  					lower := strings.ToLower(kw)
    67  					if part == "body" || part == "" || part == "all" {
    68  						builder.AddBodyKeyword(lower, ti)
    69  					}
    70  					if part == "header" || part == "all_headers" || part == "all" {
    71  						builder.AddHeaderKeyword(lower, ti)
    72  					}
    73  				}
    74  			}
    75  
    76  			if isAnd && hasNonKeywordMatcher {
    77  				hasKeyword = false
    78  				break
    79  			}
    80  			if !isAnd && hasNonKeywordMatcher {
    81  				forceNonKeyword = true
    82  				isFastPath = false
    83  			}
    84  		}
    85  
    86  		if !hasKeyword || forceNonKeyword {
    87  			builder.AddFallback(ti)
    88  		}
    89  		if hasKeyword && isFastPath {
    90  			fastPath[ti] = true
    91  		}
    92  	}
    93  
    94  	return &TemplateKeywordIndex{
    95  		dual:     builder.Build(),
    96  		fastPath: fastPath,
    97  	}
    98  }
    99  
   100  func (idx *TemplateKeywordIndex) Match(headerStr, bodyStr string) MatchResult {
   101  	result := MatchResult{
   102  		Matched:    make(map[int]bool),
   103  		NeedsCheck: make(map[int]bool),
   104  	}
   105  
   106  	allHits := idx.dual.MatchSources([]byte(headerStr), []byte(bodyStr))
   107  	for ti := range allHits {
   108  		if idx.fastPath[ti] {
   109  			result.Matched[ti] = true
   110  		} else {
   111  			result.NeedsCheck[ti] = true
   112  		}
   113  	}
   114  
   115  	return result
   116  }
   117  
   118  func extractMatcherKeywords(matcher *operators.Matcher) []string {
   119  	switch matcher.GetType() {
   120  	case operators.WordsMatcher:
   121  		var kws []string
   122  		for _, word := range matcher.Words {
   123  			if len(word) >= 3 {
   124  				kws = append(kws, word)
   125  			}
   126  		}
   127  		return kws
   128  	case operators.RegexMatcher:
   129  		var kws []string
   130  		for _, pattern := range matcher.Regex {
   131  			lits := ahocorasick.ExtractLiterals(pattern)
   132  			kws = append(kws, lits...)
   133  		}
   134  		return kws
   135  	default:
   136  		return nil
   137  	}
   138  }
   139  
   140  func resolveMatcherPart(matcher *operators.Matcher) string {
   141  	part := strings.ToLower(matcher.Part)
   142  	if part == "" {
   143  		return "body"
   144  	}
   145  	return part
   146  }