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

Code before:

    return nested_class_table

def typedef_helper(index, nested_class_table, typedef_class_table, typedef_name, name):
    # Append the current typeref name
    typedef_name.append(name)
    
    while nested_class_table[index] in nested_class_table:
        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 get_dotnet_types(pe: dnfile.dnPE) -> Iterator[DnType]:
    """get .NET types from TypeDef and TypeRef 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):
        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)
        
        typedef_token: int = calculate_dotnet_token_value(dnfile.mdtable.TypeDef.number, rid)
        yield DnType(typedef_token, typedefname, namespace=typedefnamespace)


Review comment (refactor):
This assumes that `index` is always valid - we've encountered heavily obfuscated .NET files where this wasn't necessarily the case. To avoid `IndexError` exceptions let's create a helper function called `get_dotnet_table_row` that:

1. Takes the same first two arguments as `iter_dotnet_table` plus a third argument specifying the (1-based) row index
2. Resolves the table
3. Verifies that the row index is valid for the table
4. Returns the `dnfile.base.MDTableRow` object associated with the supplied row index if the table and row index are valid, otherwise return `None`


Code after:
        name = typedef.TypeName

        # Append the current typedef name
        typedef_name.append(name)

        while nested_class_table[index] in nested_class_table:
            # Iterate through the typedef table to resolve the nested name
            table_row = get_dotnet_table_row(pe, dnfile.mdtable.TypeDef.number, nested_class_table[index])
            if table_row is None:
                return typedef.TypeNamespace, tuple(typedef_name[::-1])

            name = table_row.TypeName
            typedef_name.append(name)
            index = nested_class_table[index]

        # Document the root enclosing details
        table_row = get_dotnet_table_row(pe, dnfile.mdtable.TypeDef.number, nested_class_table[index])
        if table_row is None:
            return typedef.TypeNamespace, tuple(typedef_name[::-1])

        enclosing_name = table_row.TypeName
        typedef_name.append(enclosing_name)

        return table_row.TypeNamespace, tuple(typedef_name[::-1])

    else:
        return typedef.TypeNamespace, (typedef.TypeName,)


def resolve_nested_typeref_name(
    index: int, typeref: dnfile.mdtable.TypeRefRow, pe: dnfile.dnPE
) -> Tuple[str, Tuple[str, ...]]:
    """Resolves all nested TypeRef class names. Returns the namespace as a str and the nested TypeRef name as a tuple"""
    # If the ResolutionScope decodes to a typeRef type then it is nested
    if isinstance(typeref.ResolutionScope.table, dnfile.mdtable.TypeRef):
        typeref_name = []
        name = typeref.TypeName
        # Not appending the current typeref name to avoid potential duplicate

        # Validate index
        table_row = get_dotnet_table_row(pe, dnfile.mdtable.TypeRef.number, index)
        if table_row is None:
            return typeref.TypeNamespace, (typeref.TypeName,)

        while isinstance(table_row.ResolutionScope.table, dnfile.mdtable.TypeRef):
            # Iterate through the typeref table to resolve the nested name
            typeref_name.append(name)
            name = table_row.TypeName
            table_row = get_dotnet_table_row(pe, dnfile.mdtable.TypeRef.number, table_row.ResolutionScope.row_index)
            if table_row is None:
                return typeref.TypeNamespace, tuple(typeref_name[::-1])