From 19a5ef8a644327f2fe436a73ccea390adaf95c6e Mon Sep 17 00:00:00 2001
From: Willi Ballenthin <willi.ballenthin@gmail.com>
Date: Wed, 5 Jul 2023 18:21:03 +0200
Subject: [PATCH] import-to-ida: use existing result document json parser

---
 scripts/import-to-ida.py | 34 +++++++++++++++++-----------------
 1 file changed, 17 insertions(+), 17 deletions(-)

diff --git a/scripts/import-to-ida.py b/scripts/import-to-ida.py
index fd762f58..8b9b3471 100644
--- a/scripts/import-to-ida.py
+++ b/scripts/import-to-ida.py
@@ -29,13 +29,16 @@ Unless required by applicable law or agreed to in writing, software distributed
 See the License for the specific language governing permissions and limitations under the License.
 """
 import binascii
-import json
 import logging
 
 import ida_nalt
 import ida_funcs
 import ida_kernwin
 
+import capa.rules
+import capa.features.freeze
+import capa.render.result_document
+
 logger = logging.getLogger("capa")
 
 
@@ -65,40 +68,37 @@ def main():
     if not path:
         return 0
 
-    with open(path, "rb") as f:
-        doc = json.loads(f.read().decode("utf-8"))
-
-    if "meta" not in doc or "rules" not in doc:
-        logger.error("doesn't appear to be a capa report")
-        return -1
+    result_doc = capa.render.result_document.ResultDocument.parse_file(path)
+    meta, capabilities = result_doc.to_capa()
 
     # in IDA 7.4, the MD5 hash may be truncated, for example:
     # wanted: 84882c9d43e23d63b82004fae74ebb61
     # found: b'84882C9D43E23D63B82004FAE74EBB6\x00'
     #
     # see: https://github.com/idapython/bin/issues/11
-    a = doc["meta"]["sample"]["md5"].lower()
+    a = meta["sample"]["md5"].lower()
     b = binascii.hexlify(ida_nalt.retrieve_input_file_md5()).decode("ascii").lower()
     if not a.startswith(b):
         logger.error("sample mismatch")
         return -2
 
     rows = []
-    for rule in doc["rules"].values():
-        if rule["meta"].get("lib"):
+    for name in capabilities.keys():
+        rule = result_doc.rules[name]
+        if rule.meta.lib:
             continue
-        if rule["meta"].get("capa/subscope"):
+        if rule.meta.is_subscope_rule:
             continue
-        if rule["meta"]["scope"] != "function":
+        if rule.meta.scope != capa.rules.Scope.FUNCTION:
             continue
 
-        name = rule["meta"]["name"]
-        ns = rule["meta"].get("namespace", "")
-        for address, match in rule["matches"]:
-            if address["type"] != "absolute":
+        ns = rule.meta.namespace
+
+        for address, _ in rule.matches:
+            if address.type != capa.features.freeze.AddressType.ABSOLUTE:
                 continue
 
-            va = address["value"]
+            va = address.value
             rows.append((ns, name, va))
 
     # order by (namespace, name) so that like things show up together