From 6b41c07e5b6377b55a39e6aaad1b04f823cfdcdd Mon Sep 17 00:00:00 2001
From: OB <obarroncs@gmail.com>
Date: Wed, 22 Nov 2023 18:38:11 -0800
Subject: [PATCH] Add sigreturn for i386 and aarch64

---
 pwndbg/commands/sigreturn.py | 71 ++++++++++++++++++------------------
 1 file changed, 35 insertions(+), 36 deletions(-)

diff --git a/pwndbg/commands/sigreturn.py b/pwndbg/commands/sigreturn.py
index 4965c4634..86052eb0b 100644
--- a/pwndbg/commands/sigreturn.py
+++ b/pwndbg/commands/sigreturn.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import argparse
+from typing import Tuple
 
 import pwnlib.rop.srop
 
@@ -11,7 +12,23 @@ import pwndbg.commands
 import pwndbg.gdblib.arch
 import pwndbg.gdblib.memory
 import pwndbg.gdblib.regs
-from pwndbg.lib.regs import amd64 as amd64_regset
+from pwndbg.lib.regs import amd64, arm, aarch64, i386
+
+
+# Grab frame values from pwntools. Offsets are defined as the offset to stack pointer when syscall instruction is called
+# Offsets and names are from Linux kernel source. x86_64, for example, from CONFIG_X86_64 struct rt_sigframe (Linux Kernel /arch/x86/include/asm/sigframe.h)
+SIGRETURN_FRAME_LAYOUTS: dict[str,list[Tuple[int, str]]] = {
+    "x86-64":[(-8, "&pretcode")] + list(pwnlib.rop.srop.registers["amd64"].items()),
+    "i386":list(pwnlib.rop.srop.registers["i386"].items()),
+    "aarch64":list(pwnlib.rop.srop.registers["aarch64"].items()),
+}
+
+SIGRETURN_CORE_REGISTER: dict[str, set[str]] = {
+    "x86-64":{ *amd64.gpr, amd64.frame, amd64.stack, amd64.pc },
+    "i386": { *i386.gpr,i386.frame, i386.stack, i386.pc },
+    "aarch64": { *aarch64.gpr, "sp", "pc"},
+}
+
 
 parser = argparse.ArgumentParser(description="Display the SigreturnFrame at the specific address")
 
@@ -40,56 +57,31 @@ parser.add_argument(
 
 @pwndbg.commands.ArgparsedCommand(parser)
 @pwndbg.commands.OnlyWhenRunning
-@pwndbg.commands.OnlyWithArch(["x86-64"])
+@pwndbg.commands.OnlyWithArch(["x86-64", "i386", "aarch64"])
 def sigreturn(address: int = None, display_all=False, print_address=False):
     address = pwndbg.gdblib.regs.sp if address is None else address
 
-    arch_name = pwndbg.gdblib.arch.name
-    if arch_name == "x86-64":
-        sigreturn_x86_64(address, display_all, print_address)
-
-
-SIGRETURN_FRAME_SIZE_x86_64 = 256
-
-# Grab frame values from pwntools. Offsets are defined as the offset to stack pointer when syscall instruction is called
-# Offsets and names are from "CONFIG_X86_64 struct rt_sigframe, Linux Kernel /arch/x86/include/asm/sigframe.h
-SIGRETURN_FRAME_LAYOUT_x86_64 = sorted(
-    [(-8, "&pretcode")] + list(pwnlib.rop.srop.SigreturnFrame(arch="amd64").registers.items())
-)
-
-# Core registers
-SIGRETURN_REGISTERS_x86_64 = {
-    *amd64_regset.gpr,
-    amd64_regset.frame,
-    amd64_regset.stack,
-    amd64_regset.pc,
-}
-
-
-def print_value(string: str, address: int, print_address):
-    addr = ""
-    if print_address:
-        addr = f"{M.get(address)}: "
-    print(f"{addr}{string}")
-
-
-def sigreturn_x86_64(address: int, display_all: bool, print_address: bool):
     ptr_size = pwndbg.gdblib.arch.ptrsize
 
+    frame_layout = SIGRETURN_FRAME_LAYOUTS[pwndbg.gdblib.arch.name]
+    core_registers = SIGRETURN_CORE_REGISTER[pwndbg.gdblib.arch.name]
+
     # Offset to the stack pointer where the frame values really begins. Start reading memory there.
     # Can be negative, 0, or positive
-    frame_start_offset = SIGRETURN_FRAME_LAYOUT_x86_64[0][0]
+    frame_start_offset = frame_layout[0][0]
 
-    mem = pwndbg.gdblib.memory.read(address + frame_start_offset, SIGRETURN_FRAME_SIZE_x86_64)
+    read_size = frame_layout[-1][0] - frame_start_offset + ptr_size
 
-    for stack_offset, reg in SIGRETURN_FRAME_LAYOUT_x86_64:
+    mem = pwndbg.gdblib.memory.read(address + frame_start_offset, read_size)
+
+    for stack_offset, reg in frame_layout:
         # Subtract the offset of start of frame, to get the correct offset into "mem"
         mem_offset = stack_offset - frame_start_offset
 
         regname = C.register(reg.ljust(4).upper())
         value = pwndbg.gdblib.arch.unpack(mem[mem_offset : mem_offset + ptr_size])
 
-        if reg in SIGRETURN_REGISTERS_x86_64:
+        if reg in core_registers:
             desc = pwndbg.chain.format(value)
 
             print_value(f"{regname} {desc}", address + stack_offset, print_address)
@@ -102,3 +94,10 @@ def sigreturn_x86_64(address: int, display_all: bool, print_address: bool):
 
         elif display_all:
             print_value(f"{reg} {M.get(value)}", address + stack_offset, print_address)
+
+
+def print_value(string: str, address: int, print_address):
+    addr = ""
+    if print_address:
+        addr = f"{M.get(address)}: "
+    print(f"{addr}{string}")