Adding CPU intrinsics support

4 files changed, 272 insertions, 18 deletions

diff --git a/python/architecture.py b/python/architecture.py
index 2b422962..3ccd2b11 100644
--- a/python/architecture.py
+++ b/python/architecture.py
@@ -33,6 +33,7 @@ import callingconvention
 import platform
 import log
 import databuffer
+import types
 
 
 class _ArchitectureMetaClass(type):
@@ -128,6 +129,7 @@ class Architecture(object):
 	flags_written_by_flag_write_type = {}
 	semantic_class_for_flag_write_type = {}
 	reg_stacks = {}
+	intrinsics = {}
 	__metaclass__ = _ArchitectureMetaClass
 	next_address = 0
 
@@ -304,6 +306,27 @@ class Architecture(object):
 				top = core.BNGetArchitectureRegisterName(self.handle, info.stackTopReg)
 				self.reg_stacks[name] = function.RegisterStackInfo(storage, top_rel, top, regs[i])
 			core.BNFreeRegisterList(regs)
+
+			count = ctypes.c_ulonglong()
+			intrinsics = core.BNGetAllArchitectureIntrinsics(self.handle, count)
+			self.__dict__["intrinsics"] = {}
+			for i in xrange(0, count.value):
+				name = core.BNGetArchitectureIntrinsicName(self.handle, intrinsics[i])
+				input_count = ctypes.c_ulonglong()
+				inputs = core.BNGetArchitectureIntrinsicInputs(self.handle, intrinsics[i], input_count)
+				input_list = []
+				for j in xrange(0, input_count.value):
+					input_name = inputs[j].name
+					type_obj = types.Type(core.BNNewTypeReference(inputs[j].type), confidence = inputs[j].typeConfidence)
+					input_list.append(function.IntrinsicInput(type_obj, input_name))
+				core.BNFreeNameAndTypeList(inputs, input_count.value)
+				output_count = ctypes.c_ulonglong()
+				outputs = core.BNGetArchitectureIntrinsicOutputs(self.handle, intrinsics[i], output_count)
+				output_list = []
+				for j in xrange(0, output_count.value):
+					output_list.append(types.Type(core.BNNewTypeReference(outputs[j].type), confidence = outputs[j].confidence))
+				core.BNFreeOutputTypeList(outputs, output_count.value)
+				self.intrinsics[name] = function.IntrinsicInfo(input_list, output_list)
 		else:
 			startup._init_plugins()
 
@@ -365,6 +388,12 @@ class Architecture(object):
 			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.getIntrinsicName = self._cb.getIntrinsicName.__class__(self._get_intrinsic_name)
+			self._cb.getAllIntrinsics = self._cb.getAllIntrinsics.__class__(self._get_all_intrinsics)
+			self._cb.getIntrinsicInputs = self._cb.getIntrinsicInputs.__class__(self._get_intrinsic_inputs)
+			self._cb.freeNameAndTypeList = self._cb.freeNameAndTypeList.__class__(self._free_name_and_type_list)
+			self._cb.getIntrinsicOutputs = self._cb.getIntrinsicOutputs.__class__(self._get_intrinsic_outputs)
+			self._cb.freeTypeList = self._cb.freeTypeList.__class__(self._free_type_list)
 			self._cb.assemble = self._cb.assemble.__class__(self._assemble)
 			self._cb.isNeverBranchPatchAvailable = self._cb.isNeverBranchPatchAvailable.__class__(
 				self._is_never_branch_patch_available)
@@ -514,9 +543,28 @@ class Architecture(object):
 
 			self.__dict__["global_regs"] = self.__class__.global_regs
 
+			self._intrinsics = {}
+			self._intrinsics_by_index = {}
+			self.__dict__["intrinsics"] = self.__class__.intrinsics
+			intrinsic_index = 0
+			for intrinsic in self.__class__.intrinsics.keys():
+				if intrinsic not in self._intrinsics:
+					info = self.__class__.intrinsics[intrinsic]
+					for i in xrange(0, len(info.inputs)):
+						if isinstance(info.inputs[i], types.Type):
+							info.inputs[i] = function.IntrinsicInput(info.inputs[i])
+						elif isinstance(info.inputs[i], tuple):
+							info.inputs[i] = function.IntrinsicInput(info.inputs[i][0], info.inputs[i][1])
+					info.index = intrinsic_index
+					self._intrinsics[intrinsic] = intrinsic_index
+					self._intrinsics_by_index[intrinsic_index] = (intrinsic, info)
+					intrinsic_index += 1
+
 			self._pending_reg_lists = {}
 			self._pending_token_lists = {}
 			self._pending_condition_lists = {}
