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

@@ -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)

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"

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):