show-capabilities-by-function: use new ResultDocument

scripts/show-capabilities-by-function.py

@@ -53,6 +53,7 @@ import sys
 import logging
 import argparse
 import collections
+from typing import Dict
 
 import colorama
 
@@ -67,11 +68,12 @@ import capa.render.verbose
 import capa.features.freeze
 import capa.render.result_document as rd
 from capa.helpers import get_file_taste
+from capa.features.freeze import Address
 
 logger = logging.getLogger("capa.show-capabilities-by-function")
 
 
-def render_matches_by_function(doc):
+def render_matches_by_function(doc: rd.ResultDocument):
     """
     like:
 
@@ -90,37 +92,34 @@ def render_matches_by_function(doc):
           - send HTTP request
           - connect to HTTP server
     """
-    functions_by_bb = {}
-    functions_by_bb = {}
-    for finfo in doc["meta"]["analysis"]["layout"]["functions"]:
-        faddress = rd.deserialize_address(finfo["address"])
+    functions_by_bb: Dict[Address, Address] = {}
+    for finfo in doc.meta.analysis.layout.functions:
+        faddress = finfo.address
 
-        for bb in finfo["matched_basic_blocks"]:
-            bbaddress = rd.deserialize_address(bb["address"])
+        for bb in finfo.matched_basic_blocks:
+            bbaddress = bb.address
             functions_by_bb[bbaddress] = faddress
 
     ostream = rutils.StringIO()
 
     matches_by_function = collections.defaultdict(set)
     for rule in rutils.capability_rules(doc):
-        if rule["meta"]["scope"] == capa.rules.FUNCTION_SCOPE:
-            for addr, _ in rule["matches"]:
-                addr = rd.deserialize_address(addr)
-                matches_by_function[addr].add(rule["meta"]["name"])
-        elif rule["meta"]["scope"] == capa.rules.BASIC_BLOCK_SCOPE:
-            for addr, _ in rule["matches"]:
-                function = functions_by_bb[rd.deserialize_address(addr)]
-                matches_by_function[function].add(rule["meta"]["name"])
+        if rule.meta.scope == capa.rules.FUNCTION_SCOPE:
+            for addr, _ in rule.matches:
+                matches_by_function[addr].add(rule.meta.name)
+        elif rule.meta.scope == capa.rules.BASIC_BLOCK_SCOPE:
+            for addr, _ in rule.matches:
+                function = functions_by_bb[addr]
+                matches_by_function[function].add(rule.meta.name)
         else:
             # file scope
             pass
 
-    for f in doc["meta"]["analysis"]["feature_counts"]["functions"]:
-        addr = rd.deserialize_address(f["address"])
-        if not matches_by_function.get(addr, {}):
+    for f in doc.meta.analysis.feature_counts.functions:
+        if not matches_by_function.get(f.address, {}):
             continue
-        ostream.writeln("function at %s with %d features: " % (capa.render.verbose.format_address(addr), f["count"]))
-        for rule_name in sorted(matches_by_function[addr]):
+        ostream.writeln("function at %s with %d features: " % (capa.render.verbose.format_address(addr), f.count))
+        for rule_name in sorted(matches_by_function[f.address]):
             ostream.writeln("  - " + rule_name)
 
     return ostream.getvalue()
@@ -193,7 +192,7 @@ def main(argv=None):
     #  - when not an interactive session, and disable coloring
     # renderers should use coloring and assume it will be stripped out if necessary.
     colorama.init()
-    doc = rd.convert_capabilities_to_result_document(meta, rules, capabilities)
+    doc = rd.ResultDocument.from_capa(meta, rules, capabilities)
     print(render_matches_by_function(doc))
     colorama.deinit()