From 092faaff9b3868c8f8eeab13f4506968be6546d1 Mon Sep 17 00:00:00 2001
From: Rusty Wagner <rusty@vector35.com>
Date: Mon, 23 Oct 2017 21:45:37 -0400
Subject: Top relative register stack access as a normal register

---
 python/architecture.py | 28 +++++++++++++++++++++++-----
 python/function.py     |  3 ++-
 2 files changed, 25 insertions(+), 6 deletions(-)

(limited to 'python')

diff --git a/python/architecture.py b/python/architecture.py
index ea9d40ab..ac4b1d6e 100644
--- a/python/architecture.py
+++ b/python/architecture.py
@@ -227,10 +227,13 @@ class Architecture(object):
 				name = core.BNGetArchitectureRegisterStackName(self.handle, regs[i])
 				info = core.BNGetArchitectureRegisterStackInfo(self.handle, regs[i])
 				storage = []
-				for j in xrange(0, info.count):
+				for j in xrange(0, info.storageCount):
 					storage.append(core.BNGetArchitectureRegisterName(self.handle, info.firstStorageReg + j))
+				top_rel = []
+				for j in xrange(0, info.topRelativeCount):
+					top_rel.append(core.BNGetArchitectureRegisterName(self.handle, info.firstTopRelativeReg + j))
 				top = core.BNGetArchitectureRegisterName(self.handle, info.stackTopReg)
-				self.reg_stacks[name] = function.RegisterStackInfo(storage, top)
+				self.reg_stacks[name] = function.RegisterStackInfo(storage, top_rel, top)
 			core.BNFreeRegisterList(regs)
 		else:
 			startup._init_plugins()
@@ -313,6 +316,11 @@ class Architecture(object):
 					self._regs_by_index[reg_index] = reg
 					self.regs[reg].index = reg_index
 					reg_index += 1
+				for reg in info.top_relative_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
@@ -820,17 +828,27 @@ class Architecture(object):
 		try:
 			if reg_stack not in self._reg_stacks_by_index:
 				result[0].firstStorageReg = 0
-				result[0].count = 0
+				result[0].firstTopRelativeReg = 0
+				result[0].storageCount = 0
+				result[0].topRelativeCount = 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].storageCount = len(info.storage_regs)
+			if len(info.top_relative_regs) > 0:
+				result[0].firstTopRelativeReg = self._all_regs[info.top_relative_regs[0]]
+				result[0].topRelativeCount = len(info.top_relative_regs)
+			else:
+				result[0].firstTopRelativeReg = 0
+				result[0].topRelativeCount = 0
 			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].firstTopRelativeReg = 0
+			result[0].storageCount = 0
+			result[0].topRelativeCount = 0
 			result[0].stackTopReg = 0
 
 	def _assemble(self, ctxt, code, addr, result, errors):
diff --git a/python/function.py b/python/function.py
index 13979e37..8e058368 100644
--- a/python/function.py
+++ b/python/function.py
@@ -1788,8 +1788,9 @@ class RegisterInfo(object):
 
 
 class RegisterStackInfo(object):
-	def __init__(self, storage_regs, stack_top_reg):
+	def __init__(self, storage_regs, top_relative_regs, stack_top_reg):
 		self.storage_regs = storage_regs
+		self.top_relative_regs = top_relative_regs
 		self.stack_top_reg = stack_top_reg
 
 	def __repr__(self):
-- 
cgit v1.3.1