Merge pull request #1060 from mandiant/feat/frz-extractor-info

Feat/frz extractor info

capa/features/extractors/null.py

@@ -2,7 +2,7 @@ from typing import Dict, List, Tuple
 from dataclasses import dataclass
 
 from capa.features.common import Feature
-from capa.features.address import Address
+from capa.features.address import NO_ADDRESS, Address
 from capa.features.extractors.base_extractor import BBHandle, InsnHandle, FunctionHandle, FeatureExtractor
 
 
@@ -40,8 +40,8 @@ class NullFeatureExtractor(FeatureExtractor):
         return self.base_address
 
     def extract_global_features(self):
-        for address, feature in self.global_features:
-            yield feature, address
+        for feature in self.global_features:
+            yield feature, NO_ADDRESS
 
     def extract_file_features(self):
         for address, feature in self.file_features:

capa/features/freeze/__init__.py

@@ -18,6 +18,7 @@ import dncil.clr.token
 from pydantic import Field, BaseModel
 
 import capa.helpers
+import capa.version
 import capa.features.file
 import capa.features.insn
 import capa.features.common
@@ -194,9 +195,18 @@ class Features(BaseModel):
         allow_population_by_field_name = True
 
 
+class Extractor(BaseModel):
+    name: str
+    version: str = capa.version.__version__
+
+    class Config:
+        allow_population_by_field_name = True
+
+
 class Freeze(BaseModel):
     version: int = 2
     base_address: Address = Field(alias="base address")
+    extractor: Extractor
     features: Features
 
     class Config:
@@ -293,6 +303,7 @@ def dumps(extractor: capa.features.extractors.base_extractor.FeatureExtractor) -
     freeze = Freeze(
         version=2,
         base_address=Address.from_capa(extractor.get_base_address()),
+        extractor=Extractor(name=extractor.__class__.__name__),
         features=features,
     )
 

tests/test_freeze.py

@@ -156,13 +156,21 @@ def test_freeze_sample(tmpdir, z9324d_extractor):
     assert capa.features.freeze.main([path, o, "-v"]) == 0
 
 
-def test_freeze_load_sample(tmpdir, z9324d_extractor):
+@pytest.mark.parametrize(
+    "extractor",
+    [
+        pytest.param("z9324d_extractor"),
+    ],
+)
+def test_freeze_load_sample(tmpdir, request, extractor):
     o = tmpdir.mkdir("capa").join("test.frz")
 
+    extractor = request.getfixturevalue(extractor)
+
     with open(o.strpath, "wb") as f:
-        f.write(capa.features.freeze.dump(z9324d_extractor))
+        f.write(capa.features.freeze.dump(extractor))
 
     with open(o.strpath, "rb") as f:
         null_extractor = capa.features.freeze.load(f.read())
 
-    compare_extractors(z9324d_extractor, null_extractor)
+    compare_extractors(extractor, null_extractor)