diff --git a/capa/engine.py b/capa/engine.py index 34e9b2dd..e3698aeb 100644 --- a/capa/engine.py +++ b/capa/engine.py @@ -104,9 +104,9 @@ class Result(object): class And(Statement): """match if all of the children evaluate to True.""" - def __init__(self, *children): + def __init__(self, children): super(And, self).__init__() - self.children = list(children) + self.children = children def evaluate(self, ctx): results = [child.evaluate(ctx) for child in self.children] @@ -117,9 +117,9 @@ class And(Statement): class Or(Statement): """match if any of the children evaluate to True.""" - def __init__(self, *children): + def __init__(self, children): super(Or, self).__init__() - self.children = list(children) + self.children = children def evaluate(self, ctx): results = [child.evaluate(ctx) for child in self.children] @@ -143,7 +143,7 @@ class Not(Statement): class Some(Statement): """match if at least N of the children evaluate to True.""" - def __init__(self, count, *children): + def __init__(self, count, children): super(Some, self).__init__() self.count = count self.children = list(children) diff --git a/capa/rules.py b/capa/rules.py index 3dd92363..2b6383cb 100644 --- a/capa/rules.py +++ b/capa/rules.py @@ -265,21 +265,21 @@ def build_statements(d, scope): key = list(d.keys())[0] if key == "and": - return And(*[build_statements(dd, scope) for dd in d[key]]) + return And([build_statements(dd, scope) for dd in d[key]]) elif key == "or": - return Or(*[build_statements(dd, scope) for dd in d[key]]) + return Or([build_statements(dd, scope) for dd in d[key]]) elif key == "not": if len(d[key]) != 1: raise InvalidRule("not statement must have exactly one child statement") - return Not(*[build_statements(dd, scope) for dd in d[key]]) + return Not(build_statements(d[key][0], scope)) elif key.endswith(" or more"): count = int(key[: -len("or more")]) - return Some(count, *[build_statements(dd, scope) for dd in d[key]]) + return Some(count, [build_statements(dd, scope) for dd in d[key]]) elif key == "optional": # `optional` is an alias for `0 or more` # which is useful for documenting behaviors, # like with `write file`, we might say that `WriteFile` is optionally found alongside `CreateFileA`. - return Some(0, *[build_statements(dd, scope) for dd in d[key]]) + return Some(0, [build_statements(dd, scope) for dd in d[key]]) elif key == "function": if scope != FILE_SCOPE: @@ -288,7 +288,7 @@ def build_statements(d, scope): if len(d[key]) != 1: raise InvalidRule("subscope must have exactly one child statement") - return Subscope(FUNCTION_SCOPE, *[build_statements(dd, FUNCTION_SCOPE) for dd in d[key]]) + return Subscope(FUNCTION_SCOPE, build_statements(d[key][0], FUNCTION_SCOPE)) elif key == "basic block": if scope != FUNCTION_SCOPE: @@ -297,7 +297,7 @@ def build_statements(d, scope): if len(d[key]) != 1: raise InvalidRule("subscope must have exactly one child statement") - return Subscope(BASIC_BLOCK_SCOPE, *[build_statements(dd, BASIC_BLOCK_SCOPE) for dd in d[key]]) + return Subscope(BASIC_BLOCK_SCOPE, build_statements(d[key][0], BASIC_BLOCK_SCOPE)) elif key.startswith("count(") and key.endswith(")"): # e.g.: diff --git a/tests/test_engine.py b/tests/test_engine.py index a0fb429c..c959283e 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -22,21 +22,21 @@ def test_number(): def test_and(): - assert And(Number(1)).evaluate({Number(0): {1}}) == False - assert And(Number(1)).evaluate({Number(1): {1}}) == True - assert And(Number(1), Number(2)).evaluate({Number(0): {1}}) == False - assert And(Number(1), Number(2)).evaluate({Number(1): {1}}) == False - assert And(Number(1), Number(2)).evaluate({Number(2): {1}}) == False - assert And(Number(1), Number(2)).evaluate({Number(1): {1}, Number(2): {2}}) == True + assert And([Number(1)]).evaluate({Number(0): {1}}) == False + assert And([Number(1)]).evaluate({Number(1): {1}}) == True + assert And([Number(1), Number(2)]).evaluate({Number(0): {1}}) == False + assert And([Number(1), Number(2)]).evaluate({Number(1): {1}}) == False + assert And([Number(1), Number(2)]).evaluate({Number(2): {1}}) == False + assert And([Number(1), Number(2)]).evaluate({Number(1): {1}, Number(2): {2}}) == True def test_or(): - assert Or(Number(1)).evaluate({Number(0): {1}}) == False - assert Or(Number(1)).evaluate({Number(1): {1}}) == True - assert Or(Number(1), Number(2)).evaluate({Number(0): {1}}) == False - assert Or(Number(1), Number(2)).evaluate({Number(1): {1}}) == True - assert Or(Number(1), Number(2)).evaluate({Number(2): {1}}) == True - assert Or(Number(1), Number(2)).evaluate({Number(1): {1}, Number(2): {2}}) == True + assert Or([Number(1)]).evaluate({Number(0): {1}}) == False + assert Or([Number(1)]).evaluate({Number(1): {1}}) == True + assert Or([Number(1), Number(2)]).evaluate({Number(0): {1}}) == False + assert Or([Number(1), Number(2)]).evaluate({Number(1): {1}}) == True + assert Or([Number(1), Number(2)]).evaluate({Number(2): {1}}) == True + assert Or([Number(1), Number(2)]).evaluate({Number(1): {1}, Number(2): {2}}) == True def test_not(): @@ -45,20 +45,20 @@ def test_not(): def test_some(): - assert Some(0, Number(1)).evaluate({Number(0): {1}}) == True - assert Some(1, Number(1)).evaluate({Number(0): {1}}) == False + assert Some(0, [Number(1)]).evaluate({Number(0): {1}}) == True + assert Some(1, [Number(1)]).evaluate({Number(0): {1}}) == False - assert Some(2, Number(1), Number(2), Number(3)).evaluate({Number(0): {1}}) == False - assert Some(2, Number(1), Number(2), Number(3)).evaluate({Number(0): {1}, Number(1): {1}}) == False - assert Some(2, Number(1), Number(2), Number(3)).evaluate({Number(0): {1}, Number(1): {1}, Number(2): {1}}) == True + assert Some(2, [Number(1), Number(2), Number(3)]).evaluate({Number(0): {1}}) == False + assert Some(2, [Number(1), Number(2), Number(3)]).evaluate({Number(0): {1}, Number(1): {1}}) == False + assert Some(2, [Number(1), Number(2), Number(3)]).evaluate({Number(0): {1}, Number(1): {1}, Number(2): {1}}) == True assert ( - Some(2, Number(1), Number(2), Number(3)).evaluate( + Some(2, [Number(1), Number(2), Number(3)]).evaluate( {Number(0): {1}, Number(1): {1}, Number(2): {1}, Number(3): {1}} ) == True ) assert ( - Some(2, Number(1), Number(2), Number(3)).evaluate( + Some(2, [Number(1), Number(2), Number(3)]).evaluate( {Number(0): {1}, Number(1): {1}, Number(2): {1}, Number(3): {1}, Number(4): {1},} ) == True @@ -66,11 +66,11 @@ def test_some(): def test_complex(): - assert True == Or(And(Number(1), Number(2)), Or(Number(3), Some(2, Number(4), Number(5), Number(6))),).evaluate( - {Number(5): {1}, Number(6): {1}, Number(7): {1}, Number(8): {1}} - ) + assert True == Or( + [And([Number(1), Number(2)]), Or([Number(3), Some(2, [Number(4), Number(5), Number(6)])])] + ).evaluate({Number(5): {1}, Number(6): {1}, Number(7): {1}, Number(8): {1}}) - assert False == Or(And(Number(1), Number(2)), Or(Number(3), Some(2, Number(4), Number(5)))).evaluate( + assert False == Or([And([Number(1), Number(2)]), Or([Number(3), Some(2, [Number(4), Number(5)])])]).evaluate( {Number(5): {1}, Number(6): {1}, Number(7): {1}, Number(8): {1}} )