Handle multiple flag registers (#25)

* Handle multiple flag registers

* Explicit check for None return from SBFrame.register

Author: stephen-f0

arch/aarch64.py

@@ -1,6 +1,6 @@
 """aarch64 architecture definition."""
 
-from arch.base_arch import BaseArch
+from arch.base_arch import BaseArch, FlagRegister
 
 
 class Aarch64(BaseArch):
@@ -50,10 +50,8 @@ class Aarch64(BaseArch):
 
     gpr_key = "general"
 
-    flag_register = "cpsr"
-
     # Bitmasks used to extract flag bits from cpsr register value
-    flag_register_bit_masks = {
+    _cpsr_register_bit_masks = {
         "n": 0x80000000,
         "z": 0x40000000,
         "c": 0x20000000,
@@ -69,3 +67,7 @@ class Aarch64(BaseArch):
         "f": 0x40,
         "m": 0xF,
     }
+
+    flag_registers = [
+        FlagRegister("cpsr", _cpsr_register_bit_masks)
+    ]

arch/arm.py

@@ -1,6 +1,6 @@
 """arm architecture definition."""
 
-from arch.base_arch import BaseArch
+from arch.base_arch import BaseArch, FlagRegister
 
 
 class Arm(BaseArch):
@@ -31,10 +31,8 @@ class Arm(BaseArch):
 
     gpr_key = "general"
 
-    flag_register = "cpsr"
-
     # Bitmasks used to extract flag bits from cpsr register value
-    flag_register_bit_masks = {
+    _cpsr_register_bit_masks = {
         "n": 0x80000000,
         "z": 0x40000000,
         "c": 0x20000000,
@@ -48,3 +46,7 @@ class Arm(BaseArch):
         "f": 0x40,
         "t": 0x20,
     }
+
+    flag_registers = [
+        FlagRegister("cpsr", _cpsr_register_bit_masks)
+    ]

arch/base_arch.py

@@ -1,9 +1,17 @@
 """Base arch abstract class definition."""
 
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from typing import Dict, List
 
 
+@dataclass
+class FlagRegister:
+    """FlagRegister dataclass to store register name / bitmask associations"""
+    name: str
+    bit_masks: Dict[str, int]
+
+
 class BaseArch(ABC):
     """BaseArch abstract class definition."""
 
@@ -24,10 +32,5 @@ def gpr_key(self) -> str:
 
     @property
     @abstractmethod
-    def flag_register(self) -> str:
-        """Flag register property"""
-
-    @property
-    @abstractmethod
-    def flag_register_bit_masks(self) -> Dict[str, int]:
-        """Flag register bit mask property"""
+    def flag_registers(self) -> List[FlagRegister]:
+        """List of flag registers with associated bit masks"""

arch/i386.py

@@ -1,5 +1,5 @@
 """i386 architecture definition."""
-from arch.base_arch import BaseArch
+from arch.base_arch import BaseArch, FlagRegister
 
 
 class I386(BaseArch):
@@ -29,10 +29,8 @@ class I386(BaseArch):
 
     gpr_key = "general purpose"
 
-    flag_register = "eflags"
-
     # Bitmasks used to extract flag bits from eflags register value
-    flag_register_bit_masks = {
+    _eflags_register_bit_masks = {
         "zero": 0x40,
         "carry": 0x1,
         "parity": 0x4,
@@ -46,3 +44,7 @@ class I386(BaseArch):
         "virtual8086": 0x20000,
         "identification": 0x200000,
     }
+
+    flag_registers = [
+        FlagRegister("eflags", _eflags_register_bit_masks)
+    ]

arch/x86_64.py

@@ -1,5 +1,5 @@
 """x86_64 architecture definition."""
-from arch.base_arch import BaseArch
+from arch.base_arch import BaseArch, FlagRegister
 
 
 class X86_64(BaseArch):
@@ -31,10 +31,8 @@ class X86_64(BaseArch):
 
     gpr_key = "general purpose"
 
-    flag_register = "rflags"
-
-    # Bitmasks used to extract flag bits from rflags register value
-    flag_register_bit_masks = {
+    # Bitmasks used to extract flag bits from eflags register value
+    _eflag_register_bit_masks = {
         "zero": 0x40,
         "carry": 0x1,
         "parity": 0x4,
@@ -48,3 +46,10 @@ class X86_64(BaseArch):
         "virtualx86": 0x20000,
         "identification": 0x200000,
     }
+
+    # Whether LLDB exposes eflags or rflags varies depending on the platform
+    # rflags and eflags bit masks are identical for the lower 32-bits
+    flag_registers = [
+        FlagRegister("rflags", _eflag_register_bit_masks),
+        FlagRegister("eflags", _eflag_register_bit_masks)
+    ]

common/context_handler.py

@@ -15,7 +15,7 @@
 )
 
 from arch import get_arch
-from arch.base_arch import BaseArch
+from arch.base_arch import BaseArch, FlagRegister
 from common.constants import GLYPHS, TERM_COLOURS
 from common.util import (
     attempt_to_read_string_from_memory,
@@ -155,25 +155,22 @@ def print_register(self, register: SBValue) -> None:
 
         print(line)
 
-    def print_flags_register(self, flag_register: SBValue) -> None:
+    def print_flags_register(self, flag_register: FlagRegister) -> None:
         """Format and print the contents of the flag register."""
+        flag_value = self.frame.register[flag_register.name].GetValueAsUnsigned()
 
-        if (
-            self.old_registers.get(self.arch().flag_register)
-            == flag_register.GetValueAsUnsigned()
-        ):
+        if self.old_registers.get(flag_register.name) == flag_value:
             # No change
             highlight = TERM_COLOURS.BLUE
         else:
             # Change and highlight
             highlight = TERM_COLOURS.RED
 
-        flag_value = flag_register.GetValueAsUnsigned()
-        line = f"{highlight.value}{flag_register.GetName().ljust(7)}{TERM_COLOURS.ENDC.value}: ["
+        line = f"{highlight.value}{flag_register.name.ljust(7)}{TERM_COLOURS.ENDC.value}: ["
         line += " ".join(
             [
                 name.upper() if flag_value & bitmask else name
-                for name, bitmask in self.arch().flag_register_bit_masks.items()
+                for name, bitmask in flag_register.bit_masks.items()
             ]
         )
         line += "]"
@@ -204,7 +201,9 @@ def display_registers(self) -> None:
         for reg in get_registers(self.frame, self.arch().gpr_key):
             if reg.GetName() in self.arch().gpr_registers:
                 self.print_register(reg)
-        self.print_flags_register(self.frame.register[self.arch.flag_register])
+        for flag_register in self.arch.flag_registers:
+            if self.frame.register[flag_register.name] is not None:
+                self.print_flags_register(flag_register)
 
     def display_stack(self) -> None:
         """Print information about the contents of the top of the stack"""