Repository: mandiant/capa
File: capa/features/extractors/ghidra/file.py (Python)

Code before:
# Copyright (C) 2020 Mandiant, Inc. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
# You may obtain a copy of the License at: [package root]/LICENSE.txt
# Unless required by applicable law or agreed to in writing, software distributed under the License
#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.

import struct
from typing import Tuple, Iterator

import ghidra
from ghidra.program.flatapi import FlatProgramAPI
from ghidra.program.model.symbol import SourceType, SymbolType

import capa.features.extractors.common
import capa.features.extractors.helpers
import capa.features.extractors.strings
import capa.features.extractors.ghidra.helpers
from capa.features.file import Export, Import, Section, FunctionName
from capa.features.common import FORMAT_PE, FORMAT_ELF, Format, String, Feature, Characteristic
from capa.features.address import NO_ADDRESS, Address, FileOffsetAddress, AbsoluteVirtualAddress

currentProgram: ghidra.program.database.ProgramDB
flatapi = FlatProgramAPI(currentProgram)

MAX_OFFSET_PE_AFTER_MZ = 0x200


def check_segment_for_pe() -> Iterator[Tuple[int, int]]:
    """check segment for embedded PE

    adapted for Ghidra from:
    https://github.com/vivisect/vivisect/blob/91e8419a861f4977https://github.com/vivisect/vivisect/blob/91e8419a861f49779f18316f155311967e696836/PE/carve.py#L259f18316f155311967e696836/PE/carve.py#L25
    """

    mz_xor = [
        (
            capa.features.extractors.helpers.xor_static(b"MZ", i),
            capa.features.extractors.helpers.xor_static(b"PE", i),
            i,
        )
        for i in range(256)
    ]

    todo = []
    for mzx, pex, i in mz_xor:
        # find all segment offsets containing XOR'd "MZ" bytes
        off: ghidra.program.model.address.GenericAddress
        for off in capa.features.extractors.ghidra.helpers.find_byte_sequence(mzx):

Review comment (suggestion):
I dug and I did not find a solution for mypy complaining about undefined names. Let’s be explicit about the names Ghidrathon has injected by adding the comment # type: ignore [name-defined] to each line that uses a name injected by Ghidrathon. This solution is unfortunately far more verbose but helps us avoid creating a new object just to make mypy happy.

Code after:
# Copyright (C) 2020 Mandiant, Inc. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
# You may obtain a copy of the License at: [package root]/LICENSE.txt
# Unless required by applicable law or agreed to in writing, software distributed under the License
#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.

import struct
from typing import Tuple, Iterator

import ghidra
from ghidra.program.model.symbol import SourceType, SymbolType

import capa.features.extractors.common
import capa.features.extractors.helpers
import capa.features.extractors.strings
import capa.features.extractors.ghidra.helpers
from capa.features.file import Export, Import, Section, FunctionName
from capa.features.common import FORMAT_PE, FORMAT_ELF, Format, String, Feature, Characteristic
from capa.features.address import NO_ADDRESS, Address, FileOffsetAddress, AbsoluteVirtualAddress

currentProgram: ghidra.program.database.ProgramDB

MAX_OFFSET_PE_AFTER_MZ = 0x200


def check_segment_for_pe() -> Iterator[Tuple[int, int]]:
    """check segment for embedded PE

    adapted for Ghidra from:
    https://github.com/vivisect/vivisect/blob/91e8419a861f4977https://github.com/vivisect/vivisect/blob/91e8419a861f49779f18316f155311967e696836/PE/carve.py#L259f18316f155311967e696836/PE/carve.py#L25
    """

    mz_xor = [
        (
            capa.features.extractors.helpers.xor_static(b"MZ", i),
            capa.features.extractors.helpers.xor_static(b"PE", i),
            i,
        )
        for i in range(256)
    ]

    todo = []
    for mzx, pex, i in mz_xor:
        # find all segment offsets containing XOR'd "MZ" bytes
        off: ghidra.program.model.address.GenericAddress
        for off in capa.features.extractors.ghidra.helpers.find_byte_sequence(mzx):
            todo.append((off, mzx, pex, i))