Add missing Python APIs

6 files changed, 90 insertions, 11 deletions

diff --git a/python/architecture.py b/python/architecture.py
index 32d97b09..ea97f605 100644
--- a/python/architecture.py
+++ b/python/architecture.py
@@ -467,6 +467,8 @@ class Architecture(object):
 				token_buf[i].value = tokens[i].value
 				token_buf[i].size = tokens[i].size
 				token_buf[i].operand = tokens[i].operand
+				token_buf[i].context = tokens[i].context
+				token_buf[i].address = tokens[i].address
 			result[0] = token_buf
 			ptr = ctypes.cast(token_buf, ctypes.c_void_p)
 			self._pending_token_lists[ptr.value] = (ptr.value, token_buf)
@@ -1148,7 +1150,9 @@ class Architecture(object):
 			value = tokens[i].value
 			size = tokens[i].size
 			operand = tokens[i].operand
-			result.append(function.InstructionTextToken(token_type, text, value, size, operand))
+			context = tokens[i].context
+			address = tokens[i].address
+			result.append(function.InstructionTextToken(token_type, text, value, size, operand, context, address))
 		core.BNFreeInstructionText(tokens, count.value)
 		return result, length.value
 
diff --git a/python/basicblock.py b/python/basicblock.py
index ae9889fc..ffbcd217 100644
--- a/python/basicblock.py
+++ b/python/basicblock.py
@@ -200,7 +200,9 @@ class BasicBlock(object):
 				value = lines[i].tokens[j].value
 				size = lines[i].tokens[j].size
 				operand = lines[i].tokens[j].operand
-				tokens.append(function.InstructionTextToken(token_type, text, value, size, operand))
+				context = lines[i].tokens[j].context
+				address = lines[i].tokens[j].address
+				tokens.append(function.InstructionTextToken(token_type, text, value, size, operand, context, address))
 			result.append(function.DisassemblyTextLine(addr, tokens))
 		core.BNFreeDisassemblyTextLines(lines, count.value)
 		return result
diff --git a/python/binaryview.py b/python/binaryview.py
index cf30d78d..f8fe1fde 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -1886,7 +1886,7 @@ class BinaryView(object):
 		var = core.BNDataVariable()
 		if not core.BNGetDataVariableAtAddress(self.handle, addr, var):
 			return None
-		return DataVariable(var.address, type.Type(var.type), var.autoDiscovered)
+		return DataVariable(var.address, types.Type(var.type), var.autoDiscovered)
 
 	def get_function_at(self, addr, plat=None):
 		"""
@@ -2750,7 +2750,9 @@ class BinaryView(object):
 				value = lines[i].contents.tokens[j].value
 				size = lines[i].contents.tokens[j].size
 				operand = lines[i].contents.tokens[j].operand
-				tokens.append(function.InstructionTextToken(token_type, text, value, size, operand))
+				context = lines[i].contents.tokens[j].context
+				address = lines[i].contents.tokens[j].address
+				tokens.append(function.InstructionTextToken(token_type, text, value, size, operand, context, address))
 			contents = function.DisassemblyTextLine(addr, tokens)
 			result.append(lineardisassembly.LinearDisassemblyLine(lines[i].type, func, block, lines[i].lineOffset, contents))
 
diff --git a/python/function.py b/python/function.py
index 7d03585c..5f9f475a 100644
--- a/python/function.py
+++ b/python/function.py
@@ -26,7 +26,7 @@ import ctypes
 import _binaryninjacore as core
 from enums import (FunctionGraphType, BranchType, SymbolType, InstructionTextTokenType,
 	HighlightStandardColor, HighlightColorStyle, RegisterValueType, ImplicitRegisterExtend,
-	DisassemblyOption, IntegerDisplayType)
+	DisassemblyOption, IntegerDisplayType, InstructionTextTokenContext)
 import architecture
 import highlight
 import associateddatastore
@@ -669,7 +669,9 @@ class Function(object):
 				value = lines[i].tokens[j].value
 				size = lines[i].tokens[j].size
 				operand = lines[i].tokens[j].operand
-				tokens.append(InstructionTextToken(token_type, text, value, size, operand))
+				context = lines[i].tokens[j].context
+				address = lines[i].tokens[j].address
+				tokens.append(InstructionTextToken(token_type, text, value, size, operand, context, address))
 			result.append(tokens)
 		core.BNFreeInstructionTextLines(lines, count.value)
 		return result
@@ -916,7 +918,9 @@ class FunctionGraphBlock(object):
 				value = lines[i].tokens[j].value
 				size = lines[i].tokens[j].size
 				operand = lines[i].tokens[j].operand
-				tokens.append(InstructionTextToken(token_type, text, value, size, operand))
+				context = lines[i].tokens[j].context
+				address = lines[i].tokens[j].address
+				tokens.append(InstructionTextToken(token_type, text, value, size, operand, context, address))
 			result.append(DisassemblyTextLine(addr, tokens))
 		core.BNFreeDisassemblyTextLines(lines, count.value)
 		return result
@@ -966,7 +970,9 @@ class FunctionGraphBlock(object):
 					value = lines[i].tokens[j].value
 					size = lines[i].tokens[j].size
 					operand = lines[i].tokens[j].operand
-					tokens.append(InstructionTextToken(token_type, text, value, size, operand))
+					context = lines[i].tokens[j].context
+					address = lines[i].tokens[j].address
+					tokens.append(InstructionTextToken(token_type, text, value, size, operand, context, address))
 				yield DisassemblyTextLine(addr, tokens)
 		finally:
 			core.BNFreeDisassemblyTextLines(lines, count.value)
@@ -1237,12 +1243,15 @@ class InstructionTextToken(object):
 		========================== ============================================
 
 	"""
