AhoCorasick with regex literals (#3073)

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Change Log
 
 ## master (unreleased)
+- Aho-Corasick automaton generation and scan of regex literals to speed up rule matching corkami@google.com
 
 ### New Features
 - freeze: add `--reproducible` flag that zeros dynamic header metadata

capa/rules/__init__.py

@@ -53,7 +53,95 @@
 # This narrows candidate selection from all extracted bytes to those
 # sharing a common 4-byte prefix while keeping the implementation simple.
 # See: https://github.com/mandiant/capa/issues/2128
+
+try:
+    import ahocorasick
+except ImportError:
+    ahocorasick = None
+
 _BYTES_PREFIX_SIZE = 4
+_STRING_LITERAL_MIN = 4
+
+
+def _required_literal(pattern_str: str) -> Optional[Union[str, set[str]]]:
+    """
+    Extract required literals (longest contiguous LITERAL runs >= 4 chars) from a regex pattern.
+    If no run >= 4 chars exists or on parse error, return None.
+    """
+    try:
+        try:
+            import re._parser as re_parser
+        except ImportError:
+            import sre_parse as re_parser  # type: ignore
+
+        parsed = re_parser.parse(pattern_str)
+    except Exception:
+        return None
+
+    def walk_ast(node_list) -> set[str]:
+        runs: set[str] = set()
+        cur_run: list[str] = []
+
+        def flush():
+            nonlocal cur_run
+            if cur_run:
+                s = "".join(cur_run)
+                if len(s) >= _STRING_LITERAL_MIN:
+                    runs.add(s)
+                cur_run = []
+
+        for node in node_list:
+            op = node[0]
+            av = node[1]
+
+            if op == re_parser.LITERAL:
+                cur_run.append(chr(av))
+            elif op == re_parser.SUBPATTERN:
+                sub_ast = av[-1] if isinstance(av, (tuple, list)) else av
+                sub_runs = walk_ast(sub_ast)
+                flush()
+                runs.update(sub_runs)
+            elif op in (re_parser.MAX_REPEAT, re_parser.MIN_REPEAT):
+                min_rep, _max_rep, sub_ast = av[0], av[1], av[2]
+                if min_rep >= 1 and len(sub_ast) == 1 and sub_ast[0][0] == re_parser.LITERAL:
+                    cur_run.append(chr(sub_ast[0][1]))
+                else:
+                    flush()
+                    if min_rep >= 1:
+                        sub_runs = walk_ast(sub_ast)
+                        runs.update(sub_runs)
+            elif op == re_parser.BRANCH:
+                flush()
+                alts = av[1]
+                alt_runs_list = []
+                all_valid = True
+                for alt in alts:
+                    alt_r = walk_ast(alt)
+                    valid_r = {r for r in alt_r if len(r) >= _STRING_LITERAL_MIN}
+                    if not valid_r:
+                        all_valid = False
+                        break
+                    alt_runs_list.append(valid_r)
+                if all_valid and alt_runs_list:
+                    for ar in alt_runs_list:
+                        runs.update(ar)
+            else:
+                flush()
+
+        flush()
+        return runs
+
+    try:
+        res = walk_ast(parsed)
+        valid = {r for r in res if len(r) >= _STRING_LITERAL_MIN}
+        if not valid:
+            return None
+        if len(valid) == 1:
+            return next(iter(valid))
+        return valid
+    except Exception:
+        return None
+
 
 # these are the standard metadata fields, in the preferred order.
 # when reformatted, any custom keys will come after these.
@@ -1779,6 +1867,8 @@ class _RuleFeatureIndex:
         # Built once at index time so _match() can bucket-lookup candidate bytes patterns.
         # Key -1 holds rules whose patterns are shorter than _BYTES_PREFIX_SIZE (linear fallback).
         bytes_prefix_index: dict[int, list[tuple[str, bytes]]]
+        # Optional Aho-Corasick automaton mapping string literals (lowercased) -> list of (rule_name, feature)
+        string_literal_index: Optional[Any] = None
 
     # this routine is unstable and may change before the next major release.
     @staticmethod
@@ -1932,6 +2022,8 @@ def and_score_key(item):
         string_rules: dict[str, list[Feature]] = {}
         bytes_rules_count = 0
         bytes_prefix_index: dict[int, list[tuple[str, bytes]]] = collections.defaultdict(list)
+        string_literal_automaton = ahocorasick.Automaton() if ahocorasick is not None else None
+        string_literal_entries: dict[str, list[tuple[str, Feature]]] = collections.defaultdict(list)
 
         for rule in rules:
             rule_name = rule.meta["name"]
@@ -1977,7 +2069,30 @@ def and_score_key(item):
                 )
 
             if string_features:
