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

Code before:
    for rid, nestedclass in iter_dotnet_table(pe, dnfile.mdtable.NestedClass.number):
        assert isinstance(nestedclass, dnfile.mdtable.NestedClassRow)
        nested_class_table[nestedclass.NestedClass.row_index] = nestedclass.EnclosingClass.row_index

    return nested_class_table

def get_typedef_class_table(pe):
    typedef_class_table = []
    
    for rid, typedef in iter_dotnet_table(pe, dnfile.mdtable.TypeDef.number):
        assert isinstance(typedef, dnfile.mdtable.TypeDefRow)
        typedef_class_table.append((typedef.TypeName, typedef.TypeNamespace))
        
    return typedef_class_table
    
def typedef_helper(index, nested_class_table, typedef_class_table, n, name):
    # Append the current typeref name
    n.append(name)
    
    if nested_class_table[index] in nested_class_table:
        typedef_helper(nested_class_table[index], nested_class_table, typedef_class_table, n, typedef_class_table[nested_class_table[index]-1][0])
    else:
        # Document the root enclosing details
        n.append(typedef_class_table[nested_class_table[index]-1][0])
        namespace = typedef_class_table[nested_class_table[index]-1][1]
        n.append(namespace)
        
    return tuple(n[::-1])

def get_typeref_table(pe):
    typeref_table = []
    
    # Used to track values in typeref table
    for rid, typeref in iter_dotnet_table(pe, dnfile.mdtable.TypeRef.number):
        assert isinstance(typeref, dnfile.mdtable.TypeRefRow)
        typeref_table.append((typeref.TypeName, typeref.TypeNamespace, typeref.ResolutionScope.row_index, type(typeref.ResolutionScope.table)))
        
    return typeref_table

def typeref_helper(index, typeref_table, n, name):
    # Append the current typeref name
    n.append(name)
    
    if typeref_table[index - 1][3] == dnfile.mdtable.TypeRef:
        # Recursively call helper function with enclosing typeref details
        typeref_helper(typeref_table[index - 1][2], typeref_table, n, typeref_table[index - 1][0])
    else:
        # Document the root enclosing details
        n.append(typeref_table[index - 1][0])
        namespace = typeref_table[index - 1][1]
        n.append(namespace)

Review comment (suggestion):
See my comment about indexing the `TypeDef` table directly.

Code after:
    try:
        table = pe.net.mdtables.tables.get(table_index, [])
        return table[row_index - 1]
    except IndexError:
        return None


def resolve_nested_typedef_name(
    nested_class_table: dict, index: int, typedef: dnfile.mdtable.TypeDefRow, pe: dnfile.dnPE
) -> Tuple[str, Tuple[str, ...]]:
    """Resolves all nested TypeDef class names. Returns the namespace as a str and the nested TypeRef name as a tuple"""

    if index in nested_class_table:
        typedef_name = []
        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