-	def __init__(self, token_type, text, value = 0, size = 0, operand = 0xffffffff):
+	def __init__(self, token_type, text, value = 0, size = 0, operand = 0xffffffff,
+		context = InstructionTextTokenContext.NoTokenContext, address = 0):
 		self.type = InstructionTextTokenType(token_type)
 		self.text = text
 		self.value = value
 		self.size = size
 		self.operand = operand
+		self.context = InstructionTextTokenContext(context)
+		self.address = address
 
 	def __str__(self):
 		return self.text
diff --git a/python/lowlevelil.py b/python/lowlevelil.py
index 419e8513..29b46b65 100644
--- a/python/lowlevelil.py
+++ b/python/lowlevelil.py
@@ -187,7 +187,9 @@ class LowLevelILInstruction(object):
 			value = tokens[i].value
 			size = tokens[i].size
 			operand = tokens[i].operand
-			result.append(function.InstructionTextToken(token_type, text, value, size, operand))
+			context = tokens[i].context
+			address = tokens[i].address
+			result.append(function.InstructionTextToken(token_type, text, value, size, operand, context, address))
 		core.BNFreeInstructionText(tokens, count.value)
 		return result
 
diff --git a/python/types.py b/python/types.py
index 0470ea73..1dd2157e 100644
--- a/python/types.py
+++ b/python/types.py
@@ -22,8 +22,9 @@ import ctypes
 
 # Binary Ninja components
 import _binaryninjacore as core
-from enums import SymbolType, TypeClass, NamedTypeReferenceClass
+from enums import SymbolType, TypeClass, NamedTypeReferenceClass, InstructionTextTokenType
 import callingconvention
+import function
 
 
 class QualifiedName(object):
@@ -317,6 +318,56 @@ class Type(object):
 	def get_string_after_name(self):
 		return core.BNGetTypeStringAfterName(self.handle)
 
+	@property
+	def tokens(self):
+		"""Type string as a list of tokens (read-only)"""
+		count = ctypes.c_ulonglong()
+		tokens = core.BNGetTypeTokens(self.handle, count)
+		result = []
+		for i in xrange(0, count.value):
+			token_type = InstructionTextTokenType(tokens[i].type)
+			text = tokens[i].text
+			value = tokens[i].value
+			size = tokens[i].size
+			operand = tokens[i].operand
+			context = tokens[i].context
+			address = tokens[i].address
+			result.append(function.InstructionTextToken(token_type, text, value, size, operand, context, address))
+		core.BNFreeTokenList(tokens, count.value)
+		return result
+
+	def get_tokens_before_name(self):
+		count = ctypes.c_ulonglong()
+		tokens = core.BNGetTypeTokensBeforeName(self.handle, count)
+		result = []
+		for i in xrange(0, count.value):
+			token_type = InstructionTextTokenType(tokens[i].type)
+			text = tokens[i].text
+			value = tokens[i].value
+			size = tokens[i].size
+			operand = tokens[i].operand
+			context = tokens[i].context
+			address = tokens[i].address
+			result.append(function.InstructionTextToken(token_type, text, value, size, operand, context, address))
+		core.BNFreeTokenList(tokens, count.value)
+		return result
+
+	def get_tokens_after_name(self):
+		count = ctypes.c_ulonglong()
+		tokens = core.BNGetTypeTokensAfterName(self.handle, count)
+		result = []
+		for i in xrange(0, count.value):
+			token_type = InstructionTextTokenType(tokens[i].type)
+			text = tokens[i].text
+			value = tokens[i].value
+			size = tokens[i].size
+			operand = tokens[i].operand
+			context = tokens[i].context
+			address = tokens[i].address
+			result.append(function.InstructionTextToken(token_type, text, value, size, operand, context, address))
+		core.BNFreeTokenList(tokens, count.value)
+		return result
+
 	@classmethod
 	def void(cls):
 		return Type(core.BNCreateVoidType())
@@ -518,6 +569,9 @@ class Structure(object):
 	def remove(self, i):
 		core.BNRemoveStructureMember(self.handle, i)
 
+	def replace(self, i, t, name = ""):
+		core.BNReplaceStructureMember(self.handle, i, t.handle, name)
+
 
 class EnumerationMember(object):
 	def __init__(self, name, value, default):
@@ -565,6 +619,12 @@ class Enumeration(object):
 		else:
 			core.BNAddEnumerationMemberWithValue(self.handle, name, value)
 
+	def remove(self, i):
+		core.BNRemoveEnumerationMember(self.handle, i)
+
+	def replace(self, i, name, value):
+		core.BNReplaceEnumerationMember(self.handle, i, name, value)
+
 
 class TypeParserResult(object):
 	def __init__(self, types, variables, functions):