+			self._pending_name_and_type_lists = {}
+			self._pending_type_lists = {}
 
 	def __eq__(self, value):
 		if not isinstance(value, Architecture):
@@ -1116,6 +1164,95 @@ class Architecture(object):
 			result[0].topRelativeCount = 0
 			result[0].stackTopReg = 0
 
+	def _get_intrinsic_name(self, ctxt, intrinsic):
+		try:
+			if intrinsic in self._intrinsics_by_index:
+				return core.BNAllocString(self._intrinsics_by_index[intrinsic][0])
+			return core.BNAllocString("")
+		except (KeyError, OSError):
+			log.log_error(traceback.format_exc())
+			return core.BNAllocString("")
+
+	def _get_all_intrinsics(self, ctxt, count):
+		try:
+			regs = self._intrinsics_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_intrinsic_inputs(self, ctxt, intrinsic, count):
+		try:
+			if intrinsic in self._intrinsics_by_index:
+				inputs = self._intrinsics_by_index[intrinsic][1].inputs
+				count[0] = len(inputs)
+				input_buf = (core.BNNameAndType * len(inputs))()
+				for i in xrange(0, len(inputs)):
+					input_buf[i].name = inputs[i].name
+					input_buf[i].type = core.BNNewTypeReference(inputs[i].type.handle)
+					input_buf[i].typeConfidence = inputs[i].type.confidence
+				result = ctypes.cast(input_buf, ctypes.c_void_p)
+				self._pending_name_and_type_lists[result.value] = (result, input_buf, len(inputs))
+				return result.value
+			count[0] = 0
+			return None
+		except:
+			log.log_error(traceback.format_exc())
+			count[0] = 0
+			return None
+
+	def _free_name_and_type_list(self, ctxt, buf_raw):
+		try:
+			buf = ctypes.cast(buf_raw, ctypes.c_void_p)
+			if buf.value not in self._pending_name_and_type_lists:
+				raise ValueError("freeing name and type list that wasn't allocated")
+			name_and_types = self._pending_name_and_type_lists[buf.value][1]
+			count = self._pending_name_and_type_lists[buf.value][2]
+			for i in xrange(0, count):
+				core.BNFreeType(name_and_types[i].type)
+			del self._pending_name_and_type_lists[buf.value]
+		except (ValueError, KeyError):
+			log.log_error(traceback.format_exc())
+
+	def _get_intrinsic_outputs(self, ctxt, intrinsic, count):
+		try:
+			if intrinsic in self._intrinsics_by_index:
+				outputs = self._intrinsics_by_index[intrinsic][1].outputs
+				count[0] = len(outputs)
+				output_buf = (core.BNTypeWithConfidence * len(outputs))()
+				for i in xrange(0, len(outputs)):
+					output_buf[i].type = core.BNNewTypeReference(outputs[i].handle)
+					output_buf[i].confidence = outputs[i].confidence
+				result = ctypes.cast(output_buf, ctypes.c_void_p)
+				self._pending_type_lists[result.value] = (result, output_buf, len(outputs))
+				return result.value
+			count[0] = 0
+			return None
+		except:
+			log.log_error(traceback.format_exc())
+			count[0] = 0
+			return None
+
+	def _free_type_list(self, ctxt, buf_raw):
+		try:
+			buf = ctypes.cast(buf_raw, ctypes.c_void_p)
+			if buf.value not in self._pending_type_lists:
+				raise ValueError("freeing type list that wasn't allocated")
+			types = self._pending_type_lists[buf.value][1]
+			count = self._pending_type_lists[buf.value][2]
+			for i in xrange(0, count):
+				core.BNFreeType(types[i].type)
+			del self._pending_type_lists[buf.value]
+		except (ValueError, KeyError):
+			log.log_error(traceback.format_exc())
+
 	def _assemble(self, ctxt, code, addr, result, errors):
 		try:
 			data, error_str = self.perform_assemble(code, addr)
@@ -1701,6 +1838,23 @@ class Architecture(object):
 			return sem_group.index
 		return sem_group
 
