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

     1  package fingers
     2  
     3  import (
     4  	"github.com/chainreactors/utils/ahocorasick"
     5  )
     6  
     7  type KeywordIndex struct {
     8  	dual     *ahocorasick.DualKeywordIndex
     9  	fastPath map[int]bool // finger indices that can be resolved by AC hit alone
    10  }
    11  
    12  func NewKeywordIndex(fingers Fingers) *KeywordIndex {
    13  	builder := ahocorasick.NewDualKeywordIndexBuilder()
    14  	fastPath := make(map[int]bool)
    15  
    16  	for fi, finger := range fingers {
    17  		isFast := isSimpleFinger(finger)
    18  
    19  		for _, rule := range finger.Rules {
    20  			if rule.Regexps == nil {
    21  				continue
    22  			}
    23  
    24  			for _, body := range rule.Regexps.Body {
    25  				builder.AddBodyKeyword(body, fi)
    26  			}
    27  			for _, header := range rule.Regexps.Header {
    28  				builder.AddHeaderKeyword(header, fi)
    29  			}
    30  
    31  			if len(rule.Regexps.CompliedRegexp) > 0 || len(rule.Regexps.CompiledVulnRegexp) > 0 ||
    32  				len(rule.Regexps.MD5) > 0 || len(rule.Regexps.MMH3) > 0 || len(rule.Regexps.Cert) > 0 {
    33  				builder.AddFallback(fi)
    34  			}
    35  		}
    36  
    37  		if isFast {
    38  			fastPath[fi] = true
    39  		}
    40  	}
    41  
    42  	return &KeywordIndex{
    43  		dual:     builder.Build(),
    44  		fastPath: fastPath,
    45  	}
    46  }
    47  
    48  // isSimpleFinger returns true when a finger can be fully resolved by an AC
    49  // keyword hit: single rule, only body/header substring matchers, no regex,
    50  // no hash, no cert, no vuln patterns.
    51  func isSimpleFinger(finger *Finger) bool {
    52  	if len(finger.Rules) != 1 {
    53  		return false
    54  	}
    55  	rule := finger.Rules[0]
    56  	if rule.Regexps == nil {
    57  		return false
    58  	}
    59  	r := rule.Regexps
    60  	if len(r.CompliedRegexp) > 0 || len(r.CompiledVulnRegexp) > 0 ||
    61  		len(r.CompiledVersionRegexp) > 0 ||
    62  		len(r.MD5) > 0 || len(r.MMH3) > 0 || len(r.Cert) > 0 {
    63  		return false
    64  	}
    65  	if len(r.Body) == 0 && len(r.Header) == 0 {
    66  		return false
    67  	}
    68  	return true
    69  }
    70  
    71  func (idx *KeywordIndex) MatchCandidates(header, body []byte) map[int]bool {
    72  	return idx.dual.MatchSources(header, body)
    73  }
    74  
    75  func (idx *KeywordIndex) IsFastPath(fi int) bool {
    76  	return idx.fastPath[fi]
    77  }