Add extra_flags to register sets (#2037)

pwndbg/commands/context.py

@@ -594,10 +594,14 @@ def get_regs(regs: List[str] = None):
         change_marker = "%s" % C.config_register_changed_marker
         m = " " * len(change_marker) if reg not in changed else C.register_changed(change_marker)
 
+        bit_flags = None
         if reg in pwndbg.gdblib.regs.flags:
-            desc = C.format_flags(
-                value, pwndbg.gdblib.regs.flags[reg], pwndbg.gdblib.regs.last.get(reg, 0)
-            )
+            bit_flags = pwndbg.gdblib.regs.flags[reg]
+        elif reg in pwndbg.gdblib.regs.extra_flags:
+            bit_flags = pwndbg.gdblib.regs.extra_flags[reg]
+
+        if bit_flags:
+            desc = C.format_flags(value, bit_flags, pwndbg.gdblib.regs.last.get(reg, 0))
 
         else:
             desc = pwndbg.chain.format(value)

pwndbg/gdblib/regs.py

@@ -24,7 +24,11 @@ from pwndbg.lib.regs import reg_sets
 
 @pwndbg.gdblib.proc.OnlyWhenRunning
 def gdb_get_register(name: str):
-    return gdb.selected_frame().read_register(name)
+    frame = gdb.selected_frame()
+    try:
+        return frame.read_register(name)
+    except ValueError:
+        return frame.read_register(name.upper())
 
 
 # We need to manually make some ptrace calls to get fs/gs bases on Intel
@@ -110,6 +114,10 @@ class module(ModuleType):
     def flags(self):
         return reg_sets[pwndbg.gdblib.arch.current].flags
 
+    @property
+    def extra_flags(self):
+        return reg_sets[pwndbg.gdblib.arch.current].extra_flags
+
     @property
     def stack(self):
         return reg_sets[pwndbg.gdblib.arch.current].stack

pwndbg/lib/regs.py

@@ -55,6 +55,7 @@ class RegisterSet:
         frame: str | None = None,
         retaddr: Tuple[str, ...] = tuple(),
         flags: Dict[str, BitFlags] = {},
+        extra_flags: Dict[str, BitFlags] = {},
         gpr: Tuple[str, ...] = tuple(),
         misc: Tuple[str, ...] = tuple(),
         args: Tuple[str, ...] = tuple(),
@@ -65,6 +66,7 @@ class RegisterSet:
         self.frame = frame
         self.retaddr = retaddr
         self.flags = flags
+        self.extra_flags = extra_flags
         self.gpr = gpr
         self.misc = misc
         self.args = args
@@ -76,7 +78,9 @@ class RegisterSet:
             if reg and reg not in self.common:
                 self.common.append(reg)
 
-        self.all = {i for i in misc} | set(flags) | set(self.retaddr) | set(self.common)
+        self.all = (
+            {i for i in misc} | set(flags) | set(extra_flags) | set(self.retaddr) | set(self.common)
+        )
         self.all -= {None}
 
     def __iter__(self) -> Iterator[str]:
@@ -147,6 +151,11 @@ armcm = RegisterSet(
 aarch64 = RegisterSet(
     retaddr=("lr",),
     flags={"cpsr": aarch64_cpsr_flags},
+    extra_flags={
+        "spsr_el1": aarch64_cpsr_flags,
+        "spsr_el2": aarch64_cpsr_flags,
+        "spsr_el3": aarch64_cpsr_flags,
+    },
     # X29 is the frame pointer register (FP) but setting it
     # as frame here messes up the register order to the point
     # it's confusing. Think about improving this if frame