From c37365f0453b53e6510256b3e252bdf9c1e936cf Mon Sep 17 00:00:00 2001
From: Moritz Raabe <moritz.raabe@fireeye.com>
Date: Thu, 2 Jul 2020 18:30:44 +0200
Subject: [PATCH] fix render, cleanup feature string display

---
 capa/features/__init__.py |  7 ++++++-
 capa/render/utils.py      |  5 -----
 capa/render/vverbose.py   | 30 +++++-------------------------
 tests/test_main.py        |  8 +++++++-
 4 files changed, 18 insertions(+), 32 deletions(-)

diff --git a/capa/features/__init__.py b/capa/features/__init__.py
index 28931ad9..4ca21823 100644
--- a/capa/features/__init__.py
+++ b/capa/features/__init__.py
@@ -16,6 +16,11 @@ def bytes_to_str(b):
         return codecs.encode(b, 'hex')
 
 
+def hex_string(h):
+    """ render hex string e.g. "0a40b1" as "0A 40 B1" """
+    return ' '.join(h[i:i + 2] for i in range(0, len(h), 2)).upper()
+
+
 class Feature(object):
     def __init__(self, args, description=None):
         super(Feature, self).__init__()
@@ -106,7 +111,7 @@ class Bytes(Feature):
         return capa.engine.Result(False, self, [])
 
     def get_args_str(self):
-        return bytes_to_str(self.value).upper()
+        return hex_string(bytes_to_str(self.value))
 
     def freeze_serialize(self):
         return (self.__class__.__name__,
diff --git a/capa/render/utils.py b/capa/render/utils.py
index 712f0c77..d417d213 100644
--- a/capa/render/utils.py
+++ b/capa/render/utils.py
@@ -17,11 +17,6 @@ def hex(n):
     return '0x%X' % n
 
 
-def hex_string(h):
-    """ render hex string e.g. "0a40b1" as "0A 40 B1" """
-    return ' '.join(h[i:i + 2] for i in range(0, len(h), 2)).upper()
-
-
 def capability_rules(doc):
     """enumerate the rules in (namespace, name) order that are 'capability' rules (not lib/subscope/disposition/etc)."""
     for (_, _, rule) in sorted(map(lambda rule: (rule['meta'].get('namespace', ''), rule['meta']['name'], rule), doc.values())):
diff --git a/capa/render/vverbose.py b/capa/render/vverbose.py
index 2e5a7221..1882db69 100644
--- a/capa/render/vverbose.py
+++ b/capa/render/vverbose.py
@@ -41,16 +41,9 @@ def render_statement(ostream, match, statement, indent=0):
         # so, we have to inline some of the feature rendering here.
 
         child = statement['child']
-        if child['type'] in ('string', 'api', 'mnemonic', 'basic block', 'export', 'import', 'section', 'match', 'characteristic'):
-            value = rutils.bold2(child[child['type']])
-        elif child['type'] in ('number', 'offset'):
-            value = rutils.bold2(rutils.hex(child[child['type']]))
-        elif child['type'] == 'bytes':
-            value = rutils.bold2(rutils.hex_string(child[child['type']]))
-        else:
-            raise RuntimeError('unexpected feature type: ' + str(child))
+        value = rutils.bold2(child[child['type']])
 
-        if child['description']:
+        if child.get('description'):
             ostream.write('count(%s(%s = %s)): ' % (child['type'], value, child['description']))
         else:
             ostream.write('count(%s(%s)): ' % (child['type'], value))
@@ -81,22 +74,9 @@ def render_statement(ostream, match, statement, indent=0):
 def render_feature(ostream, match, feature, indent=0):
     ostream.write('  ' * indent)
 
-    if feature['type'] in ('string', 'api', 'mnemonic', 'basic block', 'export', 'import', 'section', 'match', 'characteristic'):
-        ostream.write(feature['type'])
-        ostream.write(': ')
-        ostream.write(rutils.bold2(feature[feature['type']]))
-    elif feature['type'] in ('number', 'offset'):
-        ostream.write(feature['type'])
-        ostream.write(': ')
-        ostream.write(rutils.bold2(rutils.hex(feature[feature['type']])))
-    elif feature['type'] == 'bytes':
-        ostream.write('bytes: ')
-        # bytes is the uppercase, hex-encoded string.
-        # it should always be an even number of characters (its hex).
-        ostream.write(rutils.bold2(rutils.hex_string(feature[feature['type']])))
-    # note that regex is found in `render_statement`
-    else:
-        raise RuntimeError('unexpected feature type: ' + str(feature))
+    ostream.write(feature['type'])
+    ostream.write(': ')
+    ostream.write(rutils.bold2(feature[feature['type']]))
 
     if 'description' in feature:
         ostream.write(' = ')
diff --git a/tests/test_main.py b/tests/test_main.py
index da1834f6..b79e4968 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -11,8 +11,11 @@ from fixtures import *
 
 
 def test_main(sample_9324d1a8ae37a36ae560c37448c9705a):
-    # tests rules can be loaded successfully
+    # tests rules can be loaded successfully and all output modes
+    assert capa.main.main([sample_9324d1a8ae37a36ae560c37448c9705a.path, '-vv']) == 0
     assert capa.main.main([sample_9324d1a8ae37a36ae560c37448c9705a.path, '-v']) == 0
+    assert capa.main.main([sample_9324d1a8ae37a36ae560c37448c9705a.path, '-j']) == 0
+    assert capa.main.main([sample_9324d1a8ae37a36ae560c37448c9705a.path]) == 0
 
 
 def test_main_single_rule(sample_9324d1a8ae37a36ae560c37448c9705a, tmpdir):
@@ -31,7 +34,10 @@ def test_main_single_rule(sample_9324d1a8ae37a36ae560c37448c9705a, tmpdir):
 
 
 def test_main_shellcode(sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32):
+    assert capa.main.main([sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32.path, '-vv', '-f', 'sc32']) == 0
     assert capa.main.main([sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32.path, '-v', '-f', 'sc32']) == 0
+    assert capa.main.main([sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32.path, '-j', '-f', 'sc32']) == 0
+    assert capa.main.main([sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32.path, '-f', 'sc32']) == 0
 
 
 def test_ruleset():