Extending register value system to include ranges and lookup table values

4 files changed, 140 insertions, 33 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index c0c9708c..c9c19953 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -1312,6 +1312,21 @@ namespace BinaryNinja
 		ArchAndAddr(Architecture* a, uint64_t addr): arch(a), address(addr) {}
 	};
 
+	struct LookupTableEntry
+	{
+		std::vector<int64_t> fromValues;
+		int64_t toValue;
+	};
+
+	struct RegisterValue
+	{
+		BNRegisterValueType state;
+		uint32_t reg; // For EntryValue and OffsetFromEntryValue, the original input register
+		int64_t value; // Offset for OffsetFromEntryValue, StackFrameOffset or RangeValue, value of register for ConstantValue
+		uint64_t rangeStart, rangeEnd, rangeStep; // Range of register, inclusive
+		std::vector<LookupTableEntry> table;
+	};
+
 	class FunctionGraph;
 
 	class Function: public CoreRefCountObject<BNFunction, BNNewFunctionReference, BNFreeFunction>
@@ -1337,10 +1352,10 @@ namespace BinaryNinja
 		Ref<LowLevelILFunction> GetLowLevelIL() const;
 		size_t GetLowLevelILForInstruction(Architecture* arch, uint64_t addr);
 		std::vector<size_t> GetLowLevelILExitsForInstruction(Architecture* arch, uint64_t addr);
-		BNRegisterValue GetRegisterValueAtInstruction(Architecture* arch, uint64_t addr, uint32_t reg);
-		BNRegisterValue GetRegisterValueAfterInstruction(Architecture* arch, uint64_t addr, uint32_t reg);
-		BNRegisterValue GetRegisterValueAtLowLevelILInstruction(size_t i, uint32_t reg);
-		BNRegisterValue GetRegisterValueAfterLowLevelILInstruction(size_t i, uint32_t reg);
+		RegisterValue GetRegisterValueAtInstruction(Architecture* arch, uint64_t addr, uint32_t reg);
+		RegisterValue GetRegisterValueAfterInstruction(Architecture* arch, uint64_t addr, uint32_t reg);
+		RegisterValue GetRegisterValueAtLowLevelILInstruction(size_t i, uint32_t reg);
+		RegisterValue GetRegisterValueAfterLowLevelILInstruction(size_t i, uint32_t reg);
 		std::vector<uint32_t> GetRegistersReadByInstruction(Architecture* arch, uint64_t addr);
 		std::vector<uint32_t> GetRegistersWrittenByInstruction(Architecture* arch, uint64_t addr);
 		std::vector<StackVariableReference> GetStackVariablesReferencedByInstruction(Architecture* arch, uint64_t addr);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 51a3afdf..0eff6aaf 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -358,14 +358,27 @@ extern "C"
 		OffsetFromEntryValue,
 		ConstantValue,
 		StackFrameOffset,
-		UndeterminedValue
+		UndeterminedValue,
+		OffsetFromUndeterminedValue,
+		SignedRangeValue,
+		UnsignedRangeValue,
+		LookupTableValue
+	};
+
+	struct BNLookupTableEntry
+	{
+		int64_t* fromValues;
+		size_t fromCount;
+		int64_t toValue;
 	};
 
 	struct BNRegisterValue
 	{
 		BNRegisterValueType state;
 		uint32_t reg; // For EntryValue and OffsetFromEntryValue, the original input register
-		int64_t value; // Offset for OffsetFromEntryValue or StackFrameOffset, value of register for ConstantValue
+		int64_t value; // Offset for OffsetFromEntryValue, StackFrameOffset or RangeValue, value of register for ConstantValue
+		uint64_t rangeStart, rangeEnd, rangeStep; // Range of register, inclusive
+		BNLookupTableEntry* table; // Number of entries in rangeEnd
 	};
 
 	struct BNRegisterOrConstant
@@ -1102,6 +1115,7 @@ extern "C"
 	                                                                      uint32_t reg);
 	BINARYNINJACOREAPI BNRegisterValue BNGetRegisterValueAtLowLevelILInstruction(BNFunction* func, size_t i, uint32_t reg);
 	BINARYNINJACOREAPI BNRegisterValue BNGetRegisterValueAfterLowLevelILInstruction(BNFunction* func, size_t i, uint32_t reg);
