vmray: clean up pydantic models and implement base address extraction

capa/features/extractors/vmray/__init__.py

@@ -17,12 +17,14 @@ class VMRayAnalysis:
         self.exports: Dict[int, str] = {}
         self.imports: Dict[int, str] = {}
         self.sections: Dict[int, str] = {}
+        self.base_address: int
 
         self.sample_file_name: str
         self.sample_file_analysis: File
         self.sample_file_static_data: Optional[StaticData]
 
         self._find_sample_file()
+        self._compute_base_address()
         self._compute_exports()
         self._compute_sections()
 
@@ -38,6 +40,10 @@ class VMRayAnalysis:
 
                 break
 
+    def _compute_base_address(self):
+        if self.sample_file_static_data and self.sample_file_static_data.pe:
+            self.base_address = self.sample_file_static_data.pe.basic_info.image_base
+
     def _compute_exports(self):
         if self.sample_file_static_data and self.sample_file_static_data.pe:
             for export in self.sample_file_static_data.pe.exports:

capa/features/extractors/vmray/extractor.py

@@ -14,7 +14,7 @@ from zipfile import ZipFile
 import capa.helpers
 import capa.features.extractors.vmray.file
 from capa.features.common import Feature
-from capa.features.address import Address
+from capa.features.address import Address, AbsoluteVirtualAddress
 from capa.features.extractors.vmray import VMRayAnalysis
 from capa.features.extractors.vmray.models import Analysis, SummaryV2
 from capa.features.extractors.base_extractor import DynamicFeatureExtractor
@@ -38,6 +38,10 @@ class VMRayExtractor(DynamicFeatureExtractor):
 
         return cls(VMRayAnalysis(sv2, flog))
 
+    def get_base_address(self) -> Address:
+        # value according to the PE header, the actual trace may use a different imagebase
+        return AbsoluteVirtualAddress(self.analysis.base_address)
+
     def extract_file_features(self) -> Iterator[Tuple[Feature, Address]]:
         yield from capa.features.extractors.vmray.file.extract_features(self.analysis)
 
@@ -50,3 +54,5 @@ if __name__ == "__main__":
     extractor = VMRayExtractor.from_archive(input_path)
     for feat, addr in extractor.extract_file_features():
         print(f"{feat} -> {addr}")
+
+    print(f"base address: {hex(extractor.get_base_address())}")

capa/features/extractors/vmray/models.py

@@ -113,7 +113,7 @@ class PEFileSection(BaseModel):
 
 class PEFile(BaseModel):
     _type: str
-    basic_info: Optional[PEFileBasicInfo] = None
+    basic_info: PEFileBasicInfo
     exports: List[PEFileExport] = []
     imports: List[PEFileImportModule] = []
     sections: List[PEFileSection] = []