Merge pull request #920 from elicn/reg-improv

Refactored QlRegisterManager

Author: xwings

qiling/arch/arm.py

@@ -24,7 +24,6 @@ def __init__(self, ql: Qiling):
         for reg_maper in reg_maps:
             self.ql.reg.expand_mapping(reg_maper)
 
-        self.ql.reg.create_reverse_mapping()
         self.ql.reg.register_sp(reg_map["sp"])
         self.ql.reg.register_pc(reg_map["pc"])
 

qiling/arch/arm64.py

@@ -23,7 +23,6 @@ def __init__(self, ql: Qiling):
         for reg_maper in reg_maps:
             self.ql.reg.expand_mapping(reg_maper)
 
-        self.ql.reg.create_reverse_mapping()
         self.ql.reg.register_sp(reg_map["sp"])
         self.ql.reg.register_pc(reg_map["pc"])
 

qiling/arch/mips.py

@@ -24,7 +24,6 @@ def __init__(self, ql: Qiling):
         for reg_maper in reg_maps:
             self.ql.reg.expand_mapping(reg_maper)
 
-        self.ql.reg.create_reverse_mapping()
         self.ql.reg.register_sp(reg_map["sp"])
         self.ql.reg.register_pc(reg_map["pc"])
 

qiling/arch/register.py

@@ -3,127 +3,157 @@
 # Cross Platform and Multi Architecture Advanced Binary Emulation Framework
 #
 
