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

Code before:
        name = typedef_class_table[nested_class_table[index]-1].TypeName
        typedef_name.append(name)
        index = nested_class_table[index]
        
    # Document the root enclosing details
    enclosing_name = typedef_class_table[nested_class_table[index]-1].TypeName
    typedef_name.append(enclosing_name)
    namespace = typedef_class_table[nested_class_table[index]-1].TypeNamespace
    
    return namespace, tuple(typedef_name[::-1])

def typeref_helper(index, typeref_table, typeref_name, name):
    # Not appending the current typeref name to avoid potential duplicate
    
    while type(typeref_table[index - 1].ResolutionScope.table) is dnfile.mdtable.TypeRef:
        # Recursively call helper function with enclosing typeref details
        typeref_name.append(name)
        name = typeref_table[index - 1].TypeName
        index = typeref_table[index - 1].ResolutionScope.row_index
    
    # Document the root enclosing details
    typeref_name.append(typeref_table[index - 1].TypeName)
    namespace = typeref_table[index - 1].TypeNamespace
     
    return namespace, tuple(typeref_name[::-1])
    
def extract_file_class_features(pe: dnfile.dnPE, **kwargs) -> Iterator[Tuple[Class, Address]]:
    """emit class features from TypeRef and TypeDef tables"""
    nested_class_table = get_nested_class_table(pe)
    typedef_class_table = pe.net.mdtables.tables.get(dnfile.mdtable.TypeDef.number, [])
    
    for rid, typedef in iter_dotnet_table(pe, dnfile.mdtable.TypeDef.number):
        # emit internal .NET classes
        assert isinstance(typedef, dnfile.mdtable.TypeDefRow)

        typedef_name = []
        typedefname = (typedef.TypeName,)
        typedefnamespace = typedef.TypeNamespace
        if rid in nested_class_table:
            typedefnamespace, typedefname = typedef_helper(rid, nested_class_table, typedef_class_table, typedef_name, typedef.TypeName)

        token = calculate_dotnet_token_value(dnfile.mdtable.TypeDef.number, rid)
        yield Class(DnType.format_name(typedefname, namespace=typedefnamespace)), DNTokenAddress(token)

    typeref_table = pe.net.mdtables.tables.get(dnfile.mdtable.TypeRef.number, [])

    for rid, typeref in iter_dotnet_table(pe, dnfile.mdtable.TypeRef.number):
        # emit external .NET classes
        assert isinstance(typeref, dnfile.mdtable.TypeRefRow)

        # If the ResolutionScope decodes to a typeRef type then it is nested

Review comment (refactor):
These appear to be the same implementations found in `capa.features.extractors.dnfile.helpers`? If so, let's import these from `capa.features.extractors.dnfile.helpers` instead of redefining them.

Code after:

    for rid, typeref in iter_dotnet_table(pe, dnfile.mdtable.TypeRef.number):
        # emit external .NET classes
        assert isinstance(typeref, dnfile.mdtable.TypeRefRow)

        typerefnamespace, typerefname = resolve_nested_typeref_name(typeref.ResolutionScope.row_index, typeref, pe)

        token = calculate_dotnet_token_value(dnfile.mdtable.TypeRef.number, rid)
        yield Class(DnType.format_name(typerefname, namespace=typerefnamespace)), DNTokenAddress(token)


def extract_file_os(**kwargs) -> Iterator[Tuple[OS, Address]]:
    yield OS(OS_ANY), NO_ADDRESS


def extract_file_arch(pe: dnfile.dnPE, **kwargs) -> Iterator[Tuple[Arch, Address]]:
    # to distinguish in more detail, see https://stackoverflow.com/a/23614024/10548020
    # .NET 4.5 added option: any CPU, 32-bit preferred
    assert pe.net is not None
    assert pe.net.Flags is not None

    if pe.net.Flags.CLR_32BITREQUIRED and pe.PE_TYPE == pefile.OPTIONAL_HEADER_MAGIC_PE:
        yield Arch(ARCH_I386), NO_ADDRESS
    elif not pe.net.Flags.CLR_32BITREQUIRED and pe.PE_TYPE == pefile.OPTIONAL_HEADER_MAGIC_PE_PLUS:
        yield Arch(ARCH_AMD64), NO_ADDRESS
    else:
        yield Arch(ARCH_ANY), NO_ADDRESS


def extract_file_strings(pe: dnfile.dnPE, **kwargs) -> Iterator[Tuple[String, Address]]:
    yield from capa.features.extractors.common.extract_file_strings(pe.__data__)


def extract_file_mixed_mode_characteristic_features(
    pe: dnfile.dnPE, **kwargs
) -> Iterator[Tuple[Characteristic, Address]]:
    if is_dotnet_mixed_mode(pe):
        yield Characteristic("mixed mode"), NO_ADDRESS


def extract_file_features(pe: dnfile.dnPE) -> Iterator[Tuple[Feature, Address]]:
    for file_handler in FILE_HANDLERS:
        for feature, addr in file_handler(pe=pe):  # type: ignore
            yield feature, addr


FILE_HANDLERS = (
    extract_file_import_names,
    extract_file_function_names,
    extract_file_strings,
    extract_file_format,