From f08c44f83c970d0a8cd562a6184e682a569584b0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Srokosz?= <psrok1@gmail.com>
Date: Thu, 1 Aug 2024 15:50:13 +0200
Subject: [PATCH] Fix regression: LF_UNION should be included in $STRUCTS (#12)

---
 drakpdb/type_info.py |  7 ++++---
 pyproject.toml       |  2 +-
 tests/conftest.py    |  3 +++
 tests/test_pdb.py    | 11 +++++++++++
 4 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/drakpdb/type_info.py b/drakpdb/type_info.py
index bfd2fbe..4d2c516 100644
--- a/drakpdb/type_info.py
+++ b/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")
         }
     }
diff --git a/pyproject.toml b/pyproject.toml
index edd704d..b9d452b 100644
--- a/pyproject.toml
+++ b/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 = [
diff --git a/tests/conftest.py b/tests/conftest.py
index 285f9cc..c0c45db 100644
--- a/tests/conftest.py
+++ b/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)
diff --git a/tests/test_pdb.py b/tests/test_pdb.py
index ee3636f..6464513 100644
--- a/tests/test_pdb.py
+++ b/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