add caching

pwndbg/commands/addsymbol.py

@@ -1,19 +1,23 @@
 from __future__ import annotations
 
 import argparse
+import hashlib
+import os
 
 import gdb
 
 import pwndbg.commands
 from pwndbg.commands import CommandCategory
-from pwndbg.gdblib.symbol import _create_symboled_elf
+from pwndbg.gdblib.symbol import create_symboled_elf
 
 parser = argparse.ArgumentParser(description="add custom symbols")
 parser.add_argument("name", type=str, help="name of the symbol")
 parser.add_argument("addr", type=int, help="addr of the symbol")
 
+SYMBOLS_CACHEDIR = pwndbg.lib.tempfile.cachedir("symbols")
 
-@pwndbg.commands.ArgparsedCommand(parser, category=CommandCategory.LINUX)
+
+@pwndbg.commands.ArgparsedCommand(parser, category=CommandCategory.MISC)
 @pwndbg.commands.OnlyWhenRunning
 def addsymbol(name, addr) -> None:
     module = pwndbg.gdblib.proc.exe
@@ -23,6 +27,20 @@ def addsymbol(name, addr) -> None:
         if module in p.objfile:
             vaddr = p.vaddr
 
-    path = _create_symboled_elf({name: addr}, base_addr=vaddr)
+    path = create_symboled_elf(
+        {name: addr},
+        base_addr=vaddr,
+        filename=os.path.join(SYMBOLS_CACHEDIR, compute_file_hash(module)),
+    )
 
     gdb.execute(f"add-symbol-file {path} {vaddr}")
+
+
+def compute_file_hash(filename: str) -> str:
+    """
+    Compute the MD5 hash of the file, return the hash
+    """
+    h = hashlib.md5()
+    with open(filename, "rb") as f:
+        h.update(f.read())
+    return h.hexdigest()

pwndbg/gdblib/symbol.py

@@ -10,7 +10,6 @@
 
 import os
 import re
-import tempfile
 from typing import Dict
 
 import gdb
@@ -49,43 +48,42 @@
 )
 
 
-def _create_symboled_elf(symbols: Dict[str, int], base_addr: int = 0, filename: str = None) -> str:
-    # TODO: cache kernel symbol elfs for kallsyms command
-    fd, pwndbg_debug_symbols_output_file = tempfile.mkstemp(prefix="symbols-", suffix=".c")
-    os.fdopen(fd, "w").write("int main(){}")
-    os.system(
-        f"gcc {pwndbg_debug_symbols_output_file} -o {pwndbg_debug_symbols_output_file[0:-2]}.debug"
-    )
-    os.unlink(f"{pwndbg_debug_symbols_output_file}")
+def create_symboled_elf(symbols: Dict[str, int], base_addr: int = 0, filename: str = None) -> str:
+    if not os.path.exists(filename + ".debug"):
+        filename = filename + ".c"
+        with open(filename, "w") as f:
+            f.write("int main(){}")
 
-    pwndbg_debug_symbols_output_file = pwndbg_debug_symbols_output_file[0:-2]
+        os.system(f"gcc {filename} -o {filename[0:-2]}.debug")
+        os.unlink(f"{filename}")
 
-    os.system(f"objcopy --only-keep-debug {pwndbg_debug_symbols_output_file}.debug")
-    os.system(f"objcopy --strip-all {pwndbg_debug_symbols_output_file}.debug")
+        filename = filename[0:-2]
 
-    elf = ELFFile(open(f"{pwndbg_debug_symbols_output_file}.debug", "rb"))
+        os.system(f"objcopy --only-keep-debug {filename}.debug")
+        os.system(f"objcopy --strip-all {filename}.debug")
 
-    required_sections = [".text", ".interp", ".rela.dyn", ".dynamic", ".bss"]
+        elf = ELFFile(open(f"{filename}.debug", "rb"))
 
-    removable_sections = ""
+        required_sections = [".text", ".interp", ".rela.dyn", ".dynamic", ".bss"]
 
-    for s in elf.iter_sections():
-        if s.name in required_sections:
-            continue
+        removable_sections = ""
 
-        removable_sections += f"--remove-section={s.name} "
+        for s in elf.iter_sections():
+            if s.name in required_sections:
+                continue
 
-    os.system(f"objcopy {removable_sections} {pwndbg_debug_symbols_output_file}.debug 2>/dev/null")
-    os.system(
-        f"objcopy --change-section-address .text={base_addr:#x} {pwndbg_debug_symbols_output_file}.debug"
-    )
+            removable_sections += f"--remove-section={s.name} "
 
+        os.system(f"objcopy {removable_sections} {filename}.debug 2>/dev/null")
+        os.system(f"objcopy --change-section-address .text={base_addr:#x} {filename}.debug")
+
+    # TODO: add more symbols in one call
     for symbol in symbols.items():
         os.system(
-            f"objcopy --add-symbol {symbol[0]}=.text:{symbol[1]:#x},global,function {pwndbg_debug_symbols_output_file}.debug"
+            f"objcopy --add-symbol {symbol[0]}=.text:{symbol[1]:#x},global,function {filename}.debug"
         )
 
-    return f"{pwndbg_debug_symbols_output_file}.debug"
+    return f"{filename}.debug"
 
 
 def _get_debug_file_directory() -> str: