From 76eba8074703d2d54ac9f20755547fa5e1ed4c9c Mon Sep 17 00:00:00 2001
From: "Matt." <dark.ryu.550@gmail.com>
Date: Sat, 17 Aug 2024 08:13:59 -0300
Subject: [PATCH] Backport use of `pwndbg.aglib.arch.update()` to current
 upstream (#2373)

---
 pwndbg/aglib/__init__.py |  4 ++++
 pwndbg/aglib/arch.py     | 45 ++++++++++++++++++++++++++++++++--------
 pwndbg/gdblib/hooks.py   |  6 ++++--
 3 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/pwndbg/aglib/__init__.py b/pwndbg/aglib/__init__.py
index e69de29bb..e227f7bb3 100644
--- a/pwndbg/aglib/__init__.py
+++ b/pwndbg/aglib/__init__.py
@@ -0,0 +1,4 @@
+from __future__ import annotations
+
+from pwndbg.aglib import arch as arch_mod
+from pwndbg.aglib.arch import arch as arch
diff --git a/pwndbg/aglib/arch.py b/pwndbg/aglib/arch.py
index 8ab54a531..df5c68393 100644
--- a/pwndbg/aglib/arch.py
+++ b/pwndbg/aglib/arch.py
@@ -1,15 +1,42 @@
 from __future__ import annotations
 
 import pwndbg
+from pwndbg.lib.arch import Arch
 
-# We will optimize this module in the future, by having it work in the same
-# way the `gdblib` version of it works, and that will come at the same
-# time this module gets expanded to have the full feature set of its `gdlib`
-# coutnerpart. For now, though, this should be good enough.
+ARCHS = (
+    "x86-64",
+    "i386",
+    "aarch64",
+    "mips",
+    "powerpc",
+    "sparc",
+    "arm",
+    "armcm",
+    "riscv:rv32",
+    "riscv:rv64",
+    "riscv",
+)
 
 
-def __getattr__(name):
-    if name == "endian":
-        return pwndbg.dbg.selected_inferior().arch().endian
-    elif name == "ptrsize":
-        return pwndbg.dbg.selected_inferior().arch().ptrsize
+# mapping between gdb and pwntools arch names
+pwnlib_archs_mapping = {
+    "x86-64": "amd64",
+    "i386": "i386",
+    "aarch64": "aarch64",
+    "mips": "mips",
+    "powerpc": "powerpc",
+    "sparc": "sparc",
+    "arm": "arm",
+    "iwmmxt": "arm",
+    "armcm": "thumb",
+    "rv32": "riscv32",
+    "rv64": "riscv64",
+}
+
+
+arch: Arch = Arch("i386", 4, "little")
+
+
+def update() -> None:
+    a = pwndbg.dbg.selected_inferior().arch()
+    arch.update(a.name, a.ptrsize, a.endian)
diff --git a/pwndbg/gdblib/hooks.py b/pwndbg/gdblib/hooks.py
index c27c6f839..91f6536f4 100644
--- a/pwndbg/gdblib/hooks.py
+++ b/pwndbg/gdblib/hooks.py
@@ -8,8 +8,9 @@ import pwndbg.gdblib.memory
 import pwndbg.gdblib.next
 import pwndbg.gdblib.tls
 import pwndbg.gdblib.typeinfo
+from pwndbg.aglib import arch_mod as arch_mod_aglib
 from pwndbg.dbg import EventType
-from pwndbg.gdblib import arch_mod
+from pwndbg.gdblib import arch_mod as arch_mod_gdblib
 
 # TODO: Combine these `update_*` hook callbacks into one method
 
@@ -25,7 +26,8 @@ def update_typeinfo() -> None:
 @pwndbg.dbg.event_handler(EventType.STOP)
 @pwndbg.dbg.event_handler(EventType.NEW_MODULE)
 def update_arch() -> None:
-    arch_mod.update()
+    arch_mod_gdblib.update()
+    arch_mod_aglib.update()
 
 
 @pwndbg.dbg.event_handler(EventType.NEW_MODULE)