-                string_rules[rule_name] = cast(list[Feature], string_features)
+                unindexed_string_features: list[Feature] = []
+                for wanted_string in string_features:
+                    indexed = False
+                    if string_literal_automaton is not None:
+                        if isinstance(wanted_string, capa.features.common.Substring):
+                            sub_val = str(wanted_string.value)
+                            if len(sub_val) >= _STRING_LITERAL_MIN:
+                                string_literal_entries[sub_val.lower()].append((rule_name, wanted_string))
+                                indexed = True
+                        elif isinstance(wanted_string, capa.features.common.Regex):
+                            req_lits = _required_literal(wanted_string.re.pattern)
+                            if req_lits:
+                                if isinstance(req_lits, str):
+                                    lit_set = {req_lits}
+                                else:
+                                    lit_set = req_lits
+                                for lit in lit_set:
+                                    if len(lit) >= _STRING_LITERAL_MIN:
+                                        string_literal_entries[lit.lower()].append((rule_name, wanted_string))
+                                        indexed = True
+                    if not indexed:
+                        unindexed_string_features.append(wanted_string)
+                if unindexed_string_features:
+                    string_rules[rule_name] = unindexed_string_features
 
             bytes_features: list[capa.features.common.Bytes] = [
                 feature for feature in features if isinstance(feature, capa.features.common.Bytes)
@@ -1995,6 +2110,13 @@ def and_score_key(item):
             for feature in hashable_features:
                 rules_by_feature[feature].add(rule_name)
 
+        if string_literal_automaton is not None and string_literal_entries:
+            for word, entries in string_literal_entries.items():
+                string_literal_automaton.add_word(word, entries)
+            string_literal_automaton.make_automaton()
+        else:
+            string_literal_automaton = None
+
         logger.debug("indexing: %d features indexed for scope %s", len(rules_by_feature), scope)
         logger.debug(
             "indexing: %d indexed features are shared by more than 3 rules",
@@ -2005,7 +2127,12 @@ def and_score_key(item):
             len(string_rules),
             bytes_rules_count,
         )
-        return RuleSet._RuleFeatureIndex(rules_by_feature, string_rules, dict(bytes_prefix_index))
+        return RuleSet._RuleFeatureIndex(
+            rules_by_feature,
+            string_rules,
+            dict(bytes_prefix_index),
+            string_literal_automaton,
+        )
 
     @staticmethod
     def _get_rules_for_scope(rules, scope) -> list[Rule]:
@@ -2160,7 +2287,7 @@ def _match(self, scope: Scope, features: FeatureSet, addr: Address) -> tuple[Fea
         # be indexed, and therefore skip the scanning here, improving performance.
         # This strategy is described here:
         # https://github.com/mandiant/capa/issues/2129
-        if feature_index.string_rules:
+        if feature_index.string_rules or feature_index.string_literal_index:
             # This is a FeatureSet that contains only String features.
             # Since we'll only be evaluating String/Regex features below, we don't care about
             # other sorts of features (Mnemonic, Number, etc.) and therefore can save some time
@@ -2176,10 +2303,27 @@ def _match(self, scope: Scope, features: FeatureSet, addr: Address) -> tuple[Fea
                     string_features[feature] = locations
 
             if string_features:
-                for rule_name, wanted_strings in feature_index.string_rules.items():
-                    for wanted_string in wanted_strings:
-                        if wanted_string.evaluate(string_features):
-                            candidate_rule_names.add(rule_name)
+                if feature_index.string_literal_index:
+                    for string_feature in string_features:
+                        haystack = str(string_feature.value).lower()
+                        for _end_index, rule_entries in feature_index.string_literal_index.iter(haystack):
+                            for rule_name, wanted_feature in rule_entries:
+                                if rule_name in candidate_rule_names:
+                                    continue
+                                if isinstance(wanted_feature, capa.features.common.Substring):
+                                    candidate_rule_names.add(rule_name)
+                                elif isinstance(wanted_feature, capa.features.common.Regex):
+                                    if wanted_feature.re.search(str(string_feature.value)):
+                                        candidate_rule_names.add(rule_name)
+
+                if feature_index.string_rules:
+                    for rule_name, wanted_strings in feature_index.string_rules.items():
+                        if rule_name in candidate_rule_names:
+                            continue
+                        for wanted_string in wanted_strings:
+                            if wanted_string.evaluate(string_features):
+                                candidate_rule_names.add(rule_name)
+                                break
 
         # Like with String/Regex features above, Bytes features cannot be matched via hash lookup.
         # To avoid a linear scan of every bytes rule against every extracted bytes feature,

pyproject.toml

@@ -85,6 +85,7 @@ dependencies = [
     "protobuf>=5",
     "msgspec>=0.18.6",
     "xmltodict>=0.13.0",
+    "pyahocorasick>=2.0",
 
     # ---------------------------------------
     # Dependencies that we develop

requirements.txt

@@ -46,4 +46,5 @@ sortedcontainers==2.4.0
 viv-utils==0.8.0
 vivisect==1.3.2
 msgspec==0.21.1
+pyahocorasick>=2.0
 bump-my-version==1.5.0

tests/test_match.py

@@ -386,3 +386,126 @@ def test_match_no_duplicate_candidate_evaluations():
 
     # Ensure Target Rule was evaluated and returned exactly ONCE
     assert len(matches["Target Rule"]) == 1
+
+
+def test_required_literal_unit():
+    assert capa.rules._required_literal("test[0-9]+") == "test"
+    assert capa.rules._required_literal("a[0-9]b") is None
+    assert capa.rules._required_literal("(test|b)") is None
+    assert capa.rules._required_literal("(test|example)") == {"test", "example"}
+
+
+def test_string_literal_prefilter_regex_confirmation():
+    rule = textwrap.dedent("""
+        rule:
+            meta:
+                name: test regex confirmation
+                scopes:
+                    static: function
+                    dynamic: process
+            features:
+                - string: /malware[0-9]+/
+        """)
+    r = capa.rules.Rule.from_yaml(rule)
+
+    feat1 = {capa.features.common.String("found malware123 sample"): {0x0}}
+    _, matches1 = match([r], feat1, 0x0)
+    assert "test regex confirmation" in matches1
+
+    feat2 = {capa.features.common.String("this malwareXYZ string"): {0x0}}
+    _, matches2 = match([r], feat2, 0x0)
+    assert "test regex confirmation" not in matches2
+
+    feat3 = {capa.features.common.String("totally clean"): {0x0}}
+    _, matches3 = match([r], feat3, 0x0)
+    assert "test regex confirmation" not in matches3
+
+
+def test_differential_parity_prefilter_on_vs_off(monkeypatch):
+    """
+    Differential parity test: Match rules with prefilter ON vs OFF (ahocorasick forced to None).
+    Expected: identical match results in both.
+    """
+    rules = [
+        capa.rules.Rule.from_yaml(
+            textwrap.dedent("""
+                rule:
+                    meta:
+                        name: Rule 1 Substring
+                        scopes:
+                            static: function
+                            dynamic: process
+                    features:
+                        - string: "specific_literal_string"
+                """)
+        ),
+        capa.rules.Rule.from_yaml(
+            textwrap.dedent("""
+                rule:
+                    meta:
+                        name: Rule 2 Regex
+                        scopes:
+                            static: function
+                            dynamic: process
+                    features:
+                        - string: /pattern_[0-9]{3}/
+                """)
+        ),
+        capa.rules.Rule.from_yaml(
+            textwrap.dedent("""
+                rule:
+                    meta:
+                        name: Rule 3 No Literal Regex
+                        scopes:
+                            static: function
+                            dynamic: process
+                    features:
+                        - string: /[a-z]/
+                """)
+        ),
+    ]
+
+    features = {
+        capa.features.common.String("here is a specific_literal_string inside"): {0x0},
+        capa.features.common.String("here is pattern_123 inside"): {0x0},
+        capa.features.common.String("here is pattern_abc inside"): {0x0},
+    }
+
+    _, matches_on = match(rules, features, 0x0)
+
+    monkeypatch.setattr(capa.rules, "ahocorasick", None)
+    _, matches_off = match(rules, features, 0x0)
+
+    assert matches_on.keys() == matches_off.keys()
+    for k in matches_on:
+        assert len(matches_on[k]) == len(matches_off[k])
+
+
+def test_string_literal_prefilter_shared_literal():
+    str_regex = textwrap.dedent("""
+        rule:
+            meta:
+                name: rule 1 shared literal
+                scopes:
+                    static: function
+                    dynamic: process
+            features:
+                - string: /.*shared_key.*/
+        """)
+    str_search = textwrap.dedent("""
+        rule:
+            meta:
+                name: rule 2 shared literal
+                scopes:
+                    static: function
+                    dynamic: process
+            features:
+                - string: /shared_key/
+        """)
+    rule_regex = capa.rules.Rule.from_yaml(str_regex)
+    rule_search = capa.rules.Rule.from_yaml(str_search)
+
+    feat = {capa.features.common.String("prefix shared_key suffix"): {0x0}}
+    _, matches = match([rule_regex, rule_search], feat, 0x0)
+    assert "rule 1 shared literal" in matches
+    assert "rule 2 shared literal" in matches