Add APIs for accessing linear disassembly

4 files changed, 311 insertions, 16 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 9e8f516c..47490857 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -612,6 +612,41 @@ namespace BinaryNinja
 		Ref<Architecture> arch;
 		uint64_t addr;
 	};
+
+	struct InstructionTextToken
+	{
+		BNInstructionTextTokenType type;
+		std::string text;
+		uint64_t value;
+
+		InstructionTextToken();
+		InstructionTextToken(BNInstructionTextTokenType type, const std::string& text, uint64_t value = 0);
+	};
+
+	struct DisassemblyTextLine
+	{
+		uint64_t addr;
+		std::vector<InstructionTextToken> tokens;
+	};
+
+	struct LinearDisassemblyPosition
+	{
+		Ref<Function> function;
+		Ref<BasicBlock> block;
+		uint64_t address;
+	};
+
+	struct LinearDisassemblyLine
+	{
+		BNLinearDisassemblyLineType type;
+		Ref<Function> function;
+		Ref<BasicBlock> block;
+		size_t lineOffset;
+		DisassemblyTextLine contents;
+	};
+
+	class DisassemblySettings;
+
 	class AnalysisCompletionEvent: public CoreRefCountObject<BNAnalysisCompletionEvent,
 		BNNewAnalysisCompletionEventReference, BNFreeAnalysisCompletionEvent>
 	{
@@ -811,6 +846,16 @@ namespace BinaryNinja
 		Ref<AnalysisCompletionEvent> AddAnalysisCompletionEvent(const std::function<void()>& callback);
 
 		BNAnalysisProgress GetAnalysisProgress();
+
+		uint64_t GetNextFunctionStartAfterAddress(uint64_t addr);
+		uint64_t GetNextBasicBlockStartAfterAddress(uint64_t addr);
+		uint64_t GetNextDataAfterAddress(uint64_t addr);
+
+		LinearDisassemblyPosition GetLinearDisassemblyPositionForAddress(uint64_t addr);
+		std::vector<LinearDisassemblyLine> GetPreviousLinearDisassemblyLines(LinearDisassemblyPosition& pos,
+			DisassemblySettings* settings);
+		std::vector<LinearDisassemblyLine> GetNextLinearDisassemblyLines(LinearDisassemblyPosition& pos,
+			DisassemblySettings* settings);
 	};
 
 	class BinaryData: public BinaryView
@@ -1035,16 +1080,6 @@ namespace BinaryNinja
 		void AddBranch(BNBranchType type, uint64_t target = 0, Architecture* arch = nullptr, bool hasDelaySlot = false);
 	};
 
-	struct InstructionTextToken
-	{
-		BNInstructionTextTokenType type;
-		std::string text;
-		uint64_t value;
-
-		InstructionTextToken();
-		InstructionTextToken(BNInstructionTextTokenType type, const std::string& text, uint64_t value = 0);
-	};
-
 	class LowLevelILFunction;
 	class FunctionRecognizer;
 	class CallingConvention;
@@ -1396,12 +1431,6 @@ namespace BinaryNinja
 		void SetMaximumSymbolWidth(size_t width);
 	};
 
-	struct DisassemblyTextLine
-	{
-		uint64_t addr;
-		std::vector<InstructionTextToken> tokens;
-	};
-
 	class Function;
 
 	struct BasicBlockEdge
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 805ff090..9784b4ab 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -169,6 +169,16 @@ extern "C"
 		AddressDisplayToken = 68
 	};
 
