load symbols from header files #1973 (#2661)

* feature: loading symbols from header files

* load symbols from header files

* fixed log warn issue

* added description for loading from header files

* added tests to test_cymbol.py

* added description in docs

* added additional test cases

* removed header file check

docs/commands/cymbol/cymbol.md

@@ -12,7 +12,7 @@ Add, show, load, edit, or delete custom structures in plain C.
 
 
 ```bash
-usage: cymbol [-h] [-a name] [-r name] [-e name] [-l name] [-s name]
+usage: cymbol [-h] [-a name] [-r name] [-e name] [-l name] [-s name] [-f filepath]
 
 ```
 ## Optional Arguments
@@ -21,6 +21,7 @@ usage: cymbol [-h] [-a name] [-r name] [-e name] [-l name] [-s name]
 | :--- | :--- | :--- | :--- |
 |`-h`|`--help`||show this help message and exit|
 |`-a`|`--add`|`None`|Add a new custom structure|
+|`-f`|`--file`|`None`|Loads custom structures from header files|
 |`-r`|`--remove`|`None`|Remove an existing custom structure|
 |`-e`|`--edit`|`None`|Edit an existing custom structure|
 |`-l`|`--load`|`None`|Load an existing custom structure|

pwndbg/commands/__init__.py

@@ -491,7 +491,7 @@ def OnlyWithResolvedHeapSyms(function: Callable[P, T]) -> Callable[P, T | None]:
     @functools.wraps(function)
     def _OnlyWithResolvedHeapSyms(*a: P.args, **kw: P.kwargs) -> T | None:
         e = log.error
-        w = log.warn
+        w = log.warning
         if (
             isinstance(pwndbg.aglib.heap.current, HeuristicHeap)
             and pwndbg.config.resolve_heap_via_heuristic == "auto"

pwndbg/commands/cymbol.py

@@ -153,6 +153,42 @@ def add_custom_structure(custom_structure_name: str) -> None:
     load_custom_structure.__wrapped__(custom_structure_name, pwndbg_custom_structure_path)
 
 
+def add_structure_from_header(header_file: str, custom_structure_name: str = None) -> None:
+    # Properly handle the provided or default name for the custom structure
+    custom_structure_name = (
+        custom_structure_name.strip()
+        if custom_structure_name
+        else os.path.splitext(os.path.basename(header_file))[0]
+    )
+
+    if not custom_structure_name:
+        print(message.error("Invalid structure name provided or generated."))
+        return
+
+    pwndbg_custom_structure_path = os.path.join(pwndbg_cachedir, custom_structure_name) + ".c"
+
+    if os.path.exists(pwndbg_custom_structure_path):
+        option = input(
+            message.notice(f"Structure '{custom_structure_name}' already exists. Overwrite? [y/n] ")
+        )
+        if option != "y":
+            return
+
+    try:
+        with open(header_file, "r") as src, open(pwndbg_custom_structure_path, "w") as f:
+            content = src.read().strip()
+            if not content:
+                print(message.notice("Header file is empty, skipping..."))
+                return
+            f.write(content)
+    except (IOError, OSError) as e:
+        print(message.error(f"Failed to process header file: {e}"))
+        return
+
+    # Avoid checking for file existance. Call the decorator wrapper directly.
+    load_custom_structure.__wrapped__(custom_structure_name, pwndbg_custom_structure_path)
+
+
 @OnlyWhenStructFileExists
 def edit_custom_structure(custom_structure_name: str, custom_structure_path: str = "") -> None:
     # Lookup an editor to use for editing the custom structure.
@@ -224,6 +260,14 @@ parser.add_argument(
     default=None,
     type=str,
 )
+parser.add_argument(
+    "-f",
+    "--file",
+    metavar="filepath",
+    help="Add a new custom structure from header file",
+    default=None,
+    type=str,
+)
 parser.add_argument(
     "-r",
     "--remove",
@@ -259,9 +303,11 @@ parser.add_argument(
 
 
 @pwndbg.commands.ArgparsedCommand(parser)
-def cymbol(add: str, remove: str, edit: str, load: str, show: str) -> None:
+def cymbol(add: str, file: str, remove: str, edit: str, load: str, show: str) -> None:
     if add:
         add_custom_structure(add)
+    elif file:
+        add_structure_from_header(file)
     elif remove:
         remove_custom_structure(remove)
     elif edit:

tests/gdb-tests/tests/test_cymbol.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import tempfile
 
 import pwndbg.aglib.dt
 import pwndbg.dbg
@@ -77,3 +78,83 @@ def test_cymbol(start_binary):
     # Test whether remove_custom_structure() works properly.
     pwndbg.commands.cymbol.remove_custom_structure("example")
     check_symbol_existance("example_t")
+
+
+def create_temp_header_file(content: str) -> str:
+    """Create a temporary header file with the given content."""
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".h") as tmp_file:
+        tmp_file.write(content.encode())
+        return tmp_file.name
+
+
+def test_cymbol_header_file(start_binary):
+    start_binary(REFERENCE_BINARY)
+
+    # Define the content of the header file
+    header_content = """
+    #include <stdint.h>
+    typedef struct example_struct_a {
+        int a;
+        char b[16];
+        char* c;
+        void* d;
+    } example_A;
+
+    typedef struct example_struct_b {
+        uint16_t X;
+    } example_B;
+
+    typedef struct example_struct_c {
+        char name[32];
+        int* data;
+        struct example_struct_a* next;
+    } example_C;
+    """
+
+    # Create a temporary header file
+    header_file_path = create_temp_header_file(header_content)
+
+    # Test adding structures from the header file
+    struct_name = "example_t"
+
+    pwndbg.commands.cymbol.add_structure_from_header(header_file_path, struct_name)
+    # Verify each structure has been loaded correctly
+    assert pwndbg.commands.cymbol.loaded_symbols.get(struct_name) is not None
+
+    # Check if the structure types match what we expect (on x86-64)
+    expected_outputs = {
+        "example_A": (
+            "example_A\n"
+            "    +0x0000 a                    : int\n"
+            "    +0x0004 b                    : char [16]\n"
+            "    +0x0018 c                    : char *\n"
+            "    +0x0020 d                    : void *"
+        ),
+        "example_B": ("example_B\n    +0x0000 X                    : uint16_t"),
+        "example_C": (
+            "example_C\n"
+            "    +0x0000 name                 : char [32]\n"
+            "    +0x0020 data                 : int *\n"
+            "    +0x0028 next                 : struct example_struct_a *"
+        ),
+    }
+
+    # Verify structure definitions
+    for struct_name, expected_output in expected_outputs.items():
+        assert expected_output == pwndbg.aglib.dt.dt(struct_name).strip()
+
+    # Test whether unload_loaded_symbol() works properly.
+    pwndbg.commands.cymbol.unload_loaded_symbol("example_t")
+    # Ensure the symbol is removed from the lookup loaded_symbols dict.
+    assert pwndbg.commands.cymbol.loaded_symbols.get("example_t") is None
+    # Ensure the symbol is no longer present in gdb.
+    check_symbol_existance("example_t")
+
+    # Load the symbol again for the next test case.
+    pwndbg.commands.cymbol.load_custom_structure("example_t")
+
+    # Test whether remove_custom_structure() works properly.
+    pwndbg.commands.cymbol.remove_custom_structure("example_t")
+    check_symbol_existance("example_t")
+    # Clean up temp files
+    os.remove(header_file_path)