diff --git a/docs/commands/cymbol/cymbol.md b/docs/commands/cymbol/cymbol.md index 602399f0f..af9ada277 100644 --- a/docs/commands/cymbol/cymbol.md +++ b/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| diff --git a/pwndbg/commands/__init__.py b/pwndbg/commands/__init__.py index 440c4de91..9349d37a0 100644 --- a/pwndbg/commands/__init__.py +++ b/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" diff --git a/pwndbg/commands/cymbol.py b/pwndbg/commands/cymbol.py index f1fb2c9a4..52e943f4a 100644 --- a/pwndbg/commands/cymbol.py +++ b/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: diff --git a/tests/gdb-tests/tests/test_cymbol.py b/tests/gdb-tests/tests/test_cymbol.py index 188d317db..288917a08 100644 --- a/tests/gdb-tests/tests/test_cymbol.py +++ b/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 + 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)