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

Code before:

def extract_file_embedded_pe(buf, **kwargs) -> Iterator[Tuple[Feature, Address]]:
    for offset, _ in pe_carve.carve(buf, 1):
        yield Characteristic("embedded pe"), FileOffsetAddress(offset)


def extract_file_export_names(vw, **kwargs) -> Iterator[Tuple[Feature, Address]]:
    for va, _, name, _ in vw.getExports():
        yield Export(name), AbsoluteVirtualAddress(va)


def extract_file_import_names(vw, **kwargs) -> Iterator[Tuple[Feature, Address]]:
    """
    extract imported function names
    1. imports by ordinal:
     - modulename.#ordinal
    2. imports by name, results in two features to support importname-only matching:
     - modulename.importname
     - importname
    """
    for va, _, _, tinfo in vw.getImports():
        # vivisect source: tinfo = "%s.%s" % (libname, impname)
        modname, impname = tinfo.split(".", 1)
        if is_viv_ord_impname(impname):
            # replace ord prefix with #
            impname = f"#{impname[len('ord') :]}"

        addr = AbsoluteVirtualAddress(va)
        for name in capa.features.extractors.helpers.generate_symbols(modname, impname):
            yield Import(name), addr


def is_viv_ord_impname(impname: str) -> bool:
    """
    return if import name matches vivisect's ordinal naming scheme `'ord%d' % ord`
    """
    if not impname.startswith("ord"):
        return False
    try:
        int(impname[len("ord") :])
    except ValueError:
        return False
    else:
        return True


def extract_file_section_names(vw, **kwargs) -> Iterator[Tuple[Feature, Address]]:
    for va, _, segname, _ in vw.getSegments():
        yield Section(segname), AbsoluteVirtualAddress(va)



Review comment (suggestion):
```suggestion
            impname = "#" + impname[len('ord'):]
```

Code after:

def extract_file_embedded_pe(buf, **kwargs) -> Iterator[Tuple[Feature, Address]]:
    for offset, _ in pe_carve.carve(buf, 1):
        yield Characteristic("embedded pe"), FileOffsetAddress(offset)


def extract_file_export_names(vw, **kwargs) -> Iterator[Tuple[Feature, Address]]:
    for va, _, name, _ in vw.getExports():
        yield Export(name), AbsoluteVirtualAddress(va)


def extract_file_import_names(vw, **kwargs) -> Iterator[Tuple[Feature, Address]]:
    """
    extract imported function names
    1. imports by ordinal:
     - modulename.#ordinal
    2. imports by name, results in two features to support importname-only matching:
     - modulename.importname
     - importname
    """
    for va, _, _, tinfo in vw.getImports():
        # vivisect source: tinfo = "%s.%s" % (libname, impname)
        modname, impname = tinfo.split(".", 1)
        if is_viv_ord_impname(impname):
            # replace ord prefix with #
            impname = "#" + impname[len("ord") :]

        addr = AbsoluteVirtualAddress(va)
        for name in capa.features.extractors.helpers.generate_symbols(modname, impname):
            yield Import(name), addr


def is_viv_ord_impname(impname: str) -> bool:
    """
    return if import name matches vivisect's ordinal naming scheme `'ord%d' % ord`
    """
    if not impname.startswith("ord"):
        return False
    try:
        int(impname[len("ord") :])
    except ValueError:
        return False
    else:
        return True


def extract_file_section_names(vw, **kwargs) -> Iterator[Tuple[Feature, Address]]:
    for va, _, segname, _ in vw.getSegments():
        yield Section(segname), AbsoluteVirtualAddress(va)