+	def get_intrinsic_name(self, intrinsic):
+		"""
+		``get_intrinsic_name`` gets an intrinsic name from an intrinsic number.
+
+		:param int intrinsic: intrinsic number
+		:return: the corresponding intrinsic string
+		:rtype: str
+		"""
+		return core.BNGetArchitectureIntrinsicName(self.handle, intrinsic)
+
+	def get_intrinsic_index(self, intrinsic):
+		if isinstance(intrinsic, str):
+			return self._intrinsics[intrinsic]
+		elif isinstance(intrinsic, lowlevelil.ILIntrinsic):
+			return intrinsic.index
+		return intrinsic
+
 	def get_flag_write_type_name(self, write_type):
 		"""
 		``get_flag_write_type_name`` gets the flag write type name for the given flag.
diff --git a/python/function.py b/python/function.py
index 10583b58..7607657f 100644
--- a/python/function.py
+++ b/python/function.py
@@ -427,7 +427,7 @@ class Function(object):
 		"""Function platform (read-only)"""
 		if self._platform:
 			return self._platform
-		else: 
+		else:
 			plat = core.BNGetFunctionPlatform(self.handle)
 			if plat is None:
 				return None
@@ -1878,6 +1878,27 @@ class RegisterStackInfo(object):
 		return "<reg stack: %d regs, stack top in %s>" % (len(self.storage_regs), self.stack_top_reg)
 
 
+class IntrinsicInput(object):
+	def __init__(self, type_obj, name=""):
+		self.name = name
+		self.type = type_obj
+
+	def __repr__(self):
+		if len(self.name) == 0:
+			return "<input: %s>" % str(self.type)
+		return "<input: %s %s>" % (str(self.type), self.name)
+
+
+class IntrinsicInfo(object):
+	def __init__(self, inputs, outputs, index=None):
+		self.inputs = inputs
+		self.outputs = outputs
+		self.index = index
+
+	def __repr__(self):
+		return "<intrinsic: %s -> %s>" % (repr(self.inputs), repr(self.outputs))
+
+
 class InstructionBranch(object):
 	def __init__(self, branch_type, target = 0, arch = None):
 		self.type = branch_type
diff --git a/python/lowlevelil.py b/python/lowlevelil.py
index a2d77c9f..33296ce8 100644
--- a/python/lowlevelil.py
+++ b/python/lowlevelil.py
@@ -131,6 +131,25 @@ class ILSemanticFlagGroup(object):
 		return self.index == other.index
 
 
+class ILIntrinsic(object):
+	def __init__(self, arch, intrinsic):
+		self.arch = arch
+		self.index = intrinsic
+		self.name = self.arch.get_intrinsic_name(self.index)
+		if self.name in self.arch.intrinsics:
+			self.inputs = self.arch.intrinsics[self.name].inputs
+			self.outputs = self.arch.intrinsics[self.name].outputs
+
+	def __str__(self):
+		return self.name
+
+	def __repr__(self):
+		return self.name
+
+	def __eq__(self, other):
+		return self.index == other.index
+
+
 class SSARegister(object):
 	def __init__(self, reg, version):
 		self.reg = reg
@@ -158,6 +177,15 @@ class SSAFlag(object):
 		return "<ssa %s version %d>" % (repr(self.flag), self.version)
 
 
+class SSARegisterOrFlag(object):
+	def __init__(self, reg_or_flag, version):
+		self.reg_or_flag = reg_or_flag
+		self.version = version
+
+	def __repr__(self):
+		return "<ssa %s version %d>" % (repr(self.reg_or_flag), self.version)
+
+
 class LowLevelILOperationAndSize(object):
 	def __init__(self, operation, size):
 		self.operation = operation
@@ -250,6 +278,8 @@ class LowLevelILInstruction(object):
 		LowLevelILOperation.LLIL_BOOL_TO_INT: [("src", "expr")],
 		LowLevelILOperation.LLIL_ADD_OVERFLOW: [("left", "expr"), ("right", "expr")],
 		LowLevelILOperation.LLIL_SYSCALL: [],
+		LowLevelILOperation.LLIL_INTRINSIC: [("output", "reg_or_flag_list"), ("intrinsic", "intrinsic"), ("param", "expr")],
+		LowLevelILOperation.LLIL_INTRINSIC_SSA: [("output", "reg_or_flag_ssa_list"), ("intrinsic", "intrinsic"), ("param", "expr")],
 		LowLevelILOperation.LLIL_BP: [],
 		LowLevelILOperation.LLIL_TRAP: [("vector", "int")],
 		LowLevelILOperation.LLIL_UNDEF: [],
@@ -292,7 +322,7 @@ class LowLevelILInstruction(object):
 		LowLevelILOperation.LLIL_SYSCALL_SSA: [("output", "expr"), ("stack", "expr"), ("param", "expr")],
 		LowLevelILOperation.LLIL_CALL_OUTPUT_SSA: [("dest_memory", "int"), ("dest", "reg_ssa_list")],
 		LowLevelILOperation.LLIL_CALL_STACK_SSA: [("src", "reg_ssa"), ("src_memory", "int")],
-		LowLevelILOperation.LLIL_CALL_PARAM_SSA: [("src", "expr_list")],
+		LowLevelILOperation.LLIL_CALL_PARAM: [("src", "expr_list")],
 		LowLevelILOperation.LLIL_LOAD_SSA: [("src", "expr"), ("src_memory", "int")],
 		LowLevelILOperation.LLIL_STORE_SSA: [("dest", "expr"), ("dest_memory", "int"), ("src_memory", "int"), ("src", "expr")],
 		LowLevelILOperation.LLIL_REG_PHI: [("dest", "reg_ssa"), ("src", "reg_ssa_list")],
@@ -336,6 +366,8 @@ class LowLevelILInstruction(object):
 				value = ILRegister(func.arch, instr.operands[i])
 			elif operand_type == "reg_stack":
 				value = ILRegisterStack(func.arch, instr.operands[i])
+			elif operand_type == "intrinsic":
+				value = ILIntrinsic(func.arch, instr.operands[i])
 			elif operand_type == "reg_ssa":
 				reg = ILRegister(func.arch, instr.operands[i])
 				i += 1
@@ -369,25 +401,36 @@ class LowLevelILInstruction(object):
 				operand_list = core.BNLowLevelILGetOperandList(func.handle, self.expr_index, i, count)
 				i += 1
 				value = []
-				for i in xrange(count.value):
-					value.append(operand_list[i])
+				for j in xrange(count.value):
+					value.append(operand_list[j])
 				core.BNLowLevelILFreeOperandList(operand_list)
 			elif operand_type == "expr_list":
 				count = ctypes.c_ulonglong()
 				operand_list = core.BNLowLevelILGetOperandList(func.handle, self.expr_index, i, count)
 				i += 1
 				value = []
-				for i in xrange(count.value):
-					value.append(LowLevelILInstruction(func, operand_list[i]))
+				for j in xrange(count.value):
+					value.append(LowLevelILInstruction(func, operand_list[j]))
+				core.BNLowLevelILFreeOperandList(operand_list)
+			elif operand_type == "reg_or_flag_list":
+				count = ctypes.c_ulonglong()
+				operand_list = core.BNLowLevelILGetOperandList(func.handle, self.expr_index, i, count)
+				i += 1
+				value = []
+				for j in xrange(count.value):
+					if (operand_list[j] & (1 << 32)) != 0:
+						value.append(ILFlag(func.arch, operand_list[j] & 0xffffffff))
+					else:
+						value.append(ILRegister(func.arch, operand_list[j] & 0xffffffff))
 				core.BNLowLevelILFreeOperandList(operand_list)
 			elif operand_type == "reg_ssa_list":
 				count = ctypes.c_ulonglong()
 				operand_list = core.BNLowLevelILGetOperandList(func.handle, self.expr_index, i, count)
 				i += 1
 				value = []
-				for i in xrange(count.value / 2):
-					reg = operand_list[i * 2]
-					reg_version = operand_list[(i * 2) + 1]
+				for j in xrange(count.value / 2):
+					reg = operand_list[j * 2]
+					reg_version = operand_list[(j * 2) + 1]
 					value.append(SSARegister(ILRegister(func.arch, reg), reg_version))
 				core.BNLowLevelILFreeOperandList(operand_list)
 			elif operand_type == "reg_stack_ssa_list":
@@ -395,9 +438,9 @@ class LowLevelILInstruction(object):
 				operand_list = core.BNLowLevelILGetOperandList(func.handle, self.expr_index, i, count)
 				i += 1
 				value = []
-				for i in xrange(count.value / 2):
-					reg_stack = operand_list[i * 2]
-					reg_version = operand_list[(i * 2) + 1]
+				for j in xrange(count.value / 2):
+					reg_stack = operand_list[j * 2]
+					reg_version = operand_list[(j * 2) + 1]
 					value.append(SSARegisterStack(ILRegisterStack(func.arch, reg_stack), reg_version))
 				core.BNLowLevelILFreeOperandList(operand_list)
 			elif operand_type == "flag_ssa_list":
@@ -405,19 +448,32 @@ class LowLevelILInstruction(object):
 				operand_list = core.BNLowLevelILGetOperandList(func.handle, self.expr_index, i, count)
 				i += 1
 				value = []
-				for i in xrange(count.value / 2):
-					flag = operand_list[i * 2]
-					flag_version = operand_list[(i * 2) + 1]
+				for j in xrange(count.value / 2):
+					flag = operand_list[j * 2]
+					flag_version = operand_list[(j * 2) + 1]
 					value.append(SSAFlag(ILFlag(func.arch, flag), flag_version))
 				core.BNLowLevelILFreeOperandList(operand_list)
+			elif operand_type == "reg_or_flag_ssa_list":
+				count = ctypes.c_ulonglong()
+				operand_list = core.BNLowLevelILGetOperandList(func.handle, self.expr_index, i, count)
+				i += 1
+				value = []
+				for j in xrange(count.value / 2):
+					if (operand_list[j * 2] & (1 << 32)) != 0:
+						reg_or_flag = ILFlag(func.arch, operand_list[j * 2] & 0xffffffff)
+					else:
+						reg_or_flag = ILRegister(func.arch, operand_list[j * 2] & 0xffffffff)
+					reg_version = operand_list[(j * 2) + 1]
+					value.append(SSARegisterOrFlag(reg_or_flag, reg_version))
+				core.BNLowLevelILFreeOperandList(operand_list)
 			elif operand_type == "reg_stack_adjust":
 				count = ctypes.c_ulonglong()
 				operand_list = core.BNLowLevelILGetOperandList(func.handle, self.expr_index, i, count)
 				i += 1
 				value = {}
-				for i in xrange(count.value / 2):
-					reg_stack = operand_list[i * 2]
-					adjust = operand_list[(i * 2) + 1]
+				for j in xrange(count.value / 2):
+					reg_stack = operand_list[j * 2]
+					adjust = operand_list[(j * 2) + 1]
 					if adjust & 0x80000000:
 						adjust |= ~0x80000000
 					value[func.arch.get_reg_stack_name(reg_stack)] = adjust
@@ -1714,6 +1770,25 @@ class LowLevelILFunction(object):
 		"""
 		return self.expr(LowLevelILOperation.LLIL_SYSCALL)
 
