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

Code before:
    # 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)
        
    return tuple(n[::-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 = get_typedef_class_table(pe)
    
    for rid, typedef in iter_dotnet_table(pe, dnfile.mdtable.TypeDef.number):
        assert isinstance(typedef, dnfile.mdtable.TypeDefRow)

        typedef_name = []
        if rid in nested_class_table:
            typedef.TypeName = 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, typedef.TypeName, namespace=typedef.TypeNamespace)


    typeref_table = get_typeref_table(pe)
    
    for rid, typeref in iter_dotnet_table(pe, dnfile.mdtable.TypeRef.number):
        assert isinstance(typeref, dnfile.mdtable.TypeRefRow)
        
        # If the ResolutionScope decodes to a typeRef type, then it is nested
        n = []
        if type(typeref.ResolutionScope.table) == dnfile.mdtable.TypeRef:
            typeref.TypeName = typeref_helper(typeref.ResolutionScope.row_index, typeref_table, n, typeref.TypeName)
              
        typeref_token: int = calculate_dotnet_token_value(dnfile.mdtable.TypeRef.number, rid)
        yield DnType(typeref_token, typeref.TypeName, namespace=typeref.TypeNamespace)
        

Review comment (refactor):
Can we instead index the `TypeDef` table directly e.g. `table[rid-1]` or [`table.get_with_row_index(rid)`](https://github.com/malwarefrank/dnfile/blob/b837ad6381a33f85ee0963cfa50f4ad91609b319/src/dnfile/base.py#L875)? I'm suggesting using direct indexing in places like https://github.com/mandiant/capa/blob/94a8b65b480fa0a5944638b0ce3c3dbfd3737bca/capa/features/extractors/dnfile/helpers.py#L345-L348

This avoids iterating the `TypeDef` table more than is needed and makes the indexing more intuitive.




Code after:
            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])

        # Document the root enclosing details
        typeref_name.append(table_row.TypeName)

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

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


def get_dotnet_nested_class_table_index(pe: dnfile.dnPE) -> Dict[int, int]:
    """Build index for EnclosingClass based off the NestedClass row index in the nestedclass table"""
    nested_class_table = {}

    # Used to find nested classes in typedef
    for _, 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