Refactor register

Author: elicn

qiling/arch/register.py

@@ -3,126 +3,149 @@
 # 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 = {}
+
+    __priv_members = ('register_mapping', 'ql', 'uc_pc', 'uc_sp')
+
+    def __init__(self, ql: Qiling):
+        self.register_mapping: MutableMapping[str, int] = {}
+        self.reverse_mapping: Mapping[int, str] = {}
         self.ql = ql
         self.uc_pc = 0
         self.uc_sp = 0
-        
 
-    def __getattribute__(self, name):
+    def __getattribute__(self, name: str) -> Any:
         name = name.lower()
-        if name in ("register_mapping", "ql", "uc_pc", "uc_sp"):
-            return super(QlRegisterManager, self).__getattribute__(name)
+
+        if name in QlRegisterManager.__priv_members:
+            return super().__getattribute__(name)
 
         elif name in self.register_mapping:
             return self.ql.uc.reg_read(self.register_mapping[name])
 
-        return super(QlRegisterManager, self).__getattribute__(name)
+        return super().__getattribute__(name)
+
 
+    def __setattr__(self, name: str, value: Any):
+        name = name.lower()
 
-    def __setattr__(self, name, value):
-        name=name.lower()
-        if name in ("register_mapping", "ql", "uc_pc", "uc_sp"):
-            super(QlRegisterManager, self).__setattr__(name, value)
+        if name in QlRegisterManager.__priv_members:
+            super().__setattr__(name, value)
 
         elif 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, extra: Mapping[str, int]) -> None:
+        """Expand registers mapping with additional ones.
+        """
 
-    def expand_mapping(self, expanded_map):
-        self.register_mapping = {**self.register_mapping, **expanded_map}
+        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)
+
+
+    def save(self) -> MutableMapping[str, Any]:
+        """Save CPU context.
+        """
+
+        return dict((reg, self.read(reg)) for reg in self.register_mapping)
 
 
-    # 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 restore(self, context: MutableMapping[str, Any] = {}) -> None:
+        """Restore CPU context.
+        """
 
+        for reg, val in context.items():
+            self.write(reg, val)
 
-    # ql.reg.restore
-    def restore(self, value = {}):
-        for reg in self.register_mapping:
-            reg_v= value[reg]
-            self.write(reg, reg_v)
 
+    # TODO: This needs to be implemented for all archs
+    def bit(self, reg: Union[str, int]) -> int:
+        """Get register size in bits.
+        """
 
-    # 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)
+        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:
         return self.ql.uc.reg_read(self.uc_pc)
 
 
     @arch_pc.setter
-    def arch_pc(self, value):
+    def arch_pc(self, value: int) -> None:
         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:
+        return self.reverse_mapping[self.uc_pc]
 
     @property
-    def arch_sp(self):
+    def arch_sp(self) -> int:
         return self.ql.uc.reg_read(self.uc_sp)
 
 
     @arch_sp.setter
-    def arch_sp(self, value):
+    def arch_sp(self, value: int) -> None:
         return self.ql.uc.reg_write(self.uc_sp, value)
 
 
-    def get_uc_reg(self, uc_reg_name):
-        return self.register_mapping.get(uc_reg_name, None)
+    def get_uc_reg(self, reg_name: str) -> int:
+        return self.register_mapping[reg_name]
 
 
     def create_reverse_mapping(self):