main.py: fix bugs for adding the cape extractor/format

This commit is contained in:
Yacine Elhamer
2023-06-22 13:55:50 +01:00
parent de2ba1ca94
commit 79ff76d124

View File

@@ -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])