From 34b689e695710d8ff5212af6582136778ad09ec1 Mon Sep 17 00:00:00 2001
From: Charif Mostafa <charif@chrf.dev>
Date: Tue, 28 May 2024 00:31:32 +0200
Subject: [PATCH] WIP

---
 pwndbg/commands/__init__.py  |  1 +
 pwndbg/commands/addsymbol.py | 28 +++++++++++++++++++++++
 pwndbg/gdblib/symbol.py      | 43 ++++++++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+)
 create mode 100644 pwndbg/commands/addsymbol.py

diff --git a/pwndbg/commands/__init__.py b/pwndbg/commands/__init__.py
index a1f3ef8b359..d604440b6f3 100644
--- a/pwndbg/commands/__init__.py
+++ b/pwndbg/commands/__init__.py
@@ -653,6 +653,7 @@ def HexOrAddressExpr(s: str) -> int:
 
 def load_commands() -> None:
     # pylint: disable=import-outside-toplevel
+    import pwndbg.commands.addsymbol
     import pwndbg.commands.ai
     import pwndbg.commands.argv
     import pwndbg.commands.aslr
diff --git a/pwndbg/commands/addsymbol.py b/pwndbg/commands/addsymbol.py
new file mode 100644
index 00000000000..ea18e9b931b
--- /dev/null
+++ b/pwndbg/commands/addsymbol.py
@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import argparse
+
+import gdb
+
+import pwndbg.commands
+from pwndbg.commands import CommandCategory
+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")
+
+
+@pwndbg.commands.ArgparsedCommand(parser, category=CommandCategory.LINUX)
+@pwndbg.commands.OnlyWhenRunning
+def addsymbol(name, addr) -> None:
+    module = pwndbg.gdblib.proc.exe
+    vaddr = 0x0
+
+    for p in pwndbg.gdblib.vmmap.get():
+        if module in p.objfile:
+            vaddr = p.vaddr
+
+    path = _create_symboled_elf({name: addr}, base_addr=vaddr)
+
+    gdb.execute(f"add-symbol-file {path} {vaddr}")
diff --git a/pwndbg/gdblib/symbol.py b/pwndbg/gdblib/symbol.py
index 70d9622297e..084124c1054 100644
--- a/pwndbg/gdblib/symbol.py
+++ b/pwndbg/gdblib/symbol.py
@@ -8,9 +8,13 @@
 
 from __future__ import annotations
 
+import os
 import re
+import tempfile
+from typing import Dict
 
 import gdb
+from elftools.elf.elffile import ELFFile
 
 import pwndbg.gdblib.android
 import pwndbg.gdblib.arch
@@ -45,6 +49,45 @@
 )
 
 
+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}")
+
+    pwndbg_debug_symbols_output_file = pwndbg_debug_symbols_output_file[0:-2]
+
+    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")
+
+    elf = ELFFile(open(f"{pwndbg_debug_symbols_output_file}.debug", "rb"))
+
+    required_sections = [".text", ".interp", ".rela.dyn", ".dynamic", ".bss"]
+
+    removable_sections = ""
+
+    for s in elf.iter_sections():
+        if s.name in required_sections:
+            continue
+
+        removable_sections += f"--remove-section={s.name} "
+
+    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"
+    )
+
+    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"
+        )
+
+    return f"{pwndbg_debug_symbols_output_file}.debug"
+
+
 def _get_debug_file_directory() -> str:
     """
     Retrieve the debug file directory path.