Repository: mandiant/capa
File: capa/features/extractors/vmray/__init__.py (Python)

Code before:

    def _compute_base_address(self):
        assert self.sample_file_static_data is not None
        if self.sample_file_static_data.pe:
            self.base_address = self.sample_file_static_data.pe.basic_info.image_base

    def _compute_exports(self):
        assert self.sample_file_static_data is not None
        if self.sample_file_static_data.pe:
            for export in self.sample_file_static_data.pe.exports:
                self.exports[export.address] = export.api.name

    def _compute_imports(self):
        assert self.sample_file_static_data is not None
        if self.sample_file_static_data.pe:
            for module in self.sample_file_static_data.pe.imports:
                for api in module.apis:
                    self.imports[api.address] = (module.dll, api.api.name)

    def _compute_sections(self):
        assert self.sample_file_static_data is not None
        if self.sample_file_static_data.pe:
            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:
            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():
            # we expect monitor IDs to be unique
            assert process.monitor_id not in self.monitor_processes

            ppid: int = (
                self.sv2.processes[process.ref_parent_process.path[1]].os_pid if process.ref_parent_process else 0
            )
            self.monitor_processes[process.monitor_id] = VMRayMonitorProcess(
                process.os_pid,
                ppid,
                process.monitor_id,
                process.origin_monitor_id,
                process.image_name,
                process.filename,
                process.cmd_line,
            )

        # not all processes are recorded in SummaryV2.json, get missing data from flog.xml, see #2394
        for monitor_process in self.flog.analysis.monitor_processes:
            vmray_monitor_process: VMRayMonitorProcess = VMRayMonitorProcess(
                monitor_process.os_pid,

Review comment (refactor):
Good catch. Let's set the default to `[]` for `sections` in the [`ElfFile` model](https://github.com/mandiant/capa/blob/1a82b9d0c56cc8b68af213739c4fc8c127d24900/capa/features/extractors/vmray/models.py#L277-L279) so we don't need checks like this throughout the code.

Code after:

    def _compute_base_address(self):
        assert self.sample_file_static_data is not None
        if self.sample_file_static_data.pe:
            self.base_address = self.sample_file_static_data.pe.basic_info.image_base

    def _compute_exports(self):
        assert self.sample_file_static_data is not None
        if self.sample_file_static_data.pe:
            for export in self.sample_file_static_data.pe.exports:
                self.exports[export.address] = export.api.name

    def _compute_imports(self):
        assert self.sample_file_static_data is not None
        if self.sample_file_static_data.pe:
            for module in self.sample_file_static_data.pe.imports:
                for api in module.apis:
                    self.imports[api.address] = (module.dll, api.api.name)

    def _compute_sections(self):
        assert self.sample_file_static_data is not None
        if self.sample_file_static_data.pe:
            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

    def _compute_monitor_processes(self):
        for process in self.sv2.processes.values():
            # we expect monitor IDs to be unique
            assert process.monitor_id not in self.monitor_processes

            ppid: int = (
                self.sv2.processes[process.ref_parent_process.path[1]].os_pid if process.ref_parent_process else 0
            )
            self.monitor_processes[process.monitor_id] = VMRayMonitorProcess(
                process.os_pid,
                ppid,
                process.monitor_id,
                process.origin_monitor_id,
                process.image_name,
                process.filename,
                process.cmd_line,
            )

        # not all processes are recorded in SummaryV2.json, get missing data from flog.xml, see #2394
        for monitor_process in self.flog.analysis.monitor_processes:
            vmray_monitor_process: VMRayMonitorProcess = VMRayMonitorProcess(
                monitor_process.os_pid,
                monitor_process.os_parent_pid,