-class QlRegisterManager():
-    """
-    This class exposes the ql.reg features that allows you to directly access
+from typing import Any, Mapping, MutableMapping, Union
+
+from qiling import Qiling
+
+class QlRegisterManager:
+    """This class exposes the ql.reg features that allows you to directly access
     or assign values to CPU registers of a particular architecture.
 
     Registers exposed are listed in the *_const.py files in the respective
     arch directories and are mapped to Unicorn Engine's definitions
     """
-    def __init__(self, ql):
-        self.register_mapping = {}
-        self.reverse_mapping = {}
+
+    def __init__(self, ql: Qiling):
+        # this funny way of initialization is used to avoid calling self setattr and
+        # getattr upon init. if it did, it would go into an endless recursion
+        self.register_mapping: MutableMapping[str, int]
+        super().__setattr__('register_mapping', {})
+
         self.ql = ql
         self.uc_pc = 0
         self.uc_sp = 0
-        
 
-    def __getattribute__(self, name):
+    def __getattr__(self, name: str) -> Any:
         name = name.lower()
-        if name in ("register_mapping", "ql", "uc_pc", "uc_sp"):
-            return super(QlRegisterManager, self).__getattribute__(name)
 
-        elif name in self.register_mapping:
+        if name in self.register_mapping:
             return self.ql.uc.reg_read(self.register_mapping[name])
 
-        return super(QlRegisterManager, self).__getattribute__(name)
+        else:
+            return super().__getattribute__(name)
 
 
-    def __setattr__(self, name, value):
-        name=name.lower()
-        if name in ("register_mapping", "ql", "uc_pc", "uc_sp"):
-            super(QlRegisterManager, self).__setattr__(name, value)
+    def __setattr__(self, name: str, value: Any):
+        name = name.lower()
 
-        elif name in self.register_mapping:
+        if name in self.register_mapping:
             self.ql.uc.reg_write(self.register_mapping[name], value)
+
         else:
-            super(QlRegisterManager, self).__setattr__(name, value)
+            super().__setattr__(name, value)
 
 
-    def expand_mapping(self, expanded_map):
-        self.register_mapping = {**self.register_mapping, **expanded_map}
+    def expand_mapping(self, extra: Mapping[str, int]) -> None:
+        """Expand registers mapping with additional ones.
+        """
+
+        self.register_mapping.update(extra)
 
 
     # read register
-    def read(self, register):
-        if isinstance(register, str):
-            register = self.register_mapping.get(register.lower(), None)
+    def read(self, register: Union[str, int]):
+        """Read a register value.
+        """
+
+        if type(register) is str:
+            register = self.register_mapping[register.lower()]
+
         return self.ql.uc.reg_read(register)
 
 
-    def write(self, register, value):
-        if isinstance(register, str):
-            register = self.register_mapping.get(register.lower(), None) 
+    def write(self, register: Union[str, int], value: int) -> None:
+        """Write a register value.
+        """
+
+        if type(register) is str:
+            register = self.register_mapping[register.lower()]
+
         return self.ql.uc.reg_write(register, value)
 
 
-    def msr(self, msr, addr= None):
-        if not addr:
+    def msr(self, msr: int, value: int = None):
+        """Read or write a model-specific register (MSR) value.
+        Intel architecture only
+        """
+
+        if value is None:
             return self.ql.uc.msr_read(msr)
-        else:
-            self.ql.uc.msr_write(msr, addr)
 
+        self.ql.uc.msr_write(msr, value)
 
-    # ql.reg.save
-    def save(self):
-        reg_dict = {}
-        for reg in self.register_mapping:
-            reg_v = self.read(reg)
-            reg_dict[reg] = reg_v
-        return reg_dict
 
+    def save(self) -> MutableMapping[str, Any]:
+        """Save CPU context.
+        """
 
-    # ql.reg.restore
-    def restore(self, value = {}):
-        for reg in self.register_mapping:
-            reg_v= value[reg]
-            self.write(reg, reg_v)
+        return dict((reg, self.read(reg)) for reg in self.register_mapping)
 
 
-    # ql.reg.bit() - Register bit
-    #FIXME: This needs to be implemented for all archs
-    def bit(self, uc_reg):
-        return self.ql.arch.get_reg_bit(uc_reg)
+    def restore(self, context: MutableMapping[str, Any] = {}) -> None:
+        """Restore CPU context.
+        """
+
+        for reg, val in context.items():
+            self.write(reg, val)
+
+
+    # TODO: This needs to be implemented for all archs
+    def bit(self, reg: Union[str, int]) -> int:
+        """Get register size in bits.
+        """
+
+        if type(reg) is str:
+            reg = self.register_mapping[reg]
+
+        return self.ql.arch.get_reg_bit(reg)
 
 
     # Generic methods to get SP and IP across Arch's #
     # These functions should only be used if the     #
     # caller is dealing with multiple Arch's         #
-    def register_sp(self, sp_id):
+    def register_sp(self, sp_id: int):
         self.uc_sp = sp_id
 
 
-    def register_pc(self, pc_id):
+    def register_pc(self, pc_id: int):
         self.uc_pc = pc_id
 
 
     @property
-    def arch_pc(self):
+    def arch_pc(self) -> int:
+        """Get the value of the architectural program counter register.
+        """
+
         return self.ql.uc.reg_read(self.uc_pc)
 
 
     @arch_pc.setter
-    def arch_pc(self, value):
+    def arch_pc(self, value: int) -> None:
+        """Set the value of the architectural program counter register.
+        """
+
         return self.ql.uc.reg_write(self.uc_pc, value)
 
     @property
-    def arch_pc_name(self):
-        return self.ql.reg.reverse_mapping[self.uc_pc]
+    def arch_pc_name(self) -> str:
+        """Get the architectural program counter register name.
+        """
+
+        return next(k for k, v in self.register_mapping.items() if v == self.uc_pc)
 
     @property
-    def arch_sp(self):
+    def arch_sp(self) -> int:
+        """Get the value of the architectural stack pointer register.
+        """
+
         return self.ql.uc.reg_read(self.uc_sp)
 
 
     @arch_sp.setter
-    def arch_sp(self, value):
-        return self.ql.uc.reg_write(self.uc_sp, value)
-
+    def arch_sp(self, value: int) -> None:
+        """Set the value of the architectural stack pointer register.
+        """
 
-    def get_uc_reg(self, uc_reg_name):
-        return self.register_mapping.get(uc_reg_name, None)
-
-
-    def create_reverse_mapping(self):
-        self.reverse_mapping = {v:k for k, v in self.register_mapping.items()}
+        return self.ql.uc.reg_write(self.uc_sp, value)

qiling/arch/x86.py

@@ -4,7 +4,6 @@
 #
 
 from struct import pack
-from typing import Union
 
 from unicorn import Uc, UC_ARCH_X86, UC_MODE_16, UC_MODE_32, UC_MODE_64
 from capstone import Cs, CS_ARCH_X86, CS_MODE_16, CS_MODE_32, CS_MODE_64
@@ -18,10 +17,7 @@
 class QlArchIntel(QlArch):
 
     # TODO: generalize this
-    def get_reg_bit(self, register: Union[str, int]) -> int:
-        if type(register) is str:
-            register = self.ql.reg.get_uc_reg(register)
-
+    def get_reg_bit(self, register: int) -> int:
         # all regs in reg_map_misc are 16 bits except of eflags
         if register == UC_X86_REG_EFLAGS:
             return self.ql.archbit
@@ -52,7 +48,6 @@ def __init__(self, ql: Qiling):
         for reg_maper in reg_maps:
             self.ql.reg.expand_mapping(reg_maper)
 
-        self.ql.reg.create_reverse_mapping()
         self.ql.reg.register_pc(reg_map_16["sp"])
         self.ql.reg.register_sp(reg_map_16["ip"])
 
@@ -87,7 +82,6 @@ def __init__(self, ql: Qiling):
         for reg_maper in reg_maps:
             self.ql.reg.expand_mapping(reg_maper)
 
-        self.ql.reg.create_reverse_mapping()
         self.ql.reg.register_sp(reg_map_32["esp"])
         self.ql.reg.register_pc(reg_map_32["eip"])
 
@@ -127,7 +121,6 @@ def __init__(self, ql: Qiling):
         for reg_maper in reg_maps:
             self.ql.reg.expand_mapping(reg_maper)
 
-        self.ql.reg.create_reverse_mapping()
         self.ql.reg.register_sp(reg_map_64["rsp"])
         self.ql.reg.register_pc(reg_map_64["rip"])