From 527e993bb446686fefb2fb528017ae13710b5dc9 Mon Sep 17 00:00:00 2001 From: William Ballenthin Date: Wed, 9 Jun 2021 22:30:43 -0600 Subject: [PATCH] engine: remove dependency on rules, fixing circular import --- capa/engine.py | 32 -------------------------------- capa/rules.py | 34 +++++++++++++++++++++++++++++++++- tests/test_engine.py | 20 ++++++++++---------- tests/test_rules.py | 6 +++--- 4 files changed, 46 insertions(+), 46 deletions(-) diff --git a/capa/engine.py b/capa/engine.py index 2356d0f5..f496a146 100644 --- a/capa/engine.py +++ b/capa/engine.py @@ -9,7 +9,6 @@ import copy import collections -import capa.rules import capa.features.common @@ -200,37 +199,6 @@ class Subscope(Statement): raise ValueError("cannot evaluate a subscope directly!") -def topologically_order_rules(rules): - """ - order the given rules such that dependencies show up before dependents. - this means that as we match rules, we can add features for the matches, and these - will be matched by subsequent rules if they follow this order. - - assumes that the rule dependency graph is a DAG. - """ - # we evaluate `rules` multiple times, so if its a generator, realize it into a list. - rules = list(rules) - namespaces = capa.rules.index_rules_by_namespace(rules) - rules = {rule.name: rule for rule in rules} - seen = set([]) - ret = [] - - def rec(rule): - if rule.name in seen: - return - - for dep in rule.get_dependencies(namespaces): - rec(rules[dep]) - - ret.append(rule) - seen.add(rule.name) - - for rule in rules.values(): - rec(rule) - - return ret - - def match(rules, features, va): """ Args: diff --git a/capa/rules.py b/capa/rules.py index 184c046e..beb371c3 100644 --- a/capa/rules.py +++ b/capa/rules.py @@ -856,6 +856,38 @@ def index_rules_by_namespace(rules): return dict(namespaces) +def topologically_order_rules(rules): + """ + order the given rules such that dependencies show up before dependents. + this means that as we match rules, we can add features for the matches, and these + will be matched by subsequent rules if they follow this order. + + assumes that the rule dependency graph is a DAG. + """ + # we evaluate `rules` multiple times, so if its a generator, realize it into a list. + rules = list(rules) + namespaces = index_rules_by_namespace(rules) + rules = {rule.name: rule for rule in rules} + seen = set([]) + ret = [] + + def rec(rule): + if rule.name in seen: + return + + for dep in rule.get_dependencies(namespaces): + rec(rules[dep]) + + ret.append(rule) + seen.add(rule.name) + + for rule in rules.values(): + rec(rule) + + return ret + + + class RuleSet(object): """ a ruleset is initialized with a collection of rules, which it verifies and sorts into scopes. @@ -918,7 +950,7 @@ class RuleSet(object): continue scope_rules.update(get_rules_and_dependencies(rules, rule.name)) - return get_rules_with_scope(capa.engine.topologically_order_rules(scope_rules), scope) + return get_rules_with_scope(topologically_order_rules(scope_rules), scope) @staticmethod def _extract_subscope_rules(rules): diff --git a/tests/test_engine.py b/tests/test_engine.py index 5f0d7cce..642ddb49 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -266,7 +266,7 @@ def test_match_matched_rules(): ] features, matches = capa.engine.match( - capa.engine.topologically_order_rules(rules), + capa.rules.topologically_order_rules(rules), {capa.features.insn.Number(100): {1}}, 0x0, ) @@ -276,7 +276,7 @@ def test_match_matched_rules(): # the ordering of the rules must not matter, # the engine should match rules in an appropriate order. features, matches = capa.engine.match( - capa.engine.topologically_order_rules(reversed(rules)), + capa.rules.topologically_order_rules(reversed(rules)), {capa.features.insn.Number(100): {1}}, 0x0, ) @@ -324,28 +324,28 @@ def test_regex(): ), ] features, matches = capa.engine.match( - capa.engine.topologically_order_rules(rules), + capa.rules.topologically_order_rules(rules), {capa.features.insn.Number(100): {1}}, 0x0, ) assert capa.features.common.MatchedRule("test rule") not in features features, matches = capa.engine.match( - capa.engine.topologically_order_rules(rules), + capa.rules.topologically_order_rules(rules), {capa.features.common.String("aaaa"): {1}}, 0x0, ) assert capa.features.common.MatchedRule("test rule") not in features features, matches = capa.engine.match( - capa.engine.topologically_order_rules(rules), + capa.rules.topologically_order_rules(rules), {capa.features.common.String("aBBBBa"): {1}}, 0x0, ) assert capa.features.common.MatchedRule("test rule") not in features features, matches = capa.engine.match( - capa.engine.topologically_order_rules(rules), + capa.rules.topologically_order_rules(rules), {capa.features.common.String("abbbba"): {1}}, 0x0, ) @@ -370,7 +370,7 @@ def test_regex_ignorecase(): ), ] features, matches = capa.engine.match( - capa.engine.topologically_order_rules(rules), + capa.rules.topologically_order_rules(rules), {capa.features.common.String("aBBBBa"): {1}}, 0x0, ) @@ -393,7 +393,7 @@ def test_regex_complex(): ), ] features, matches = capa.engine.match( - capa.engine.topologically_order_rules(rules), + capa.rules.topologically_order_rules(rules), {capa.features.common.String(r"Hardware\Key\key with spaces\some value"): {1}}, 0x0, ) @@ -451,7 +451,7 @@ def test_match_namespace(): ] features, matches = capa.engine.match( - capa.engine.topologically_order_rules(rules), + capa.rules.topologically_order_rules(rules), {capa.features.insn.API("CreateFile"): {1}}, 0x0, ) @@ -463,7 +463,7 @@ def test_match_namespace(): assert capa.features.common.MatchedRule("file/create/CreateFile") in features features, matches = capa.engine.match( - capa.engine.topologically_order_rules(rules), + capa.rules.topologically_order_rules(rules), {capa.features.insn.API("WriteFile"): {1}}, 0x0, ) diff --git a/tests/test_rules.py b/tests/test_rules.py index 7dc1b666..12791d2e 100644 --- a/tests/test_rules.py +++ b/tests/test_rules.py @@ -751,14 +751,14 @@ def test_regex_values_always_string(): ), ] features, matches = capa.engine.match( - capa.engine.topologically_order_rules(rules), + capa.rules.topologically_order_rules(rules), {capa.features.common.String("123"): {1}}, 0x0, ) - assert capa.features.comm.MatchedRule("test rule") in features + assert capa.features.common.MatchedRule("test rule") in features features, matches = capa.engine.match( - capa.engine.topologically_order_rules(rules), + capa.rules.topologically_order_rules(rules), {capa.features.common.String("0x123"): {1}}, 0x0, )