gitea-mirror/capa
1
0
Fork 0
mirror of https://github.com/mandiant/capa.git synced 2025-12-23 07:28:34 -08:00
Code Issues Packages Projects Releases Wiki Activity

FeatureExtractor alias: fix mypy typing issues by adding ininstance-based assert statements

Browse Source
This commit is contained in:
Yacine Elhamer
2023-06-26 22:46:27 +01:00
parent 63e4d3d5eb
commit b172f9a354
4 changed files with 15 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
scripts/profile-time.py
View File

@@ -46,7 +46,7 @@ import capa.helpers
import capa.features
import capa.features.common
import capa.features.freeze
from capa.features.extractors.base_extractor import StaticFeatureExtractor
from capa.features.extractors.base_extractor import FeatureExtractor, StaticFeatureExtractor
logger = logging.getLogger("capa.profile")
@@ -104,13 +104,14 @@ def main(argv=None):
args.format == capa.features.common.FORMAT_AUTO and capa.features.freeze.is_freeze(taste)
):
with open(args.sample, "rb") as f:
extractor = capa.features.freeze.load(f.read())
extractor: FeatureExtractor = capa.features.freeze.load(f.read())
assert isinstance(extractor, StaticFeatureExtractor)
else:
extractor = capa.main.get_extractor(
args.sample, args.format, args.os, capa.main.BACKEND_VIV, sig_paths, should_save_workspace=False
)
assert isinstance(extractor, StaticFeatureExtractor)
with tqdm.tqdm(total=args.number * args.repeat) as pbar:
def do_iteration():

4
scripts/show-capabilities-by-function.py
View File

@@ -70,7 +70,7 @@ import capa.render.result_document as rd
from capa.helpers import get_file_taste
from capa.features.common import FORMAT_AUTO
from capa.features.freeze import Address
from capa.features.extractors.base_extractor import StaticFeatureExtractor
from capa.features.extractors.base_extractor import FeatureExtractor, StaticFeatureExtractor
logger = logging.getLogger("capa.show-capabilities-by-function")
@@ -161,7 +161,7 @@ def main(argv=None):
if (args.format == "freeze") or (args.format == FORMAT_AUTO and capa.features.freeze.is_freeze(taste)):
format_ = "freeze"
with open(args.sample, "rb") as f:
extractor = capa.features.freeze.load(f.read())
extractor: FeatureExtractor = capa.features.freeze.load(f.read())
else:
format_ = args.format
should_save_workspace = os.environ.get("CAPA_SAVE_WORKSPACE") not in ("0", "no", "NO", "n", None)

8
scripts/show-features.py
View File

@@ -80,8 +80,8 @@ import capa.render.verbose as v
import capa.features.common
import capa.features.freeze
import capa.features.address
import capa.features.extractors.base_extractor
from capa.helpers import log_unsupported_runtime_error
from capa.features.extractors.base_extractor import FeatureExtractor, StaticFeatureExtractor
logger = logging.getLogger("capa.show-features")
@@ -117,14 +117,13 @@ def main(argv=None):
args.format == capa.features.common.FORMAT_AUTO and capa.features.freeze.is_freeze(taste)
):
with open(args.sample, "rb") as f:
extractor = capa.features.freeze.load(f.read())
extractor: FeatureExtractor = capa.features.freeze.load(f.read())
else:
should_save_workspace = os.environ.get("CAPA_SAVE_WORKSPACE") not in ("0", "no", "NO", "n", None)
try:
extractor = capa.main.get_extractor(
args.sample, args.format, args.os, args.backend, sig_paths, should_save_workspace
)
assert isinstance(extractor, capa.features.extractors.base_extractor.StaticFeatureExtractor)
except capa.exceptions.UnsupportedFormatError:
capa.helpers.log_unsupported_format_error()
return -1
@@ -132,6 +131,7 @@ def main(argv=None):
log_unsupported_runtime_error()
return -1
assert isinstance(extractor, StaticFeatureExtractor)
for feature, addr in extractor.extract_global_features():
print(f"global: {format_address(addr)}: {feature}")
@@ -190,7 +190,7 @@ def ida_main():
return 0
def print_features(functions, extractor: capa.features.extractors.base_extractor.FeatureExtractor):
def print_features(functions, extractor: StaticFeatureExtractor):
for f in functions:
if extractor.is_library_function(f.address):
function_name = extractor.get_function_name(f.address)