From 55720ddbfd067f2920fa7a31c1e5c8b0d3cec105 Mon Sep 17 00:00:00 2001
From: mr-tz <moritz.raabe@mandiant.com>
Date: Thu, 12 Dec 2024 09:43:45 +0000
Subject: [PATCH] make more fields optional for more flexible model

tmp
---
 capa/features/extractors/vmray/__init__.py | 9 +++++----
 capa/features/extractors/vmray/models.py   | 6 +++---
 tests/test_vmray_model.py                  | 4 ++--
 3 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/capa/features/extractors/vmray/__init__.py b/capa/features/extractors/vmray/__init__.py
index af59a3208..71d9afd10 100644
--- 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
index 027680586..36cd261e3 100644
--- 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
index c693b6631..58d8a9ccc 100644
--- 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():