+	BINARYNINJACOREAPI void BNFreeRegisterValue(BNRegisterValue* value);
 	BINARYNINJACOREAPI uint32_t* BNGetRegistersReadByInstruction(BNFunction* func, BNArchitecture* arch, uint64_t addr,
 	                                                             size_t* count);
 	BINARYNINJACOREAPI uint32_t* BNGetRegistersWrittenByInstruction(BNFunction* func, BNArchitecture* arch, uint64_t addr,
diff --git a/function.cpp b/function.cpp
index b8429998..7cc30351 100644
--- a/function.cpp
+++ b/function.cpp
@@ -143,27 +143,56 @@ vector<size_t> Function::GetLowLevelILExitsForInstruction(Architecture* arch, ui
 }
 
 
-BNRegisterValue Function::GetRegisterValueAtInstruction(Architecture* arch, uint64_t addr, uint32_t reg)
+static RegisterValue GetRegisterValueFromAPIObject(BNRegisterValue& value)
 {
-	return BNGetRegisterValueAtInstruction(m_object, arch->GetObject(), addr, reg);
+	RegisterValue result;
+	result.state = value.state;
+	result.reg = value.reg;
+	result.value = value.value;
+	result.rangeStart = value.rangeStart;
+	result.rangeEnd = value.rangeEnd;
+	result.rangeStep = value.rangeStep;
+	if (value.state == LookupTableValue)
+	{
+		for (size_t i = 0; i < (size_t)value.rangeEnd; i++)
+		{
+			LookupTableEntry entry;
+			entry.fromValues.insert(entry.fromValues.end(), &value.table[i].fromValues[0],
+				&value.table[i].fromValues[value.table[i].fromCount]);
+			entry.toValue = value.table[i].toValue;
+			result.table.push_back(entry);
+		}
+	}
+	BNFreeRegisterValue(&value);
+	return result;
+}
+
+
+RegisterValue Function::GetRegisterValueAtInstruction(Architecture* arch, uint64_t addr, uint32_t reg)
+{
+	BNRegisterValue value = BNGetRegisterValueAtInstruction(m_object, arch->GetObject(), addr, reg);
+	return GetRegisterValueFromAPIObject(value);
 }
 
 
-BNRegisterValue Function::GetRegisterValueAfterInstruction(Architecture* arch, uint64_t addr, uint32_t reg)
+RegisterValue Function::GetRegisterValueAfterInstruction(Architecture* arch, uint64_t addr, uint32_t reg)
 {
-	return BNGetRegisterValueAfterInstruction(m_object, arch->GetObject(), addr, reg);
+	BNRegisterValue value = BNGetRegisterValueAfterInstruction(m_object, arch->GetObject(), addr, reg);
+	return GetRegisterValueFromAPIObject(value);
 }
 
 
-BNRegisterValue Function::GetRegisterValueAtLowLevelILInstruction(size_t i, uint32_t reg)
+RegisterValue Function::GetRegisterValueAtLowLevelILInstruction(size_t i, uint32_t reg)
 {
-	return BNGetRegisterValueAtLowLevelILInstruction(m_object, i, reg);
+	BNRegisterValue value = BNGetRegisterValueAtLowLevelILInstruction(m_object, i, reg);
+	return GetRegisterValueFromAPIObject(value);
 }
 
 
-BNRegisterValue Function::GetRegisterValueAfterLowLevelILInstruction(size_t i, uint32_t reg)
+RegisterValue Function::GetRegisterValueAfterLowLevelILInstruction(size_t i, uint32_t reg)
 {
-	return BNGetRegisterValueAfterLowLevelILInstruction(m_object, i, reg);
+	BNRegisterValue value = BNGetRegisterValueAfterLowLevelILInstruction(m_object, i, reg);
+	return GetRegisterValueFromAPIObject(value);
 }
 
 
diff --git a/python/__init__.py b/python/__init__.py
index cdbce895..dfd76b90 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -2003,28 +2003,69 @@ class Enumeration(object):
 		else:
 			core.BNAddEnumerationMemberWithValue(self.handle, name, value)
 
+class LookupTableEntry:
+	def __init__(self, from_values, to_value):
+		self.from_values = from_values
+		self.to_value = to_value
+
+	def __repr__(self):
+		return "[%s] -> 0x%x" % (', '.join(["0x%x" % i for i in self.from_values]), self.to_value)
+
 class RegisterValue:
-	def __init__(self, arch, value_type, reg, value):
-		self.type = core.BNRegisterValueType_names[value_type]
-		if value_type == EntryValue:
-			self.reg = arch.get_reg_name(reg)
-		elif value_type == OffsetFromEntryValue:
-			self.reg = arch.get_reg_name(reg)
-			self.offset = value
-		elif value_type == ConstantValue:
-			self.value = value
-		elif value_type == StackFrameOffset:
-			self.offset = value
+	def __init__(self, arch, value):
+		self.type = value.state
+		if value.state == EntryValue:
+			self.reg = arch.get_reg_name(value.reg)
+		elif value.state == OffsetFromEntryValue:
+			self.reg = arch.get_reg_name(value.reg)
+			self.offset = value.value
+		elif value.state == ConstantValue:
+			self.value = value.value
+		elif value.state == StackFrameOffset:
+			self.offset = value.value
+		elif value.state == SignedRangeValue:
+			self.offset = value.value
+			self.start = value.rangeStart
+			self.end = value.rangeEnd
+			self.step = value.rangeStep
+			if self.start & (1 << 63):
+				self.start |= ~((1 << 63) - 1)
+			if self.end & (1 << 63):
+				self.end |= ~((1 << 63) - 1)
+		elif value.state == UnsignedRangeValue:
+			self.offset = value.value
+			self.start = value.rangeStart
+			self.end = value.rangeEnd
+			self.step = value.rangeStep
+		elif value.state == LookupTableValue:
+			self.table = []
+			self.mapping = {}
+			for i in xrange(0, value.rangeEnd):
+				from_list = []
+				for j in xrange(0, value.table[i].fromCount):
+					from_list.append(value.table[i].fromValues[j])
+					self.mapping[value.table[i].fromValues[j]] = value.table[i].toValue
+				self.table.append(LookupTableEntry(from_list, value.table[i].toValue))
+		elif value.state == OffsetFromUndeterminedValue:
+			self.offset = value.value
 
 	def __repr__(self):
-		if self.type == "EntryValue":
+		if self.type == EntryValue:
 			return "<entry %s>" % self.reg
-		if self.type == "OffsetFromEntryValue":
+		if self.type == OffsetFromEntryValue:
 			return "<entry %s + 0x%x>" % (self.reg, self.offset)
-		if self.type == "ConstantValue":
+		if self.type == ConstantValue:
 			return "<const 0x%x>" % self.value
-		if self.type == "StackFrameOffset":
+		if self.type == StackFrameOffset:
 			return "<stack frame offset 0x%x>" % self.offset
+		if (self.type == SignedRangeValue) or (self.type == UnsignedRangeValue):
+			if self.step == 1:
+				return "<range: 0x%x-0x%x>" % (self.start, self.end)
+			return "<range: 0x%x-0x%x, step 0x%x>" % (self.start, self.end, self.step)
+		if self.type == LookupTableValue:
+			return "<table: %s>" % ', '.join([repr(i) for i in self.table])
+		if self.type == OffsetFromUndeterminedValue:
+			return "<undetermined with offset 0x%x>" % self.offset
 		return "<undetermined>"
 
 class StackVariable:
@@ -2236,25 +2277,33 @@ class Function(object):
 		if isinstance(reg, str):
 			reg = arch.regs[reg].index
 		value = core.BNGetRegisterValueAtInstruction(self.handle, arch.handle, addr, reg)
-		return RegisterValue(arch, value.state, value.reg, value.value)
+		result = RegisterValue(arch, value)
+		core.BNFreeRegisterValue(value)
+		return result
 
 	def get_reg_value_after(self, arch, addr, reg):
 		if isinstance(reg, str):
 			reg = arch.regs[reg].index
 		value = core.BNGetRegisterValueAfterInstruction(self.handle, arch.handle, addr, reg)
-		return RegisterValue(arch, value.state, value.reg, value.value)
+		result = RegisterValue(arch, value)
+		core.BNFreeRegisterValue(value)
+		return result
 
 	def get_reg_value_at_low_level_il_instruction(self, i, reg):
 		if isinstance(reg, str):
 			reg = self.arch.regs[reg].index
 		value = core.BNGetRegisterValueAtInstruction(self.handle, self.arch.handle, i, reg)
-		return RegisterValue(self.arch, value.state, value.reg, value.value)
+		result = RegisterValue(self.arch, value)
+		core.BNFreeRegisterValue(value)
+		return result
 
 	def get_reg_value_after_low_level_il_instruction(self, i, reg):
 		if isinstance(reg, str):
 			reg = self.arch.regs[reg].index
 		value = core.BNGetRegisterValueAfterInstruction(self.handle, self.arch.handle, i, reg)
-		return RegisterValue(self.arch, value.state, value.reg, value.value)
+		result = RegisterValue(self.arch, value)
+		core.BNFreeRegisterValue(value)
+		return result
 
 	def get_regs_read_by(self, arch, addr):
 		count = ctypes.c_ulonglong()