From 79ff76d124dcaf57443d49bc38ecc3ee5a701c27 Mon Sep 17 00:00:00 2001 From: Yacine Elhamer Date: Thu, 22 Jun 2023 13:55:50 +0100 Subject: [PATCH] main.py: fix bugs for adding the cape extractor/format --- capa/main.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/capa/main.py b/capa/main.py index 7b7af961..0b6372a2 100644 --- a/capa/main.py +++ b/capa/main.py @@ -43,7 +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.cape.extractor import capa.features.extractors.common import capa.features.extractors.pefile import capa.features.extractors.dnfile_ @@ -525,7 +525,8 @@ def get_extractor( UnsupportedArchError UnsupportedOSError """ - if format_ not in (FORMAT_SC32, FORMAT_SC64): + + if format_ not in (FORMAT_SC32, FORMAT_SC64, FORMAT_CAPE): if not is_supported_format(path): raise UnsupportedFormatError() @@ -535,13 +536,13 @@ def get_extractor( if os_ == OS_AUTO and not is_supported_os(path): raise UnsupportedOSError() - elif format_ == FORMAT_CAPE: - import capa.features.extractors.cape + if format_ == FORMAT_CAPE: + import capa.features.extractors.cape.extractor import json with open(path, "r+", encoding="utf-8") as f: report = json.load(f) - return capa.features.extractors.cape.from_report(report) + return capa.features.extractors.cape.extractor.CapeExtractor.from_report(report) if format_ == FORMAT_DOTNET: import capa.features.extractors.dnfile.extractor @@ -613,7 +614,7 @@ def get_file_extractors(sample: str, format_: str) -> List[FeatureExtractor]: with open(sample, "r+", encoding="utf-8") as f: report = json.load(f) - file_extractors.append(capa.features.extractors.cape.from_report(report)) + file_extractors.append(capa.features.extractors.cape.extractor.CapeExtractor.from_report(report)) return file_extractors @@ -921,7 +922,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_CAPE, "CAPE sandbox report"), (FORMAT_FREEZE, "features previously frozen by capa"), ] format_help = ", ".join([f"{f[0]}: {f[1]}" for f in formats])