From bee91583e57780ec59691ec8d13a923d68282c2f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ana=20Mar=C3=ADa=20Mart=C3=ADnez=20G=C3=B3mez?= Date: Mon, 27 Jul 2020 15:09:25 +0200 Subject: [PATCH] Enable descriptions for statement nodes Enable descriptions for statement nodes such as and and or. Use of case in: fireeye/capa-rules/pull/51 Documentation should be added in capa-rules. --- capa/engine.py | 30 +++++++++++++++++------------- capa/render/__init__.py | 2 ++ capa/rules.py | 18 +++++++++--------- 3 files changed, 28 insertions(+), 22 deletions(-) diff --git a/capa/engine.py b/capa/engine.py index e3698aeb..3d7b5ca7 100644 --- a/capa/engine.py +++ b/capa/engine.py @@ -20,12 +20,16 @@ class Statement(object): and to declare the interface method `evaluate` """ - def __init__(self): + def __init__(self, description=None): super(Statement, self).__init__() self.name = self.__class__.__name__ + self.description = description def __str__(self): - return "%s(%s)" % (self.name.lower(), ",".join(map(str, self.get_children()))) + if self.description: + return "%s(%s = %s)" % (self.name.lower(), ",".join(map(str, self.get_children())), self.description) + else: + return "%s(%s)" % (self.name.lower(), ",".join(map(str, self.get_children()))) def __repr__(self): return str(self) @@ -104,8 +108,8 @@ class Result(object): class And(Statement): """match if all of the children evaluate to True.""" - def __init__(self, children): - super(And, self).__init__() + def __init__(self, children, description=None): + super(And, self).__init__(description=description) self.children = children def evaluate(self, ctx): @@ -117,8 +121,8 @@ class And(Statement): class Or(Statement): """match if any of the children evaluate to True.""" - def __init__(self, children): - super(Or, self).__init__() + def __init__(self, children, description=None): + super(Or, self).__init__(description=description) self.children = children def evaluate(self, ctx): @@ -130,8 +134,8 @@ class Or(Statement): class Not(Statement): """match only if the child evaluates to False.""" - def __init__(self, child): - super(Not, self).__init__() + def __init__(self, child, description=None): + super(Not, self).__init__(description=description) self.child = child def evaluate(self, ctx): @@ -143,10 +147,10 @@ class Not(Statement): class Some(Statement): """match if at least N of the children evaluate to True.""" - def __init__(self, count, children): - super(Some, self).__init__() + def __init__(self, count, children, description=None): + super(Some, self).__init__(description=description) self.count = count - self.children = list(children) + self.children = children def evaluate(self, ctx): results = [child.evaluate(ctx) for child in self.children] @@ -161,8 +165,8 @@ class Some(Statement): class Range(Statement): """match if the child is contained in the ctx set with a count in the given range.""" - def __init__(self, child, min=None, max=None): - super(Range, self).__init__() + def __init__(self, child, min=None, max=None, description=None): + super(Range, self).__init__(description=description) self.child = child self.min = min if min is not None else 0 self.max = max if max is not None else (1 << 64 - 1) diff --git a/capa/render/__init__.py b/capa/render/__init__.py index 6c8b4664..ceb79d5a 100644 --- a/capa/render/__init__.py +++ b/capa/render/__init__.py @@ -28,6 +28,8 @@ def convert_statement_to_result_document(statement): """ statement_type = statement.name.lower() result = {"type": statement_type} + if statement.description: + result["description"] = statement.description if statement_type == "some" and statement.count == 0: result["type"] = "optional" diff --git a/capa/rules.py b/capa/rules.py index 2b6383cb..92cd0c56 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]], description=d.get("description")) 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]], description=d.get("description")) elif key == "not": if len(d[key]) != 1: raise InvalidRule("not statement must have exactly one child statement") - return Not(build_statements(d[key][0], scope)) + return Not(build_statements(d[key][0], scope), description=d.get("description")) 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]], description=d.get("description")) 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]], description=d.get("description")) elif key == "function": if scope != FILE_SCOPE: @@ -338,18 +338,18 @@ def build_statements(d, scope): count = d[key] if isinstance(count, int): - return Range(feature, min=count, max=count) + return Range(feature, min=count, max=count, description=d.get("description")) elif count.endswith(" or more"): min = parse_int(count[: -len(" or more")]) max = None - return Range(feature, min=min, max=max) + return Range(feature, min=min, max=max, description=d.get("description")) elif count.endswith(" or fewer"): min = None max = parse_int(count[: -len(" or fewer")]) - return Range(feature, min=min, max=max) + return Range(feature, min=min, max=max, description=d.get("description")) elif count.startswith("("): min, max = parse_range(count) - return Range(feature, min=min, max=max) + return Range(feature, min=min, max=max, description=d.get("description")) else: raise InvalidRule("unexpected range: %s" % (count)) elif key == "string" and not isinstance(d[key], six.string_types):