From 53d0aa531e89971a7d609220d677486f3c39c1a8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Luis=20Di=20Biase?= <josx@interorganic.com.ar>
Date: Tue, 12 Sep 2023 18:40:39 -0300
Subject: [PATCH] Refactor reset_architecture: check for param forced, binary
 running, elf header
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: José Luis Di Biase <josx@interorganic.com.ar>
---
 gef.py | 55 ++++++++++++++++---------------------------------------
 1 file changed, 16 insertions(+), 39 deletions(-)

diff --git a/gef.py b/gef.py
index 5c187451f..f35a31b61 100644
--- a/gef.py
+++ b/gef.py
@@ -2150,34 +2150,6 @@ def checksec(filename: str) -> Dict[str, bool]:
     return Elf(filename).checksec
 
 
-@lru_cache()
-def get_arch() -> str:
-    """Return the binary's architecture."""
-    if is_alive():
-        arch = gdb.selected_frame().architecture()
-        return arch.name()
-
-    arch_str = gdb.execute("show architecture", to_string=True).strip()
-    pat = "The target architecture is set automatically (currently "
-    if arch_str.startswith(pat):
-        arch_str = arch_str[len(pat):].rstrip(")")
-        return arch_str
-
-    pat = "The target architecture is assumed to be "
-    if arch_str.startswith(pat):
-        return arch_str[len(pat):]
-
-    pat = "The target architecture is set to "
-    if arch_str.startswith(pat):
-        # GDB version >= 10.1
-        if '"auto"' in arch_str:
-            return re.findall(r"currently \"(.+)\"", arch_str)[0]
-        return re.findall(r"\"(.+)\"", arch_str)[0]
-
-    # Unknown, we throw an exception to be safe
-    raise RuntimeError(f"Unknown architecture: {arch_str}")
-
-
 @deprecated("Use `gef.binary.entry_point` instead")
 def get_entry_point() -> Optional[int]:
     """Return the binary entry point."""
@@ -3679,20 +3651,25 @@ def reset_architecture(arch: Optional[str] = None) -> None:
             raise OSError(f"Specified arch {arch.upper()} is not supported")
         return
 
-    gdb_arch = get_arch()
+    # check for bin running
+    if is_alive():
+        try:
+            arch = gdb.selected_frame().architecture()
+            gef.arch = arches[arch.name()]()
+        except KeyError:
+            raise OSError(f"Running Binary arch {arch.name().upper()} is not supported")
+        return
 
-    preciser_arch = next((a for a in arches.values() if a.supports_gdb_arch(gdb_arch)), None)
-    if preciser_arch:
-        gef.arch = preciser_arch()
+    # check for elf header architecture
+    if gef.binary.e_machine:
+        try:
+            gef.arch = arches[gef.binary.e_machine]()
+        except KeyError:
+            raise OSError(f"CPU type is currently not supported: {gef.binary.e_machine}")
         return
 
-    # last resort, use the info from elf header to find it from the known architectures
-    try:
-        arch_name = gef.binary.e_machine if gef.binary else gdb_arch
-        gef.arch = arches[arch_name]()
-    except KeyError:
-        raise OSError(f"CPU type is currently not supported: {get_arch()}")
-    return
+    # Unknown, we throw an exception to be safe
+    raise RuntimeError(f"Unknown architecture")
 
 
 @lru_cache()