dotnet update: emit enclosing class information for nested classes

capa/features/extractors/dnfile/helpers.py

@@ -108,6 +108,10 @@ def get_dotnet_managed_imports(pe: dnfile.dnPE) -> Iterator[DnType]:
             TypeName (index into String heap)
             TypeNamespace (index into String heap)
     """
+    #typeref_table = get_typeref_table(pe)
+    #typeref_class_names = []
+    #typeref_assembled_class_names = {}
+
     for rid, member_ref in iter_dotnet_table(pe, dnfile.mdtable.MemberRef.number):
         assert isinstance(member_ref, dnfile.mdtable.MemberRefRow)
 
@@ -131,6 +135,9 @@ def get_dotnet_managed_imports(pe: dnfile.dnPE) -> Iterator[DnType]:
             # remove get_/set_ from MemberRef name
             member_ref_name = member_ref_name[4:]
 
+        #typeref_class_names.append((member_ref.Class.row.TypeName, member_ref.Class.row.TypeNamespace))
+        #member_ref.Class.row.TypeNamespace, member_ref.Class.row.TypeName, typeref_assembled_class_names = is_nested_helper(rid, typeref_table, typeref_class_names, member_ref.Class.row, typeref_assembled_class_names, False)
+
         yield DnType(
             token,
             member_ref.Class.row.TypeName,
@@ -188,6 +195,10 @@ def get_dotnet_managed_methods(pe: dnfile.dnPE) -> Iterator[DnType]:
             TypeNamespace (index into String heap)
             MethodList (index into MethodDef table; it marks the first of a contiguous run of Methods owned by this Type)
     """
+    #nested_class_table = get_nested_class_table(pe)
+    #typedef_class_names = []
+    #typedef_assembled_class_names = {}
+
     accessor_map: Dict[int, str] = {}
     for methoddef, methoddef_access in get_dotnet_methoddef_property_accessors(pe):
         accessor_map[methoddef] = methoddef_access
@@ -211,6 +222,9 @@ def get_dotnet_managed_methods(pe: dnfile.dnPE) -> Iterator[DnType]:
                 # remove get_/set_
                 method_name = method_name[4:]
 
+            #typedef_class_names.append((method_name, typedef.TypeNamespace))
+            #typedef.TypeNamespace, typedef.TypeName, typedef_assembled_class_names = is_nested_helper(idx, nested_class_table, typedef_class_names, typedef, typedef_assembled_class_names, True)
+
             yield DnType(token, typedef.TypeName, namespace=typedef.TypeNamespace, member=method_name, access=access)
 
 
@@ -299,22 +313,102 @@ def get_dotnet_unmanaged_imports(pe: dnfile.dnPE) -> Iterator[DnUnmanagedMethod]
         # like kernel32.CreateFileA
         yield DnUnmanagedMethod(token, module, method)
 
+def get_nested_class_table(pe):
+    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_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.struct.ResolutionScope_CodedIndex))
+
+    return typeref_table
+
+def is_nested_helper(rid, table, class_names, typeX, assembled_class_names, is_typedef):
+    if is_typedef:
+        name = typeX.TypeName
+        space = typeX.TypeNamespace
+        nested_classes = table
+
+        if rid in nested_classes:
+            space = class_names[nested_classes[rid]-1][1]
+
+            enclosing_class = class_names[nested_classes[rid]-1][0]
+            nested_class = class_names[rid-1][0]
+
+            if enclosing_class in assembled_class_names:
+                enclosing_class = f"{assembled_class_names[enclosing_class]}"
+                assembled_class_names[nested_class] = enclosing_class
+
+            for i in class_names:
+                if i[0] == enclosing_class.split('/')[0]:
+                    space = i[1]
+
+            name = (enclosing_class, nested_class)
+            assembled_class_names[name[1]] = f"{name[0]}/{name[1]}"
+
+    else:
+        name = typeX.TypeName.split('`')[0] if '`' in typeX.TypeName else typeX.TypeName
+        space = typeX.TypeNamespace
+        typeref_table = table
+
+        if typeX.struct.ResolutionScope_CodedIndex <= len(typeref_table):
+            space = typeref_table[typeX.struct.ResolutionScope_CodedIndex-1][1]
+
+            enclosing_class = typeref_table[typeX.struct.ResolutionScope_CodedIndex-1][0]
+            nested_class = name
+
+            if enclosing_class in assembled_class_names:
+                enclosing_class = f"{assembled_class_names[enclosing_class]}"
+                assembled_class_names[nested_class] = enclosing_class
+
+            for i in class_names:
+                if i[0] == enclosing_class.split('/')[0]:
+                    space = i[1]
+
+            name = (enclosing_class, nested_class)
+            assembled_class_names[name[1]] = f"{name[0]}/{name[1]}"
+
+    return space, name, assembled_class_names
 
 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_names = []
+    typedef_assembled_class_names = {}
+
     for rid, typedef in iter_dotnet_table(pe, dnfile.mdtable.TypeDef.number):
         assert isinstance(typedef, dnfile.mdtable.TypeDefRow)
 
+        typedef_class_names.append((typedef.TypeName, typedef.TypeNamespace))
+        typedef.TypeNamespace, typedef.TypeName, typedef_assembled_class_names = is_nested_helper(rid, nested_class_table, typedef_class_names, typedef, typedef_assembled_class_names, True)
+
         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)
+    typeref_class_names = []
+    typeref_assembled_class_names = {}
+
     for rid, typeref in iter_dotnet_table(pe, dnfile.mdtable.TypeRef.number):
         assert isinstance(typeref, dnfile.mdtable.TypeRefRow)
-
+
+        typeref_class_names.append((typeref.TypeName, typeref.TypeNamespace))
+        typeref.TypeNamespace, typeref.TypeName, typeref_assembled_class_names = is_nested_helper(rid, typeref_table, typeref_class_names, typeref, typeref_assembled_class_names, False)
+
         typeref_token: int = calculate_dotnet_token_value(dnfile.mdtable.TypeRef.number, rid)
         yield DnType(typeref_token, typeref.TypeName, namespace=typeref.TypeNamespace)
-
-
+
 def calculate_dotnet_token_value(table: int, rid: int) -> int:
     return ((table & 0xFF) << Token.TABLE_SHIFT) | (rid & Token.RID_MASK)