Fetch C struct as Python dictionary (#2082)

Add fetch_struct_as_dictionary

---------

Co-authored-by: Gulshan Singh <gsingh2011@gmail.com>

pwndbg/gdblib/memory.py

@@ -5,6 +5,8 @@ Reading, writing, and describing memory.
 from __future__ import annotations
 
 import re
+from typing import Dict
+from typing import Union
 
 import gdb
 
@@ -17,6 +19,9 @@ import pwndbg.lib.memory
 from pwndbg.lib.memory import PAGE_MASK
 from pwndbg.lib.memory import PAGE_SIZE
 
+GdbDict = Dict[str, Union["GdbDict", int]]
+
+
 MMAP_MIN_ADDR = 0x8000
 
 
@@ -360,3 +365,55 @@ def update_min_addr() -> None:
     global MMAP_MIN_ADDR
     if pwndbg.gdblib.qemu.is_qemu_kernel():
         MMAP_MIN_ADDR = 0
+
+
+def fetch_struct_as_dictionary(
+    struct_name: str,
+    struct_address: int,
+    include_only_fields: set[str] = set(),
+    exclude_fields: set[str] = set(),
+) -> GdbDict:
+    struct_type = gdb.lookup_type("struct " + struct_name)
+    fetched_struct = poi(struct_type, struct_address)
+
+    return pack_struct_into_dictionary(fetched_struct, include_only_fields, exclude_fields)
+
+
+def pack_struct_into_dictionary(
+    fetched_struct: gdb.Value,
+    include_only_fields: set[str] = set(),
+    exclude_fields: set[str] = set(),
+) -> GdbDict:
+    struct_as_dictionary = {}
+
+    if len(include_only_fields) != 0:
+        for field_name in include_only_fields:
+            key = field_name
+            value = convert_gdb_value_to_python_value(fetched_struct[field_name])
+            struct_as_dictionary[key] = value
+    else:
+        for field in fetched_struct.type.fields():
+            if field.name is None:
+                # Flatten anonymous structs/unions
+                anon_type = convert_gdb_value_to_python_value(fetched_struct[field])
+                assert isinstance(anon_type, dict)
+                struct_as_dictionary.update(anon_type)
+            elif field.name not in exclude_fields:
+                key = field.name
+                value = convert_gdb_value_to_python_value(fetched_struct[field])
+                struct_as_dictionary[key] = value
+
+    return struct_as_dictionary
+
+
+def convert_gdb_value_to_python_value(gdb_value: gdb.Value) -> int | GdbDict:
+    gdb_type = gdb_value.type.strip_typedefs()
+
+    if gdb_type.code == gdb.TYPE_CODE_PTR:
+        return int(gdb_value)
+    elif gdb_type.code == gdb.TYPE_CODE_INT:
+        return int(gdb_value)
+    elif gdb_type.code == gdb.TYPE_CODE_STRUCT:
+        return pack_struct_into_dictionary(gdb_value)
+
+    raise NotImplementedError

tests/gdb-tests/tests/binaries/nested_structs.c

@@ -0,0 +1,78 @@
+/* This program initializes some nested C structs.
+ * Useful for testing pwndbg commands that operate on structs.
+ */
+
+/* Can a command deal with nested typedefs?
+ * mydef_outer -> mydef_inner -> int
+ */
+typedef int mydef_inner;
+typedef mydef_inner mydef_outer;
+
+/* Can a command deal with anonymous structs?
+ * ISO C11 says anonymous_i & anonymous_j fields should be accessible like this:
+ * inner_struct.anonymous_i
+ */
+struct inner_struct
+{
+    int inner_a;
+    mydef_outer inner_b; // int
+
+    struct
+    {
+        int anonymous_i;
+        int anonymous_j;
+    };
+};
+
+/* Can a command deal with nested named structs and nested anonymous structs?
+ * The anonymous_nested field should be accessible like this:
+ * outer_struct.anonymous_nested
+ */
+struct outer_struct
+{
+    int outer_x;
+    mydef_outer outer_y; // int
+
+    struct inner_struct inner;
+
+    struct
+    {
+        int anonymous_k;
+        int anonymous_l;
+
+        struct
+        {
+            int anonymous_nested;
+        };
+    };
+
+    int outer_z;
+};
+
+// Set a breakpoint on this function to stop in the important places.
+void break_here(void) {}
+
+struct outer_struct outer;
+
+int main(void)
+{
+    // Initialize outer_struct fields with arbitrary values.
+    outer.outer_x = 1;
+    outer.outer_y = 2;
+    outer.outer_z = 5;
+
+    outer.inner.inner_a = 3;
+    outer.inner.inner_b = 4;
+
+    outer.inner.anonymous_i = 42;
+    outer.inner.anonymous_j = 44;
+
+    outer.anonymous_nested = 100;
+
+    outer.anonymous_k = 82;
+    outer.anonymous_l = 84;
+
+    break_here();
+
+    return 0;
+}

tests/gdb-tests/tests/test_memory.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
+import gdb
+
 import pwndbg.gdblib.memory
 import pwndbg.gdblib.stack
 import tests
 
 REFERENCE_BINARY = tests.binaries.get("reference-binary.out")
+NESTED_STRUCTS_BINARY = tests.binaries.get("nested_structs.out")
 
 
 def test_memory_read_write(start_binary):
@@ -30,3 +33,87 @@ def test_memory_read_write(start_binary):
     assert pwndbg.gdblib.memory.read(stack_addr, len(val) + 4) == bytearray(
         "Z" * 8 + "YYXX", "utf8"
     )
+
+
+def test_fetch_struct_as_dictionary(start_binary):
+    """
+    Test pwndbg.gdblib.memory.fetch_struct_as_dictionary()
+    Ensure it can handle nested structs, anonymous structs & nested typedefs.
+    """
+    start_binary(NESTED_STRUCTS_BINARY)
+    gdb.execute("break break_here")
+    gdb.execute("continue")
+
+    expected_result = {
+        "outer_x": 1,
+        "outer_y": 2,
+        "inner": {"inner_a": 3, "inner_b": 4, "anonymous_i": 42, "anonymous_j": 44},
+        "anonymous_k": 82,
+        "anonymous_l": 84,
+        "anonymous_nested": 100,
+        "outer_z": 5,
+    }
+
+    struct_address = pwndbg.gdblib.symbol.address("outer")
+    assert struct_address is not None
+
+    result = pwndbg.gdblib.memory.fetch_struct_as_dictionary("outer_struct", struct_address)
+
+    assert result == expected_result
+
+
+def test_fetch_struct_as_dictionary_include_filter(start_binary):
+    """
+    Test pwndbg.gdblib.memory.fetch_struct_as_dictionary()
+    Ensure its include_only_fields filter works.
+    """
+    start_binary(NESTED_STRUCTS_BINARY)
+    gdb.execute("break break_here")
+    gdb.execute("continue")
+
+    expected_result = {
+        "outer_x": 1,
+        "inner": {"inner_a": 3, "inner_b": 4, "anonymous_i": 42, "anonymous_j": 44},
+        "anonymous_k": 82,
+        "anonymous_nested": 100,
+    }
+
+    struct_address = pwndbg.gdblib.symbol.address("outer")
+    assert struct_address is not None
+
+    result = pwndbg.gdblib.memory.fetch_struct_as_dictionary(
+        "outer_struct",
+        struct_address,
+        include_only_fields={"outer_x", "inner", "anonymous_k", "anonymous_nested"},
+    )
+
+    assert result == expected_result
+
+
+def test_fetch_struct_as_dictionary_exclude_filter(start_binary):
+    """
+    Test pwndbg.gdblib.memory.fetch_struct_as_dictionary()
+    Ensure its exclude_fields filter works.
+    Note that the exclude filter cannot filter fields of anonymous structs.
+    """
+    start_binary(NESTED_STRUCTS_BINARY)
+    gdb.execute("break break_here")
+    gdb.execute("continue")
+
+    expected_result = {
+        "outer_y": 2,
+        "anonymous_k": 82,
+        "anonymous_l": 84,
+        "anonymous_nested": 100,
+    }
+
+    struct_address = pwndbg.gdblib.symbol.address("outer")
+    assert struct_address is not None
+
+    result = pwndbg.gdblib.memory.fetch_struct_as_dictionary(
+        "outer_struct",
+        struct_address,
+        exclude_fields={"outer_x", "inner", "outer_z"},
+    )
+
+    assert result == expected_result