From de2ba1ca9430894d6d43bf816c3ee9e274798b17 Mon Sep 17 00:00:00 2001
From: Yacine Elhamer <elhamer.yacine@gmail.com>
Date: Thu, 22 Jun 2023 12:55:39 +0100
Subject: [PATCH] add the cape report format to main and across several other
 locations

---
 capa/features/common.py |  1 +
 capa/helpers.py         |  7 ++++++-
 capa/main.py            | 18 ++++++++++++++++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/capa/features/common.py b/capa/features/common.py
index 5060ebaa..d3c1aa32 100644
--- a/capa/features/common.py
+++ b/capa/features/common.py
@@ -449,6 +449,7 @@ VALID_FORMAT = (FORMAT_PE, FORMAT_ELF, FORMAT_DOTNET)
 FORMAT_AUTO = "auto"
 FORMAT_SC32 = "sc32"
 FORMAT_SC64 = "sc64"
+FORMAT_CAPE = "cape"
 FORMAT_FREEZE = "freeze"
 FORMAT_RESULT = "result"
 FORMAT_UNKNOWN = "unknown"
diff --git a/capa/helpers.py b/capa/helpers.py
index c03e0553..d06c6676 100644
--- a/capa/helpers.py
+++ b/capa/helpers.py
@@ -14,10 +14,11 @@ from typing import NoReturn
 import tqdm
 
 from capa.exceptions import UnsupportedFormatError
-from capa.features.common import FORMAT_PE, FORMAT_SC32, FORMAT_SC64, FORMAT_DOTNET, FORMAT_UNKNOWN, Format
+from capa.features.common import FORMAT_PE, FORMAT_SC32, FORMAT_SC64, FORMAT_CAPE, FORMAT_DOTNET, FORMAT_UNKNOWN, Format
 
 EXTENSIONS_SHELLCODE_32 = ("sc32", "raw32")
 EXTENSIONS_SHELLCODE_64 = ("sc64", "raw64")
+EXTENSIONS_CAPE = ("json", "json_")
 EXTENSIONS_ELF = "elf_"
 
 logger = logging.getLogger("capa")
@@ -57,6 +58,10 @@ def get_format_from_extension(sample: str) -> str:
         return FORMAT_SC32
     elif sample.endswith(EXTENSIONS_SHELLCODE_64):
         return FORMAT_SC64
+    elif sample.endswith(EXTENSIONS_CAPE):
+        # once we have support for more sandboxes that use json-formatted reports,
+        # we update this logic to ask the user to explicity specify the format
+        return FORMAT_CAPE
     return FORMAT_UNKNOWN
 
 
diff --git a/capa/main.py b/capa/main.py
index bdf0cec3..7b7af961 100644
--- a/capa/main.py
+++ b/capa/main.py
@@ -43,6 +43,7 @@ import capa.render.vverbose
 import capa.features.extractors
 import capa.render.result_document
 import capa.render.result_document as rdoc
+import capa.features.extractors.cape
 import capa.features.extractors.common
 import capa.features.extractors.pefile
 import capa.features.extractors.dnfile_
@@ -71,6 +72,7 @@ from capa.features.common import (
     FORMAT_AUTO,
     FORMAT_SC32,
     FORMAT_SC64,
+    FORMAT_CAPE,
     FORMAT_DOTNET,
     FORMAT_FREEZE,
     FORMAT_RESULT,
@@ -533,6 +535,14 @@ def get_extractor(
         if os_ == OS_AUTO and not is_supported_os(path):
             raise UnsupportedOSError()
 
+    elif format_ == FORMAT_CAPE:
+        import capa.features.extractors.cape
+        import json
+
+        with open(path, "r+", encoding="utf-8") as f:
+            report = json.load(f)
+        return capa.features.extractors.cape.from_report(report)
+
     if format_ == FORMAT_DOTNET:
         import capa.features.extractors.dnfile.extractor
 
@@ -598,6 +608,13 @@ def get_file_extractors(sample: str, format_: str) -> List[FeatureExtractor]:
     elif format_ == capa.features.extractors.common.FORMAT_ELF:
         file_extractors.append(capa.features.extractors.elffile.ElfFeatureExtractor(sample))
 
+    if format_ == FORMAT_CAPE:
+        import json
+
+        with open(sample, "r+", encoding="utf-8") as f:
+            report = json.load(f)
+        file_extractors.append(capa.features.extractors.cape.from_report(report))
+
     return file_extractors
 
 
@@ -904,6 +921,7 @@ def install_common_args(parser, wanted=None):
             (FORMAT_ELF, "Executable and Linkable Format"),
             (FORMAT_SC32, "32-bit shellcode"),
             (FORMAT_SC64, "64-bit shellcode"),
+            (FORMAT_CAPE, "CAPE sandbox report")
             (FORMAT_FREEZE, "features previously frozen by capa"),
         ]
         format_help = ", ".join([f"{f[0]}: {f[1]}" for f in formats])