make more fields optional for more flexible model

tmp
mandiant · Dec 12, 2024 · 55720dd · 55720dd
1 parent 893378c
commit 55720dd
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 9 deletions.
diff --git a/capa/features/extractors/vmray/__init__.py b/capa/features/extractors/vmray/__init__.py
@@ -36,8 +36,8 @@ class VMRayMonitorProcess:
     monitor_id: int  # unique ID assigned to process by VMRay
     origin_monitor_id: int  # unique VMRay ID of parent process
     image_name: str
-    filename: str
-    cmd_line: str
+    filename: Optional[str] = ""
+    cmd_line: Optional[str] = ""
 
 
 class VMRayAnalysis:
@@ -151,8 +151,9 @@ def _compute_sections(self):
             for pefile_section in self.sample_file_static_data.pe.sections:
                 self.sections[pefile_section.virtual_address] = pefile_section.name
         elif self.sample_file_static_data.elf:
-            for elffile_section in self.sample_file_static_data.elf.sections:
-                self.sections[elffile_section.header.sh_addr] = elffile_section.header.sh_name
+            if self.sample_file_static_data.elf.sections:
+                for elffile_section in self.sample_file_static_data.elf.sections:
+                    self.sections[elffile_section.header.sh_addr] = elffile_section.header.sh_name
 
     def _compute_monitor_processes(self):
         for process in self.sv2.processes.values():

diff --git a/capa/features/extractors/vmray/models.py b/capa/features/extractors/vmray/models.py
@@ -276,7 +276,7 @@ class ElfFileHeader(BaseModel):
 
 class ElfFile(BaseModel):
     # file_header: ElfFileHeader
-    sections: list[ElfFileSection]
+    sections: Optional[list[ElfFileSection]] = None
 
 
 class StaticData(BaseModel):
@@ -316,9 +316,9 @@ class Process(BaseModel):
     # monitor_reason: str
     origin_monitor_id: int  # VMRay ID of parent process
     os_pid: int
-    filename: SanitizedString
+    filename: Optional[SanitizedString] = ""
     image_name: str
-    cmd_line: SanitizedString
+    cmd_line: Optional[SanitizedString] = ""
     ref_parent_process: Optional[GenericReference] = None
 
 

diff --git a/tests/test_vmray_model.py b/tests/test_vmray_model.py
@@ -103,8 +103,8 @@ def test_vmray_model_elffile():
         """
     )
 
-    assert elffile.sections[0].header.sh_name == "abcd1234"
-    assert elffile.sections[0].header.sh_addr == 2863311530
+    assert elffile.sections and elffile.sections[0].header.sh_name == "abcd1234"
+    assert elffile.sections and elffile.sections[0].header.sh_addr == 2863311530
 
 
 def test_vmray_model_pefile():