+	def intrinsic(self, outputs, intrinsic, params):
+		"""
+		``intrinsic`` return an intrinsic expression.
+
+		:return: an intrinsic expression.
+		:rtype: LowLevelILExpr
+		"""
+		output_list = []
+		for output in outputs:
+			if isinstance(output, ILFlag):
+				output_list.append((1 << 32) | output.index)
+			else:
+				output_list.append(output.index)
+		param_list = []
+		for param in params:
+			param_list.append(param.index)
+		return self.expr(LowLevelILOperation.LLIL_INTRINSIC, len(outputs), self.add_operand_list(output_list),
+			self.arch.get_intrinsic_index(intrinsic), len(params), self.add_operand_list(param_list))
+
 	def breakpoint(self):
 		"""
 		``breakpoint`` returns a processor breakpoint expression.
diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py
index 8594ea36..3c4ba3fd 100644
--- a/python/mediumlevelil.py
+++ b/python/mediumlevelil.py
@@ -150,6 +150,8 @@ class MediumLevelILInstruction(object):
 		MediumLevelILOperation.MLIL_SYSCALL_UNTYPED: [("output", "expr"), ("params", "expr"), ("stack", "expr")],
 		MediumLevelILOperation.MLIL_BP: [],
 		MediumLevelILOperation.MLIL_TRAP: [("vector", "int")],
+		MediumLevelILOperation.MLIL_INTRINSIC: [("output", "var_list"), ("intrinsic", "intrinsic"), ("params", "expr_list")],
+		MediumLevelILOperation.MLIL_INTRINSIC_SSA: [("output", "var_ssa_list"), ("intrinsic", "intrinsic"), ("params", "expr_list")],
 		MediumLevelILOperation.MLIL_UNDEF: [],
 		MediumLevelILOperation.MLIL_UNIMPL: [],
 		MediumLevelILOperation.MLIL_UNIMPL_MEM: [("src", "expr")],
@@ -223,6 +225,8 @@ class MediumLevelILInstruction(object):
 					value = instr.operands[i]
 			elif operand_type == "expr":
 				value = MediumLevelILInstruction(func, instr.operands[i])
+			elif operand_type == "intrinsic":
+				value = lowlevelil.ILIntrinsic(func.arch, instr.operands[i])
 			elif operand_type == "var":
 				value = function.Variable.from_identifier(self.function.source_function, instr.operands[i])
 			elif operand_type == "var_ssa":