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

Code before:
            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 enclosing_and_nested_classes_index_table(pe: dnfile.dnPE):
    """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 _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_dotnet_types(pe: dnfile.dnPE) -> Iterator[DnType]:
    """get .NET types from TypeDef and TypeRef tables"""
    nested_class_table = enclosing_and_nested_classes_index_table(pe)

    for rid, typedef in iter_dotnet_table(pe, dnfile.mdtable.TypeDef.number):
        assert isinstance(typedef, dnfile.mdtable.TypeDefRow)

        typedefnamespace, typedefname = resolve_nested_typedef_name(nested_class_table, rid, typedef, pe)

        typedef_token: int = calculate_dotnet_token_value(dnfile.mdtable.TypeDef.number, rid)
        yield DnType(typedef_token, typedefname, namespace=typedefnamespace)

    for rid, typeref in iter_dotnet_table(pe, dnfile.mdtable.TypeRef.number):
        assert isinstance(typeref, dnfile.mdtable.TypeRefRow)

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

        typeref_token: int = calculate_dotnet_token_value(dnfile.mdtable.TypeRef.number, rid)
        yield DnType(typeref_token, typerefname, namespace=typerefnamespace)


Review comment (suggestion):
```suggestion
    # Used to find nested classes in TypeDef table
```

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

    return nested_class_table


def get_dotnet_types(pe: dnfile.dnPE) -> Iterator[DnType]:
    """get .NET types from TypeDef and TypeRef tables"""
    nested_class_table = get_dotnet_nested_class_table_index(pe)

    for rid, typedef in iter_dotnet_table(pe, dnfile.mdtable.TypeDef.number):
        assert isinstance(typedef, dnfile.mdtable.TypeDefRow)

        typedefnamespace, typedefname = resolve_nested_typedef_name(nested_class_table, rid, typedef, pe)

        typedef_token: int = calculate_dotnet_token_value(dnfile.mdtable.TypeDef.number, rid)
        yield DnType(typedef_token, typedefname, namespace=typedefnamespace)

    for rid, typeref in iter_dotnet_table(pe, dnfile.mdtable.TypeRef.number):
        assert isinstance(typeref, dnfile.mdtable.TypeRefRow)

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

        typeref_token: int = calculate_dotnet_token_value(dnfile.mdtable.TypeRef.number, rid)