Merge pull request #468 from fireeye/features/support-string-values-special-chars

add support for string features with special characters e.g. '\n'

capa/features/__init__.py

@@ -38,6 +38,20 @@ def hex_string(h):
     return " ".join(h[i : i + 2] for i in range(0, len(h), 2)).upper()
 
 
+def escape_string(s):
+    """escape special characters"""
+    s = repr(s)
+    if not s.startswith(('"', "'")):
+        # u'hello\r\nworld' -> hello\\r\\nworld
+        s = s[2:-1]
+    else:
+        # 'hello\r\nworld' -> hello\\r\\nworld
+        s = s[1:-1]
+    s = s.replace("\\'", "'")  # repr() may escape "'" in some edge cases, remove
+    s = s.replace('"', '\\"')  # repr() does not escape '"', add
+    return s
+
+
 class Feature(object):
     def __init__(self, value, arch=None, description=None):
         """

capa/ida/plugin/model.py

@@ -488,13 +488,17 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
 
         @param feature: capa feature read from doc
         """
-        if feature[feature["type"]]:
+        key = feature["type"]
+        value = feature[feature["type"]]
+        if value:
+            if key == "string":
+                value = '"%s"' % capa.features.escape_string(value)
             if feature.get("description", ""):
-                return "%s(%s = %s)" % (feature["type"], feature[feature["type"]], feature["description"])
+                return "%s(%s = %s)" % (key, value, feature["description"])
             else:
-                return "%s(%s)" % (feature["type"], feature[feature["type"]])
+                return "%s(%s)" % (key, value)
         else:
-            return "%s" % feature["type"]
+            return "%s" % key
 
     def render_capa_doc_feature_node(self, parent, feature, locations, doc):
         """process capa doc feature node
@@ -551,7 +555,9 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
             )
 
         if feature["type"] == "regex":
-            return CapaExplorerStringViewItem(parent, display, location, feature["match"])
+            return CapaExplorerStringViewItem(
+                parent, display, location, '"%s"' % capa.features.escape_string(feature["match"])
+            )
 
         if feature["type"] == "basicblock":
             return CapaExplorerBlockItem(parent, location)
@@ -576,7 +582,9 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
 
         if feature["type"] in ("string",):
             # display string preview
-            return CapaExplorerStringViewItem(parent, display, location, feature[feature["type"]])
+            return CapaExplorerStringViewItem(
+                parent, display, location, '"%s"' % capa.features.escape_string(feature[feature["type"]])
+            )
 
         if feature["type"] in ("import", "export"):
             # display no preview

capa/ida/plugin/view.py

@@ -624,11 +624,23 @@ class CapaExplorerRulgenEditor(QtWidgets.QTreeWidget):
 
         # single features
         for (k, v) in filter(lambda t: t[1] == 1, counted):
-            self.new_feature_node(self.root, ("- %s: %s" % (k.name.lower(), k.get_value_str()), ""))
+            if isinstance(k, (capa.features.String,)):
+                value = '"%s"' % capa.features.escape_string(k.get_value_str())
+            else:
+                value = k.get_value_str()
+            self.new_feature_node(self.root, ("- %s: %s" % (k.name.lower(), value), ""))
 
         # n > 1 features
         for (k, v) in filter(lambda t: t[1] > 1, counted):
-            self.new_feature_node(self.root, ("- count(%s): %d" % (str(k), v), ""))
+            if k.value:
+                if isinstance(k, (capa.features.String,)):
+                    value = '"%s"' % capa.features.escape_string(k.get_value_str())
+                else:
+                    value = k.get_value_str()
+                display = "- count(%s(%s)): %d" % (k.name.lower(), value, v)
+            else:
+                display = "- count(%s): %d" % (k.name.lower(), v)
+            self.new_feature_node(self.root, (display, ""))
 
         self.expandAll()
         self.update_preview()
@@ -882,16 +894,20 @@ class CapaExplorerRulegenFeatures(QtWidgets.QTreeWidget):
         def format_address(e):
             return "%X" % e if e else ""
 
+        def format_feature(feature):
+            """ """
+            name = feature.name.lower()
+            value = feature.get_value_str()
+            if isinstance(feature, (capa.features.String,)):
+                value = '"%s"' % capa.features.escape_string(value)
+            return "%s(%s)" % (name, value)
+
         for (feature, eas) in sorted(features.items(), key=lambda k: sorted(k[1])):
             if isinstance(feature, capa.features.basicblock.BasicBlock):
                 # filter basic blocks for now, we may want to add these back in some time
                 # in the future
                 continue
 
-            if isinstance(feature, capa.features.String):
-                # strip string for display
-                feature.value = feature.value.strip()
-
             # level 0
             if type(feature) not in self.parent_items:
                 self.parent_items[type(feature)] = self.new_parent_node(parent, (feature.name.lower(),))
@@ -900,20 +916,22 @@ class CapaExplorerRulegenFeatures(QtWidgets.QTreeWidget):
             if feature not in self.parent_items:
                 if len(eas) > 1:
                     self.parent_items[feature] = self.new_parent_node(
-                        self.parent_items[type(feature)], (str(feature),), feature=feature
+                        self.parent_items[type(feature)], (format_feature(feature),), feature=feature
                     )
                 else:
                     self.parent_items[feature] = self.new_leaf_node(
-                        self.parent_items[type(feature)], (str(feature),), feature=feature
+                        self.parent_items[type(feature)], (format_feature(feature),), feature=feature
                     )
 
             # level n > 1
             if len(eas) > 1:
                 for ea in sorted(eas):
-                    self.new_leaf_node(self.parent_items[feature], (str(feature), format_address(ea)), feature=feature)
+                    self.new_leaf_node(
+                        self.parent_items[feature], (format_feature(feature), format_address(ea)), feature=feature
+                    )
             else:
                 ea = eas.pop()
-                for (i, v) in enumerate((str(feature), format_address(ea))):
+                for (i, v) in enumerate((format_feature(feature), format_address(ea))):
                     self.parent_items[feature].setText(i, v)
                 self.parent_items[feature].setData(0, 0x100, feature)
 

capa/render/vverbose.py

@@ -56,7 +56,11 @@ def render_statement(ostream, match, statement, indent=0):
         child = statement["child"]
 
         if child[child["type"]]:
-            value = rutils.bold2(child[child["type"]])
+            if child["type"] == "string":
+                value = '"%s"' % capa.features.escape_string(child[child["type"]])
+            else:
+                value = child[child["type"]]
+            value = rutils.bold2(value)
             if child.get("description"):
                 ostream.write("count(%s(%s = %s)): " % (child["type"], value, child["description"]))
             else:
@@ -90,6 +94,9 @@ def render_feature(ostream, match, feature, indent=0):
         key = "string"  # render string for regex to mirror the rule source
         value = feature["match"]  # the match provides more information than the value for regex
 
+    if key == "string":
+        value = '"%s"' % capa.features.escape_string(value)
+
     ostream.write(key)
     ostream.write(": ")
 

scripts/lint.py

@@ -25,6 +25,8 @@ import argparse
 import itertools
 import posixpath
 
+import ruamel.yaml
+
 import capa.main
 import capa.rules
 import capa.engine
@@ -342,6 +344,32 @@ class FormatIncorrect(Lint):
         return False
 
 
+class FormatStringQuotesIncorrect(Lint):
+    name = "rule string quotes incorrect"
+
+    def check_rule(self, ctx, rule):
+        events = capa.rules.Rule._get_ruamel_yaml_parser().parse(rule.definition)
+        for key in events:
+            if not (isinstance(key, ruamel.yaml.ScalarEvent) and key.value == "string"):
+                continue
+            value = next(events)  # assume value is next event
+            if not isinstance(value, ruamel.yaml.ScalarEvent):
+                # ignore non-scalar
+                continue
+            if value.value.startswith("/") and value.value.endswith(("/", "/i")):
+                # ignore regex for now
+                continue
+            if value.style is None:
+                # no quotes
+                self.recommendation = 'add double quotes to "%s"' % value.value
+                return True
+            if value.style == "'":
+                # single quote
+                self.recommendation = 'change single quotes to double quotes for "%s"' % value.value
+                return True
+        return False
+
+
 def run_lints(lints, ctx, rule):
     for lint in lints:
         if lint.check_rule(ctx, rule):
@@ -402,6 +430,7 @@ def lint_features(ctx, rule):
 FORMAT_LINTS = (
     FormatLineFeedEOL(),
     FormatSingleEmptyLineEOF(),
+    FormatStringQuotesIncorrect(),
     FormatIncorrect(),
 )
 

tests/test_rules.py

@@ -681,6 +681,25 @@ def test_explicit_string_values_int():
     assert (String("0x123") in children) == True
 
 
+def test_string_values_special_characters():
+    rule = textwrap.dedent(
+        """
+        rule:
+            meta:
+                name: test rule
+            features:
+                - or:
+                    - string: "hello\\r\\nworld"
+                    - string: "bye\\nbye"
+                      description: "test description"
+        """
+    )
+    r = capa.rules.Rule.from_yaml(rule)
+    children = list(r.statement.get_children())
+    assert (String("hello\r\nworld") in children) == True
+    assert (String("bye\nbye") in children) == True
+
+
 def test_regex_values_always_string():
     rules = [
         capa.rules.Rule.from_yaml(