Fix regression: LF_UNION should be included in $STRUCTS (#12)

drakpdb/type_info.py

@@ -107,7 +107,7 @@ def process_structure_member(member):
 
 
 def process_structure(struct):
-    if struct.leaf_type != "LF_STRUCTURE":
+    if struct.leaf_type not in ["LF_STRUCTURE", "LF_UNION"]:
         # Unhandled type of structure
         return [0, {}]
     if not hasattr(struct.fieldlist, "substructs"):
@@ -129,7 +129,8 @@ def process_tpi(pdb):
     """
     return {
         "$STRUCTS": {
-            name: process_structure(structure)
-            for name, structure in pdb.STREAM_TPI.structures.items()
+            type_info.name: process_structure(type_info)
+            for type_info in pdb.STREAM_TPI.types.values()
+            if hasattr(type_info, "name")
         }
     }

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "drakpdb"
-version = "0.2.1"
+version = "0.2.2"
 description = "Helper library to generate DRAKVUF profiles."
 readme = "README.md"
 classifiers = [

tests/conftest.py

@@ -8,3 +8,6 @@ def pytest_generate_tests(metafunc):
     if 'pdb_file' in metafunc.fixturenames:
         pdbs_files = list(pdbs_dir.glob("*.pdb"))
         metafunc.parametrize('pdb_file', pdbs_files)
+    if 'nt_kernel_pdb_file' in metafunc.fixturenames:
+        pdbs_files = list(pdbs_dir.glob("ntkrnlmp*.pdb"))
+        metafunc.parametrize('nt_kernel_pdb_file', pdbs_files)

tests/test_pdb.py

@@ -7,3 +7,14 @@ def test_pdb_profile(pdb_file):
     assert profile["$FUNCTIONS"]
     assert profile["$CONSTANTS"]
     assert profile["$METADATA"]
+
+
+def test_kernel_types(nt_kernel_pdb_file):
+    profile = make_pdb_profile(nt_kernel_pdb_file)
+    structs = profile["$STRUCTS"]
+    # There are parsed structs
+    assert structs
+    # including unions like _HANDLE_TABLE_ENTRY
+    assert "_HANDLE_TABLE_ENTRY" in structs
+    # Size of structure is more than 0
+    assert structs["_HANDLE_TABLE_ENTRY"][0] > 0