Register annotation handlers for armcm (#2388)

pwndbg/emu/emulator.py

@@ -82,6 +82,7 @@ arch_to_UC = {
     "mips": U.UC_ARCH_MIPS,
     "sparc": U.UC_ARCH_SPARC,
     "arm": U.UC_ARCH_ARM,
+    "armcm": U.UC_ARCH_ARM,
     "aarch64": U.UC_ARCH_ARM64,
     # 'powerpc': U.UC_ARCH_PPC,
     "rv32": U.UC_ARCH_RISCV,
@@ -95,6 +96,7 @@ arch_to_UC_consts = {
     "mips": parse_consts(U.mips_const),
     "sparc": parse_consts(U.sparc_const),
     "arm": parse_consts(U.arm_const),
+    "armcm": parse_consts(U.arm_const),
     "aarch64": parse_consts(U.arm64_const),
     "rv32": parse_consts(U.riscv_const),
     "rv64": parse_consts(U.riscv_const),
@@ -110,6 +112,7 @@ arch_to_reg_const_map = {
     "mips": create_reg_to_const_map(arch_to_UC_consts["mips"]),
     "sparc": create_reg_to_const_map(arch_to_UC_consts["sparc"]),
     "arm": create_reg_to_const_map(arch_to_UC_consts["arm"]),
+    "armcm": create_reg_to_const_map(arch_to_UC_consts["armcm"]),
     "aarch64": create_reg_to_const_map(
         arch_to_UC_consts["aarch64"], {"CPSR": U.arm64_const.UC_ARM64_REG_NZCV}
     ),
@@ -171,11 +174,18 @@ arch_to_SYSCALL = {
     U.UC_ARCH_RISCV: [C.riscv_const.RISCV_INS_ECALL],
 }
 
+ARM_BANNED_INSTRUCTIONS = {
+    C.arm.ARM_INS_MRC,
+    C.arm.ARM_INS_MRRC,
+    C.arm.ARM_INS_MRC2,
+    C.arm.ARM_INS_MRRC2,
+}
 # We stop emulation when hitting these instructions, since they depend on co-processors or other information
 # unavailable to the emulator
 BANNED_INSTRUCTIONS = {
     "mips": {C.mips.MIPS_INS_RDHWR},
-    "arm": {C.arm.ARM_INS_MRC, C.arm.ARM_INS_MRRC, C.arm.ARM_INS_MRC2, C.arm.ARM_INS_MRRC2},
+    "arm": ARM_BANNED_INSTRUCTIONS,
+    "armcm": ARM_BANNED_INSTRUCTIONS,
 }
 
 # https://github.com/unicorn-engine/unicorn/issues/550

pwndbg/gdblib/disasm/arm.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from typing import Callable
 from typing import Dict
+from typing import Literal
 
 from capstone import *  # noqa: F403
 from capstone.arm import *  # noqa: F403
@@ -72,10 +73,13 @@ ARM_MATH_INSTRUCTIONS = {
 }
 
 
+# This class enhances both ARM A-profile and ARM M-profile (Cortex-M)
 class DisassemblyAssistant(pwndbg.gdblib.disasm.arch.DisassemblyAssistant):
-    def __init__(self, architecture: str) -> None:
+    def __init__(self, architecture: str, flags_reg: Literal["cpsr", "xpsr"]) -> None:
         super().__init__(architecture)
 
+        self.flags_reg = flags_reg
+
         self.annotation_handlers: Dict[int, Callable[[PwndbgInstruction, Emulator], None]] = {
             # MOV
             ARM_INS_MOV: self._common_move_annotator,
@@ -85,13 +89,13 @@ class DisassemblyAssistant(pwndbg.gdblib.disasm.arch.DisassemblyAssistant):
             # MOVN
             ARM_INS_MVN: self._common_generic_register_destination,
             # CMP
-            ARM_INS_CMP: self._common_cmp_annotator_builder("cpsr", "-"),
+            ARM_INS_CMP: self._common_cmp_annotator_builder(flags_reg, "-"),
             # CMN
-            ARM_INS_CMN: self._common_cmp_annotator_builder("cpsr", "+"),
+            ARM_INS_CMN: self._common_cmp_annotator_builder(flags_reg, "+"),
             # TST (bitwise "and")
-            ARM_INS_TST: self._common_cmp_annotator_builder("cpsr", "&"),
+            ARM_INS_TST: self._common_cmp_annotator_builder(flags_reg, "&"),
             # TEQ (bitwise exclusive "or")
-            ARM_INS_TEQ: self._common_cmp_annotator_builder("cpsr", "^"),
+            ARM_INS_TEQ: self._common_cmp_annotator_builder(flags_reg, "^"),
         }
 
     @override
@@ -143,11 +147,7 @@ class DisassemblyAssistant(pwndbg.gdblib.disasm.arch.DisassemblyAssistant):
         if instruction.address != pwndbg.gdblib.regs.pc:
             return InstructionCondition.UNDETERMINED
 
-        value = (
-            pwndbg.gdblib.regs.cpsr
-            if pwndbg.gdblib.arch.current == "arm"
-            else pwndbg.gdblib.regs.xpsr
-        )
+        value = pwndbg.gdblib.regs[self.flags_reg]
 
         N = (value >> 31) & 1
         Z = (value >> 30) & 1
@@ -291,4 +291,6 @@ class DisassemblyAssistant(pwndbg.gdblib.disasm.arch.DisassemblyAssistant):
         return target
 
 
-assistant = DisassemblyAssistant("arm")
+# Register the assistant for both ARM A-profile and ARM M-profile
+assistant = DisassemblyAssistant("arm", "cpsr")
+assistant = DisassemblyAssistant("armcm", "xpsr")