+	enum BNLinearDisassemblyLineType
+	{
+		CodeDisassemblyLineType,
+		DataVariableLineType,
+		HexDumpLineType,
+		FunctionHeaderLineType,
+		FunctionContinuationLineType,
+		StackVariableLineType
+	};
+
 	enum BNSymbolType
 	{
 		FunctionSymbol = 0,
@@ -299,6 +309,9 @@ extern "C"
 	{
 		ShowAddress = 0,
 
+		// Linear disassembly options
+		GroupLinearDisassemblyFunctions = 64,
+
 		// Debugging options
 		ShowBasicBlockRegisterState = 128,
 		ShowFlagUsage = 129
@@ -580,6 +593,22 @@ extern "C"
 		size_t count;
 	};
 
+	struct BNLinearDisassemblyLine
+	{
+		BNLinearDisassemblyLineType type;
+		BNFunction* function;
+		BNBasicBlock* block;
+		size_t lineOffset;
+		BNDisassemblyTextLine contents;
+	};
+
+	struct BNLinearDisassemblyPosition
+	{
+		BNFunction* function;
+		BNBasicBlock* block;
+		uint64_t address;
+	};
+
 	struct BNReferenceSource
 	{
 		BNFunction* func;
@@ -1231,6 +1260,19 @@ extern "C"
 
 	BINARYNINJACOREAPI BNAnalysisProgress BNGetAnalysisProgress(BNBinaryView* view);
 
+	BINARYNINJACOREAPI uint64_t BNGetNextFunctionStartAfterAddress(BNBinaryView* view, uint64_t addr);
+	BINARYNINJACOREAPI uint64_t BNGetNextBasicBlockStartAfterAddress(BNBinaryView* view, uint64_t addr);
+	BINARYNINJACOREAPI uint64_t BNGetNextDataAfterAddress(BNBinaryView* view, uint64_t addr);
+
+	BINARYNINJACOREAPI BNLinearDisassemblyPosition BNGetLinearDisassemblyPositionForAddress(BNBinaryView* view,
+		uint64_t addr);
+	BINARYNINJACOREAPI void BNFreeLinearDisassemblyPosition(BNLinearDisassemblyPosition* pos);
+	BINARYNINJACOREAPI BNLinearDisassemblyLine* BNGetPreviousLinearDisassemblyLines(BNBinaryView* view,
+		BNLinearDisassemblyPosition* pos, BNDisassemblySettings* settings, size_t* count);
+	BINARYNINJACOREAPI BNLinearDisassemblyLine* BNGetNextLinearDisassemblyLines(BNBinaryView* view,
+		BNLinearDisassemblyPosition* pos, BNDisassemblySettings* settings, size_t* count);
+	BINARYNINJACOREAPI void BNFreeLinearDisassemblyLines(BNLinearDisassemblyLine* lines, size_t count);
+
 	// Disassembly settings
 	BINARYNINJACOREAPI BNDisassemblySettings* BNCreateDisassemblySettings(void);
 	BINARYNINJACOREAPI BNDisassemblySettings* BNNewDisassemblySettingsReference(BNDisassemblySettings* settings);
diff --git a/binaryview.cpp b/binaryview.cpp
index a3abc02a..437c57a6 100644
--- a/binaryview.cpp
+++ b/binaryview.cpp
@@ -1100,6 +1100,116 @@ BNAnalysisProgress BinaryView::GetAnalysisProgress()
 }
 
 
+uint64_t BinaryView::GetNextFunctionStartAfterAddress(uint64_t addr)
+{
+	return BNGetNextFunctionStartAfterAddress(m_object, addr);
+}
+
+
+uint64_t BinaryView::GetNextBasicBlockStartAfterAddress(uint64_t addr)
+{
+	return BNGetNextBasicBlockStartAfterAddress(m_object, addr);
+}
+
+
+uint64_t BinaryView::GetNextDataAfterAddress(uint64_t addr)
+{
+	return BNGetNextDataAfterAddress(m_object, addr);
+}
+
+
+LinearDisassemblyPosition BinaryView::GetLinearDisassemblyPositionForAddress(uint64_t addr)
+{
+	BNLinearDisassemblyPosition pos = BNGetLinearDisassemblyPositionForAddress(m_object, addr);
+
+	LinearDisassemblyPosition result;
+	result.function = pos.function ? new Function(pos.function) : nullptr;
+	result.block = pos.block ? new BasicBlock(pos.block) : nullptr;
+	result.address = pos.address;
+	return result;
+}
+
+
+vector<LinearDisassemblyLine> BinaryView::GetPreviousLinearDisassemblyLines(LinearDisassemblyPosition& pos,
+	DisassemblySettings* settings)
+{
+	BNLinearDisassemblyPosition linearPos;
+	linearPos.function = pos.function ? BNNewFunctionReference(pos.function->GetObject()) : nullptr;
+	linearPos.block = pos.block ? BNNewBasicBlockReference(pos.block->GetObject()) : nullptr;
+	linearPos.address = pos.address;
+
+	size_t count;
+	BNLinearDisassemblyLine* lines = BNGetPreviousLinearDisassemblyLines(m_object, &linearPos,
+		settings ? settings->GetObject() : nullptr, &count);
+
+	vector<LinearDisassemblyLine> result;
+	for (size_t i = 0; i < count; i++)
+	{
+		LinearDisassemblyLine line;
+		line.function = lines[i].function ? new Function(BNNewFunctionReference(lines[i].function)) : nullptr;
+		line.block = lines[i].block ? new BasicBlock(BNNewBasicBlockReference(lines[i].block)) : nullptr;
+		line.lineOffset = lines[i].lineOffset;
+		line.contents.addr = lines[i].contents.addr;
+		for (size_t j = 0; j < lines[i].contents.count; j++)
+		{
+			InstructionTextToken token;
+			token.type = lines[i].contents.tokens[j].type;
+			token.text = lines[i].contents.tokens[j].text;
+			token.value = lines[i].contents.tokens[j].value;
+			line.contents.tokens.push_back(token);
+		}
+		result.push_back(line);
+	}
+
+	pos.function = linearPos.function ? new Function(linearPos.function) : nullptr;
+	pos.block = linearPos.block ? new BasicBlock(linearPos.block) : nullptr;
+	pos.address = linearPos.address;
+
+	BNFreeLinearDisassemblyLines(lines, count);
+	return result;
+}
+
+
+vector<LinearDisassemblyLine> BinaryView::GetNextLinearDisassemblyLines(LinearDisassemblyPosition& pos,
+	DisassemblySettings* settings)
+{
+	BNLinearDisassemblyPosition linearPos;
+	linearPos.function = pos.function ? BNNewFunctionReference(pos.function->GetObject()) : nullptr;
+	linearPos.block = pos.block ? BNNewBasicBlockReference(pos.block->GetObject()) : nullptr;
+	linearPos.address = pos.address;
+
+	size_t count;
+	BNLinearDisassemblyLine* lines = BNGetNextLinearDisassemblyLines(m_object, &linearPos,
+		settings ? settings->GetObject() : nullptr, &count);
+
+	vector<LinearDisassemblyLine> result;
+	for (size_t i = 0; i < count; i++)
+	{
+		LinearDisassemblyLine line;
+		line.function = lines[i].function ? new Function(BNNewFunctionReference(lines[i].function)) : nullptr;
+		line.block = lines[i].block ? new BasicBlock(BNNewBasicBlockReference(lines[i].block)) : nullptr;
+		line.lineOffset = lines[i].lineOffset;
+		line.contents.addr = lines[i].contents.addr;
+		for (size_t j = 0; j < lines[i].contents.count; j++)
+		{
+			InstructionTextToken token;
+			token.type = lines[i].contents.tokens[j].type;
+			token.text = lines[i].contents.tokens[j].text;
+			token.value = lines[i].contents.tokens[j].value;
+			line.contents.tokens.push_back(token);
+		}
+		result.push_back(line);
+	}
+
+	pos.function = linearPos.function ? new Function(linearPos.function) : nullptr;
+	pos.block = linearPos.block ? new BasicBlock(linearPos.block) : nullptr;
+	pos.address = linearPos.address;
+
+	BNFreeLinearDisassemblyLines(lines, count);
+	return result;
+}
+
+
 BinaryData::BinaryData(FileMetadata* file): BinaryView(BNCreateBinaryDataView(file->GetObject()))
 {
 }
diff --git a/python/__init__.py b/python/__init__.py
index 0c983fe2..1037e530 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -676,6 +676,25 @@ class AnalysisProgress(object):
 	def __repr__(self):
 		return "<progress: %s>" % str(self)
 
+class LinearDisassemblyPosition(object):
+	def __init__(self, func, block, addr):
+		self.function = func
+		self.block = block
+		self.address = addr
+
+class LinearDisassemblyLine(object):
+	def __init__(self, func, block, line_offset, contents):
+		self.function = func
+		self.block = block
+		self.line_offset = line_offset
+		self.contents = contents
+
+	def __str__(self):
+		return str(self.contents)
+
+	def __repr__(self):
+		return repr(self.contents)
+
 class BinaryView(object):
 	name = None
 	"""Binary View name"""
@@ -969,6 +988,11 @@ class BinaryView(object):
 		result = core.BNGetAnalysisProgress(self.handle)
 		return AnalysisProgress(result.state, result.count, result.total)
 
+	@property
+	def linear_disassembly(self):
+		"""Iterator for all lines in the linear disassembly of the view"""
+		return self.get_linear_disassembly(None)
+
 	def __len__(self):
 		return int(core.BNGetViewLength(self.handle))
 
@@ -1518,6 +1542,96 @@ class BinaryView(object):
 	def add_analysis_completion_event(self, callback):
 		return AnalysisCompletionEvent(self, callback)
 
+	def get_next_function_start_after(self, addr):
+		return core.BNGetNextFunctionStartAfterAddress(self.handle, addr)
+
+	def get_next_basic_block_start_after(self, addr):
+		return core.BNGetNextBasicBlockStartAfterAddress(self.handle, addr)
+
+	def get_next_data_after(self, addr):
+		return core.BNGetNextDataAfterAddress(self.handle, addr)
+
+	def get_linear_disassembly_position_at(self, addr):
+		pos = core.BNGetLinearDisassemblyPositionForAddress(self.handle, addr)
+		func = None
+		block = None
+		if pos.function:
+			func = Function(self, pos.function)
+		if pos.block:
+			block = BasicBlock(self, pos.block)
+		return LinearDisassemblyPosition(func, block, pos.address)
+
+	def _get_linear_disassembly_lines(self, api, pos, settings):
+		pos_obj = core.BNLinearDisassemblyPosition()
+		pos_obj.function = None
+		pos_obj.block = None
+		pos_obj.address = pos.address
+		if pos.function is not None:
+			pos_obj.function = core.BNNewFunctionReference(pos.function.handle)
+		if pos.block is not None:
+			pos_obj.block = core.BNNewBasicBlockReference(pos.block.handle)
+
+		if settings is not None:
+			settings = settings.handle
+
+		count = ctypes.c_ulonglong(0)
+		lines = api(self.handle, pos_obj, settings, count)
+
+		result = []
+		for i in xrange(0, count.value):
+			func = None
+			block = None
+			if lines[i].function:
+				func = Function(self, core.BNNewFunctionReference(lines[i].function))
+			if lines[i].block:
+				block = BasicBlock(self, core.BNNewBasicBlockReference(lines[i].block))
+			addr = lines[i].contents.addr
+			tokens = []
+			for j in xrange(0, lines[i].contents.count):
+				token_type = core.BNInstructionTextTokenType_names[lines[i].contents.tokens[j].type]
+				text = lines[i].contents.tokens[j].text
+				value = lines[i].contents.tokens[j].value
+				tokens.append(InstructionTextToken(token_type, text, value))
+			contents = DisassemblyTextLine(addr, tokens)
+			result.append(LinearDisassemblyLine(func, block, lines[i].lineOffset, contents))
+
+		func = None
+		block = None
+		if pos_obj.function:
+			func = Function(self, pos_obj.function)
+		if pos_obj.block:
+			block = BasicBlock(self, pos_obj.block)
+		pos.function = func
+		pos.block = block
+		pos.address = pos_obj.address
+
+		core.BNFreeLinearDisassemblyLines(lines, count.value)
+		return result
+
+	def get_previous_linear_disassembly_lines(self, pos, settings):
+		return self._get_linear_disassembly_lines(core.BNGetPreviousLinearDisassemblyLines, pos, settings)
+
+	def get_next_linear_disassembly_lines(self, pos, settings):
+		return self._get_linear_disassembly_lines(core.BNGetNextLinearDisassemblyLines, pos, settings)
+
+	def get_linear_disassembly(self, settings):
+		"""Gets an iterator for all lines in the linear disassembly of the view for the given disassembly settings"""
+		class LinearDisassemblyIterator(object):
+			def __init__(self, view, settings):
+				self.view = view
+				self.settings = settings
+
+			def __iter__(self):
+				pos = self.view.get_linear_disassembly_position_at(self.view.start)
+				while True:
+					lines = self.view.get_next_linear_disassembly_lines(pos, self.settings)
+					if len(lines) == 0:
+						break
+					for line in lines:
+						yield line
+
+		return iter(LinearDisassemblyIterator(self, settings))
+
 	def __setattr__(self, name, value):
 		try:
 			object.__setattr__(self,name,value)