Add APIs to access and update portions of the function type, and added new APIs for global registers and implicit incoming state in calling conventions

4 files changed, 363 insertions, 14 deletions

diff --git a/python/architecture.py b/python/architecture.py
index b5f56767..61559934 100644
--- a/python/architecture.py
+++ b/python/architecture.py
@@ -116,6 +116,7 @@ class Architecture(object):
 	regs = {}
 	stack_pointer = None
 	link_reg = None
+	global_regs = []
 	flags = []
 	flag_write_types = []
 	flag_roles = {}
@@ -208,6 +209,13 @@ class Architecture(object):
 				core.BNFreeRegisterList(flags)
 				self._flags_written_by_flag_write_type[self._flag_write_types[write_type]] = flag_indexes
 				self.__dict__["flags_written_by_flag_write_type"][write_type] = flag_names
+
+			count = ctypes.c_ulonglong()
+			regs = core.BNGetArchitectureGlobalRegisters(self.handle, count)
+			self.__dict__["global_regs"] = []
+			for i in xrange(0, count.value):
+				self.global_regs.append(core.BNGetArchitectureRegisterName(self.handle, regs[i]))
+			core.BNFreeRegisterList(regs)
 		else:
 			startup._init_plugins()
 
@@ -250,6 +258,7 @@ class Architecture(object):
 			self._cb.getStackPointerRegister = self._cb.getStackPointerRegister.__class__(
 				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.assemble = self._cb.assemble.__class__(self._assemble)
 			self._cb.isNeverBranchPatchAvailable = self._cb.isNeverBranchPatchAvailable.__class__(
 				self._is_never_branch_patch_available)
@@ -330,6 +339,8 @@ class Architecture(object):
 					flags.append(self._flags[flag])
 				self._flags_written_by_flag_write_type[self._flag_write_types[write_type]] = flags
 
+			self.__dict__["global_regs"] = self.__class__.global_regs
+
 			self._pending_reg_lists = {}
 			self._pending_token_lists = {}
 
@@ -719,6 +730,20 @@ class Architecture(object):
 			log.log_error(traceback.format_exc())
 			return 0
 
+	def _get_global_registers(self, ctxt, count):
+		try:
+			count[0] = len(self.__class__.global_regs)
+			reg_buf = (ctypes.c_uint * len(self.__class__.global_regs))()
+			for i in xrange(0, len(self.__class__.global_regs)):
+				reg_buf[i] = self._all_regs[self.__class__.global_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 _assemble(self, ctxt, code, addr, result, errors):
 		try:
 			data, error_str = self.perform_assemble(code, addr)
diff --git a/python/callingconvention.py b/python/callingconvention.py
index 21c8c95a..609c71b0 100644
--- a/python/callingconvention.py
+++ b/python/callingconvention.py
@@ -26,6 +26,8 @@ import _binaryninjacore as core
 import architecture
 import log
 import types
+import function
+import binaryview
 
 
 class CallingConvention(object):
@@ -38,6 +40,8 @@ class CallingConvention(object):
 	int_return_reg = None
 	high_int_return_reg = None
 	float_return_reg = None
+	global_pointer_reg = None
+	implicitly_defined_regs = []
 
 	_registered_calling_conventions = []
 
@@ -58,6 +62,10 @@ class CallingConvention(object):
 			self._cb.getIntegerReturnValueRegister = self._cb.getIntegerReturnValueRegister.__class__(self._get_int_return_reg)
 			self._cb.getHighIntegerReturnValueRegister = self._cb.getHighIntegerReturnValueRegister.__class__(self._get_high_int_return_reg)
 			self._cb.getFloatReturnValueRegister = self._cb.getFloatReturnValueRegister.__class__(self._get_float_return_reg)
+			self._cb.getGlobalPointerRegister = self._cb.getGlobalPointerRegister.__class__(self._get_global_pointer_reg)
+			self._cb.getImplicitlyDefinedRegisters = self._cb.getImplicitlyDefinedRegisters.__class__(self._get_implicitly_defined_regs)
+			self._cb.getIncomingRegisterValue = self._cb.getIncomingRegisterValue.__class__(self._get_incoming_reg_value)
+			self._cb.getIncomingFlagValue = self._cb.getIncomingFlagValue.__class__(self._get_incoming_flag_value)
 			self.handle = core.BNCreateCallingConvention(arch.handle, name, self._cb)
 			self.__class__._registered_calling_conventions.append(self)
 		else:
@@ -112,6 +120,21 @@ class CallingConvention(object):
 			else:
 				self.__dict__["float_return_reg"] = self.arch.get_reg_name(reg)
 
+			reg = core.BNGetGlobalPointerRegister(self.handle)
+			if reg == 0xffffffff:
+				self.__dict__["global_pointer_reg"] = None
+			else:
+				self.__dict__["global_pointer_reg"] = self.arch.get_reg_name(reg)
+
+			count = ctypes.c_ulonglong()
+			regs = core.BNGetImplicitlyDefinedRegisters(self.handle, count)
+			result = []
+			arch = self.arch
+			for i in xrange(0, count.value):
+				result.append(arch.get_reg_name(regs[i]))
+			core.BNFreeRegisterList(regs, count.value)
+			self.__dict__["implicitly_defined_regs"] = result
+
 		self.confidence = confidence
 
 	def __del__(self):
@@ -220,12 +243,76 @@ class CallingConvention(object):
 			log.log_error(traceback.format_exc())
 			return False
 
+	def _get_global_pointer_reg(self, ctxt):
+		try:
+			if self.__class__.global_pointer_reg is None:
+				return 0xffffffff
+			return self.arch.regs[self.__class__.global_pointer_reg].index
+		except:
+			log.log_error(traceback.format_exc())
+			return False
+
+	def _get_implicitly_defined_regs(self, ctxt, count):
+		try:
+			regs = self.__class__.implicitly_defined_regs
+			count[0] = len(regs)
+			reg_buf = (ctypes.c_uint * len(regs))()
+			for i in xrange(0, len(regs)):
+				reg_buf[i] = self.arch.regs[regs[i]].index
+			result = ctypes.cast(reg_buf, ctypes.c_void_p)
+			self._pending_reg_lists[result.value] = (result, reg_buf)
+			return result.value
+		except:
+			log.log_error(traceback.format_exc())
+			count[0] = 0
+			return None
+
+	def _get_incoming_reg_value(self, ctxt, reg, func):
+		try:
+			func_obj = function.Function(binaryview.BinaryView(handle = core.BNGetFunctionData(func)),
+				core.BNNewFunctionReference(func))
+			reg_name = self.arch.get_reg_name(reg)
+			return self.perform_get_incoming_reg_value(reg_name, func_obj)._to_api_object()
+		except:
+			log.log_error(traceback.format_exc())
+			return function.RegisterValue()._to_api_object()
+
+	def _get_incoming_flag_value(self, ctxt, reg, func):
+		try:
+			func_obj = function.Function(binaryview.BinaryView(handle = core.BNGetFunctionData(func)),
+				core.BNNewFunctionReference(func))
+			reg_name = self.arch.get_reg_name(reg)
+			return self.perform_get_incoming_flag_value(reg_name, func_obj)._to_api_object()
+		except:
+			log.log_error(traceback.format_exc())
+			return function.RegisterValue()._to_api_object()
+
 	def __repr__(self):
 		return "<calling convention: %s %s>" % (self.arch.name, self.name)
 
 	def __str__(self):
 		return self.name
 
+	def perform_get_incoming_reg_value(self, reg, func):
+		return function.RegisterValue()
+
+	def perform_get_incoming_flag_value(self, reg, func):
+		return function.RegisterValue()
+
 	def with_confidence(self, confidence):
 		return CallingConvention(self.arch, handle = core.BNNewCallingConventionReference(self.handle),
 			confidence = confidence)
+
+	def get_incoming_reg_value(self, reg, func):
+		reg_num = self.arch.get_reg_index(reg)
+		func_handle = None
+		if func is not None:
+			func_handle = func.handle
+		return function.RegisterValue(self.arch, core.BNGetIncomingRegisterValue(self.handle, reg_num, func_handle))
+
+	def get_incoming_flag_value(self, flag, func):
+		reg_num = self.arch.get_flag_index(flag)
+		func_handle = None
+		if func is not None:
+			func_handle = func.handle
+		return function.RegisterValue(self.arch, core.BNGetIncomingFlagValue(self.handle, reg_num, func_handle))
diff --git a/python/function.py b/python/function.py
index b14ba53e..cf9bd759 100644
--- a/python/function.py
+++ b/python/function.py
@@ -37,6 +37,7 @@ import lowlevelil
 import mediumlevelil
 import binaryview
 import log
+import callingconvention
 
 
 class LookupTableEntry(object):
@@ -49,26 +50,50 @@ class LookupTableEntry(object):
 
 
 class RegisterValue(object):
-	def __init__(self, arch, value):
-		self.type = RegisterValueType(value.state)
-		if value.state == RegisterValueType.EntryValue:
-			self.reg = arch.get_reg_name(value.value)
-		elif value.state == RegisterValueType.ConstantValue:
-			self.value = value.value
-		elif value.state == RegisterValueType.StackFrameOffset:
-			self.offset = value.value
+	def __init__(self, arch = None, value = None):
+		if value is None:
+			self.type = RegisterValueType.UndeterminedValue
+		else:
+			self.type = RegisterValueType(value.state)
+			if value.state == RegisterValueType.EntryValue:
+				self.arch = arch
+				if arch is not None:
+					self.reg = arch.get_reg_name(value.value)
+				else:
+					self.reg = value.value
+			elif (value.state == RegisterValueType.ConstantValue) or (value.state == RegisterValueType.ConstantPointerValue):
+				self.value = value.value
+			elif value.state == RegisterValueType.StackFrameOffset:
+				self.offset = value.value
 
 	def __repr__(self):
 		if self.type == RegisterValueType.EntryValue:
 			return "<entry %s>" % self.reg
 		if self.type == RegisterValueType.ConstantValue:
 			return "<const %#x>" % self.value
+		if self.type == RegisterValueType.ConstantPointerValue:
+			return "<const ptr %#x>" % self.value
 		if self.type == RegisterValueType.StackFrameOffset:
 			return "<stack frame offset %#x>" % self.offset
 		if self.type == RegisterValueType.ReturnAddressValue:
 			return "<return address>"
 		return "<undetermined>"
 
+	def _to_api_object(self):
+		result = core.BNRegisterValue()
+		result.state = self.type
+		result.value = 0
+		if self.type == RegisterValueType.EntryValue:
+			if self.arch is not None:
+				result.value = self.arch.get_reg_index(self.reg)
+			else:
+				result.value = self.reg
+		elif (self.type == RegisterValueType.ConstantValue) or (self.type == RegisterValueType.ConstantPointerValue):
+			result.value = self.value
+		elif self.type == RegisterValueType.StackFrameOffset:
+			result.value = self.offset
+		return result
+
 
 class ValueRange(object):
 	def __init__(self, start, end, step):
@@ -240,6 +265,25 @@ class IndirectBranchInfo(object):
 		return "<branch %s:%#x -> %s:%#x>" % (self.source_arch.name, self.source_addr, self.dest_arch.name, self.dest_addr)
 
 
+class ParameterVariables(object):
+	def __init__(self, var_list, confidence = types.max_confidence):
+		self.vars = var_list
+		self.confidence = confidence
+
+	def __repr__(self):
+		return repr(self.vars)
+
+	def __iter__(self):
+		for var in self.vars:
+			yield var
+
+	def __getitem__(self, idx):
+		return self.vars[idx]
+
+	def with_confidence(self, confidence):
+		return ParameterVariables(list(self.vars), confidence = confidence)
+
+
 class _FunctionAssociatedDataStore(associateddatastore._AssociatedDataStore):
 	_defaults = {}
 
@@ -335,8 +379,19 @@ class Function(object):
 
 	@property
 	def can_return(self):
-		"""Whether function can return (read-only)"""
-		return core.BNCanFunctionReturn(self.handle)
+		"""Whether function can return"""
+		result = core.BNCanFunctionReturn(self.handle)
+		return types.BoolWithConfidence(result.value, confidence = result.confidence)
+
+	@can_return.setter
+	def can_return(self, value):
+		bc = core.BNBoolWithConfidence()
+		bc.value = bool(value)
+		if hasattr(value, 'confidence'):
+			bc.confidence = value.confidence
+		else:
+			bc.confidence = types.max_confidence
+		core.BNSetUserFunctionCanReturn(self.handle, bc)
 
 	@property
 	def explicitly_defined_type(self):
@@ -452,6 +507,97 @@ class Function(object):
 		core.BNFreeAnalysisPerformanceInfo(info, count.value)
 		return result
 
+	@property
+	def type_tokens(self):
+		"""Text tokens for this function's prototype"""
+		return self.get_type_tokens()[0].tokens
+
+	@property
+	def return_type(self):
+		"""Return type of the function"""
+		result = core.BNGetFunctionReturnType(self.handle)
+		if not result.type:
+			return None
+		return types.Type(result.type, confidence = result.confidence)
+
+	@return_type.setter
+	def return_type(self, value):
+		type_conf = core.BNTypeWithConfidence()
+		if value is None:
+			type_conf.type = None
+			type_conf.confidence = 0
+		else:
+			type_conf.type = value.handle
+			type_conf.confidence = value.confidence
+		core.BNSetUserFunctionReturnType(self.handle, type_conf)
+
+	@property
+	def calling_convention(self):
+		"""Calling convention used by the function"""
+		result = core.BNGetFunctionCallingConvention(self.handle)
+		if not result.convention:
+			return None
+		return callingconvention.CallingConvention(None, handle = result.convention, confidence = result.confidence)
+
+	@calling_convention.setter
+	def calling_convention(self, value):
+		conv_conf = core.BNCallingConventionWithConfidence()
+		if value is None:
+			conv_conf.convention = None
+			conv_conf.confidence = 0
+		else:
+			conv_conf.convention = value.handle
+			conv_conf.confidence = value.confidence
+		core.BNSetUserFunctionCallingConvention(self.handle, conv_conf)
+
+	@property
+	def parameter_vars(self):
+		"""List of variables for the incoming function parameters"""
+		result = core.BNGetFunctionParameterVariables(self.handle)
+		var_list = []
+		for i in xrange(0, result.count):
+			var_list.append(Variable(self, result.vars[i].type, result.vars[i].index, result.vars[i].storage))
+		confidence = result.confidence
+		core.BNFreeParameterVariables(result)
+		return ParameterVariables(var_list, confidence = confidence)
+
+	@parameter_vars.setter
+	def parameter_vars(self, value):
+		if value is None:
+			var_list = []
+		else:
+			var_list = list(value)
+		var_conf = core.BNParameterVariablesWithConfidence()
+		var_conf.vars = (core.BNVariable * len(var_list))()
+		var_conf.count = len(var_list)
+		for i in xrange(0, len(var_list)):
+			var_conf.vars[i].type = var_list[i].source_type
+			var_conf.vars[i].index = var_list[i].index
+			var_conf.vars[i].storage = var_list[i].storage
+		if value is None:
+			var_conf.confidence = 0
+		elif hasattr(value, 'confidence'):
+			var_conf.confidence = value.confidence
+		else:
+			var_conf.confidence = types.max_confidence
+		core.BNSetUserFunctionParameterVariables(self.handle, var_conf)
+
+	@property
+	def has_variable_arguments(self):
+		"""Whether the function takes a variable number of arguments"""
+		result = core.BNFunctionHasVariableArguments(self.handle)
+		return types.BoolWithConfidence(result.value, confidence = result.confidence)
+
+	@has_variable_arguments.setter
+	def has_variable_arguments(self, value):
+		bc = core.BNBoolWithConfidence()
+		bc.value = bool(value)
+		if hasattr(value, 'confidence'):
+			bc.confidence = value.confidence
+		else:
+			bc.confidence = types.max_confidence
+		core.BNSetUserFunctionHasVariableArguments(self.handle, bc)
+
 	def __iter__(self):
 		count = ctypes.c_ulonglong()
 		blocks = core.BNGetFunctionBasicBlockList(self.handle, count)
@@ -759,6 +905,64 @@ class Function(object):
 	def set_user_type(self, value):
 		core.BNSetFunctionUserType(self.handle, value.handle)
 
+	def set_auto_return_type(self, value):
+		type_conf = core.BNTypeWithConfidence()
+		if value is None:
+			type_conf.type = None
+			type_conf.confidence = 0
+		else:
+			type_conf.type = value.handle
+			type_conf.confidence = value.confidence
+		core.BNSetAutoFunctionReturnType(self.handle, type_conf)
+
+	def set_auto_calling_convention(self, value):
+		conv_conf = core.BNCallingConventionWithConfidence()
+		if value is None:
+			conv_conf.convention = None
+			conv_conf.confidence = 0
+		else:
+			conv_conf.convention = value.handle
+			conv_conf.confidence = value.confidence
+		core.BNSetAutoFunctionCallingConvention(self.handle, conv_conf)
+
+	def set_auto_parameter_vars(self, value):
+		if value is None:
+			var_list = []
+		else:
+			var_list = list(value)
+		var_conf = core.BNParameterVariablesWithConfidence()
+		var_conf.vars = (core.BNVariable * len(var_list))()
+		var_conf.count = len(var_list)
+		for i in xrange(0, len(var_list)):
+			var_conf.vars[i].type = var_list[i].source_type
+			var_conf.vars[i].index = var_list[i].index
+			var_conf.vars[i].storage = var_list[i].storage
+		if value is None:
+			var_conf.confidence = 0
+		elif hasattr(value, 'confidence'):
+			var_conf.confidence = value.confidence
+		else:
+			var_conf.confidence = types.max_confidence
+		core.BNSetAutoFunctionParameterVariables(self.handle, var_conf)
+
+	def set_auto_has_variable_arguments(self, value):
+		bc = core.BNBoolWithConfidence()
+		bc.value = bool(value)
+		if hasattr(value, 'confidence'):
+			bc.confidence = value.confidence
+		else:
+			bc.confidence = types.max_confidence
+		core.BNSetAutoFunctionHasVariableArguments(self.handle, bc)
+
+	def set_auto_can_return(self, value):
+		bc = core.BNBoolWithConfidence()
+		bc.value = bool(value)
+		if hasattr(value, 'confidence'):
+			bc.confidence = value.confidence
+		else:
+			bc.confidence = types.max_confidence
+		core.BNSetAutoFunctionCanReturn(self.handle, bc)
+
 	def get_int_display_type(self, instr_addr, value, operand, arch=None):
 		if arch is None:
 			arch = self.arch
@@ -932,6 +1136,29 @@ class Function(object):
 		core.BNFreeVariableNameAndType(found_var)
 		return result
 
+	def get_type_tokens(self, settings=None):
+		if settings is not None:
+			settings = settings.handle
+		count = ctypes.c_ulonglong()
+		lines = core.BNGetFunctionTypeTokens(self.handle, settings, count)
+		result = []
+		for i in xrange(0, count.value):
+			addr = lines[i].addr
+			tokens = []
+			for j in xrange(0, lines[i].count):
+				token_type = InstructionTextTokenType(lines[i].tokens[j].type)
+				text = lines[i].tokens[j].text
+				value = lines[i].tokens[j].value
+				size = lines[i].tokens[j].size
+				operand = lines[i].tokens[j].operand
+				context = lines[i].tokens[j].context
+				confidence = lines[i].tokens[j].confidence
+				address = lines[i].tokens[j].address
+				tokens.append(InstructionTextToken(token_type, text, value, size, operand, context, address, confidence))
+			result.append(DisassemblyTextLine(addr, tokens))
+		core.BNFreeDisassemblyTextLines(lines, count.value)
+		return result
+
 
 class AdvancedFunctionAnalysisDataRequestor(object):
 	def __init__(self, func = None):
diff --git a/python/types.py b/python/types.py
index 47d99bee..e297ddd2 100644
--- a/python/types.py
+++ b/python/types.py
@@ -279,7 +279,7 @@ class Type(object):
 		result = core.BNGetTypeCallingConvention(self.handle)
 		if not result.convention:
 			return None
-		return callingconvention.CallingConvention(None, handle = result, confidence = result.confidence)
+		return callingconvention.CallingConvention(None, handle = result.convention, confidence = result.confidence)
 
 	@property
 	def parameters(self):
@@ -295,7 +295,8 @@ class Type(object):
 	@property
 	def has_variable_arguments(self):
 		"""Whether type has variable arguments (read-only)"""
-		return core.BNTypeHasVariableArguments(self.handle)
+		result = core.BNTypeHasVariableArguments(self.handle)
+		return BoolWithConfidence(result.value, confidence = result.confidence)
 
 	@property
 	def can_return(self):
@@ -505,7 +506,7 @@ class Type(object):
 		return Type(core.BNCreateArrayType(type_conf, count))
 
 	@classmethod
-	def function(self, ret, params, calling_convention=None, variable_arguments=False):
+	def function(self, ret, params, calling_convention=None, variable_arguments=None):
 		"""
 		``function`` class method for creating an function Type.
 
@@ -537,8 +538,17 @@ class Type(object):
 			conv_conf.convention = calling_convention.handle
 			conv_conf.confidence = calling_convention.confidence
 
+		if variable_arguments is None:
+			variable_arguments = BoolWithConfidence(False, confidence = 0)
+		elif not isinstance(variable_arguments, BoolWithConfidence):
+			variable_arguments = BoolWithConfidence(variable_arguments)
+
+		vararg_conf = core.BNBoolWithConfidence()
+		vararg_conf.value = variable_arguments.value
+		vararg_conf.confidence = variable_arguments.confidence
+
 		return Type(core.BNCreateFunctionType(ret_conf, conv_conf, param_buf, len(params),
-			variable_arguments))
+			vararg_conf))
 
 	@classmethod
 	def generate_auto_type_id(self, source, name):