Adding support for register stacks in IL (for x87)

1 files changed, 101 insertions, 0 deletions

diff --git a/python/architecture.py b/python/architecture.py
index 72403fec..289b0abd 100644
--- a/python/architecture.py
+++ b/python/architecture.py
@@ -122,6 +122,7 @@ class Architecture(object):
 	flag_roles = {}
 	flags_required_for_flag_condition = {}
 	flags_written_by_flag_write_type = {}
+	reg_stacks = {}
 	__metaclass__ = _ArchitectureMetaClass
 	next_address = 0
 
@@ -216,6 +217,19 @@ class Architecture(object):
 			for i in xrange(0, count.value):
 				self.global_regs.append(core.BNGetArchitectureRegisterName(self.handle, regs[i]))
 			core.BNFreeRegisterList(regs)
+
+			count = ctypes.c_ulonglong()
+			regs = core.BNGetAllArchitectureRegisterStacks(self.handle, count)
+			self.__dict__["reg_stacks"] = {}
+			for i in xrange(0, count.value):
+				name = core.BNGetArchitectureRegisterStackName(self.handle, regs[i])
+				info = core.BNGetArchitectureRegisterStackInfo(self.handle, regs[i])
+				storage = []
+				for j in xrange(0, info.count):
+					storage.append(core.BNGetArchitectureRegisterName(self.handle, info.firstStorageReg + j))
+				top = core.BNGetArchitectureRegisterName(self.handle, info.stackTopReg)
+				self.reg_stacks[name] = function.RegisterStackInfo(storage, top)
+			core.BNFreeRegisterList(regs)
 		else:
 			startup._init_plugins()
 
@@ -259,6 +273,9 @@ class Architecture(object):
 				self._get_stack_pointer_register)
 			self._cb.getLinkRegister = self._cb.getLinkRegister.__class__(self._get_link_register)
 			self._cb.getGlobalRegisters = self._cb.getGlobalRegisters.__class__(self._get_global_registers)
+			self._cb.getRegisterStackName = self._cb.getRegisterStackName.__class__(self._get_register_stack_name)
+			self._cb.getAllRegisterStacks = self._cb.getAllRegisterStacks.__class__(self._get_all_register_stacks)
+			self._cb.getRegisterStackInfo = self._cb.getRegisterStackInfo.__class__(self._get_register_stack_info)
 			self._cb.assemble = self._cb.assemble.__class__(self._assemble)
 			self._cb.isNeverBranchPatchAvailable = self._cb.isNeverBranchPatchAvailable.__class__(
 				self._is_never_branch_patch_available)
@@ -280,6 +297,25 @@ class Architecture(object):
 			self._regs_by_index = {}
 			self.__dict__["regs"] = self.__class__.regs
 			reg_index = 0
+
+			# Registers used for storage in register stacks must be sequential, so allocate these in order first
+			self._all_reg_stacks = {}
+			self._reg_stacks_by_index = {}
+			self.__dict__["reg_stacks"] = self.__class__.reg_stacks
+			reg_stack_index = 0
+			for reg_stack in self.reg_stacks:
+				info = self.reg_stacks[reg_stack]
+				for reg in info.storage_regs:
+					self._all_regs[reg] = reg_index
+					self._regs_by_index[reg_index] = reg
+					self.regs[reg].index = reg_index
+					reg_index += 1
+				if reg_stack not in self._all_reg_stacks:
+					self._all_reg_stacks[reg_stack] = reg_stack_index
+					self._reg_stacks_by_index[reg_stack_index] = reg_stack
+					self.reg_stacks[reg_stack].index = reg_stack_index
+					reg_stack_index += 1
+
 			for reg in self.regs:
 				info = self.regs[reg]
 				if reg not in self._all_regs:
@@ -744,6 +780,47 @@ class Architecture(object):
 			count[0] = 0
 			return None
 
+	def _get_register_stack_name(self, ctxt, reg_stack):
+		try:
+			if reg_stack in self._reg_stacks_by_index:
+				return core.BNAllocString(self._reg_stacks_by_index[reg_stack])
+			return core.BNAllocString("")
+		except (KeyError, OSError):
+			log.log_error(traceback.format_exc())
+			return core.BNAllocString("")
+
+	def _get_all_register_stacks(self, ctxt, count):
+		try:
+			regs = self._reg_stacks_by_index.keys()
+			count[0] = len(regs)
+			reg_buf = (ctypes.c_uint * len(regs))()
+			for i in xrange(0, len(regs)):
+				reg_buf[i] = regs[i]
+			result = ctypes.cast(reg_buf, ctypes.c_void_p)
+			self._pending_reg_lists[result.value] = (result, reg_buf)
+			return result.value
+		except KeyError:
+			log.log_error(traceback.format_exc())
+			count[0] = 0
+			return None
+
+	def _get_register_stack_info(self, ctxt, reg_stack, result):
+		try:
+			if reg_stack not in self._reg_stacks_by_index:
+				result[0].firstStorageReg = 0
+				result[0].count = 0
+				result[0].stackTopReg = 0
+				return
+			info = self.__class__.regs[self._reg_stacks_by_index[reg_stack]]
+			result[0].firstStorageReg = self._all_regs[info.storage_regs[0]]
+			result[0].count = len(info.storage_regs)
+			result[0].stackTopReg = self._all_regs[info.stack_top_reg]
+		except KeyError:
+			log.log_error(traceback.format_exc())
+			result[0].firstStorageReg = 0
+			result[0].count = 0
+			result[0].stackTopReg = 0
+
 	def _assemble(self, ctxt, code, addr, result, errors):
 		try:
 			data, error_str = self.perform_assemble(code, addr)
@@ -1251,6 +1328,23 @@ class Architecture(object):
 		"""
 		return core.BNGetArchitectureRegisterName(self.handle, reg)
 
+	def get_reg_stack_name(self, reg_stack):
+		"""
+		``get_reg_stack_name`` gets a register stack name from a register stack number.
+
+		:param int reg_stack: register stack number
+		:return: the corresponding register string
+		:rtype: str
+		"""
+		return core.BNGetArchitectureRegisterStackName(self.handle, reg_stack)
+
+	def get_reg_stack_for_reg(self, reg):
+		reg = self.get_reg_index(reg)
+		result = core.BNGetArchitectureRegisterStackForRegister(self.handle, reg)
+		if result == 0xffffffff:
+			return None
+		return self.get_reg_stack_name(result)
+
 	def get_flag_name(self, flag):
 		"""
 		``get_flag_name`` gets a flag name from a flag number.
@@ -1268,6 +1362,13 @@ class Architecture(object):
 			return reg.index
 		return reg
 
+	def get_reg_stack_index(self, reg_stack):
+		if isinstance(reg_stack, str):
+			return self.reg_stacks[reg_stack].index
+		elif isinstance(reg_stack, lowlevelil.ILRegisterStack):
+			return reg_stack.index
+		return reg_stack
+
 	def get_flag_index(self, flag):
 		if isinstance(flag, str):
 			return self._flags[flag]