From d27bf51d4b7c2170d00dca7cce97c7526c866868 Mon Sep 17 00:00:00 2001
From: Brian Potchik <brian@vector35.com>
Date: Wed, 23 Dec 2020 23:35:48 -0500
Subject: Add GetLowLevelILInstructionsForAddress API to retrieve all LLIL for
 an address.

---
 binaryninjaapi.h   |  1 +
 binaryninjacore.h  |  4 ++--
 function.cpp       | 14 ++++++++++++++
 python/function.py | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 224ea01c..5b3307a4 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -3047,6 +3047,7 @@ __attribute__ ((format (printf, 1, 2)))
 		Ref<LowLevelILFunction> GetLowLevelIL() const;
 		Ref<LowLevelILFunction> GetLowLevelILIfAvailable() const;
 		size_t GetLowLevelILForInstruction(Architecture* arch, uint64_t addr);
+		std::set<size_t> GetLowLevelILInstructionsForAddress(Architecture* arch, uint64_t addr);
 		std::vector<size_t> GetLowLevelILExitsForInstruction(Architecture* arch, uint64_t addr);
 		RegisterValue GetRegisterValueAtInstruction(Architecture* arch, uint64_t addr, uint32_t reg);
 		RegisterValue GetRegisterValueAfterInstruction(Architecture* arch, uint64_t addr, uint32_t reg);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 72f661a9..e90c4294 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -3076,8 +3076,8 @@ __attribute__ ((format (printf, 1, 2)))
 	BINARYNINJACOREAPI BNLowLevelILFunction* BNGetFunctionLowLevelIL(BNFunction* func);
 	BINARYNINJACOREAPI BNLowLevelILFunction* BNGetFunctionLowLevelILIfAvailable(BNFunction* func);
 	BINARYNINJACOREAPI size_t BNGetLowLevelILForInstruction(BNFunction* func, BNArchitecture* arch, uint64_t addr);
-	BINARYNINJACOREAPI size_t* BNGetLowLevelILExitsForInstruction(BNFunction* func, BNArchitecture* arch, uint64_t addr,
-	                                                              size_t* count);
+	BINARYNINJACOREAPI size_t* BNGetLowLevelILInstructionsForAddress(BNFunction* func, BNArchitecture* arch, uint64_t addr, size_t* count);
+	BINARYNINJACOREAPI size_t* BNGetLowLevelILExitsForInstruction(BNFunction* func, BNArchitecture* arch, uint64_t addr, size_t* count);
 	BINARYNINJACOREAPI void BNFreeILInstructionList(size_t* list);
 	BINARYNINJACOREAPI BNMediumLevelILFunction* BNGetFunctionMediumLevelIL(BNFunction* func);
 	BINARYNINJACOREAPI BNMediumLevelILFunction* BNGetFunctionMediumLevelILIfAvailable(BNFunction* func);
diff --git a/function.cpp b/function.cpp
index 340f1e53..d7345248 100644
--- a/function.cpp
+++ b/function.cpp
@@ -307,6 +307,20 @@ size_t Function::GetLowLevelILForInstruction(Architecture* arch, uint64_t addr)
 }
 
 
+set<size_t> Function::GetLowLevelILInstructionsForAddress(Architecture* arch, uint64_t addr)
+{
+	size_t count;
+	size_t* instrs = BNGetLowLevelILInstructionsForAddress(m_object, arch->GetObject(), addr, &count);
+
+	set<size_t> result;
+	for (size_t i = 0; i < count; i++)
+		result.insert(instrs[i]);
+
+	BNFreeILInstructionList(instrs);
+	return result;
+}
+
+
 vector<size_t> Function::GetLowLevelILExitsForInstruction(Architecture* arch, uint64_t addr)
 {
 	size_t count;
diff --git a/python/function.py b/python/function.py
index d986fcca..13d73435 100644
--- a/python/function.py
+++ b/python/function.py
@@ -2040,6 +2040,52 @@ class Function(object):
 
 		return self.llil[idx]
 
+	def get_llil_at(self, addr, arch=None):
+		"""
+		``get_llil_at`` gets the LowLevelILInstruction corresponding to the given virtual address
+
+		:param int addr: virtual address of the function to be queried
+		:param Architecture arch: (optional) Architecture for the given function
+		:rtype: LowLevelILInstruction
+		:Example:
+
+			>>> func = bv.functions[0]
+			>>> func.get_llil_at(func.start)
+			<il: push(rbp)>
+		"""
+		if arch is None:
+			arch = self.arch
+
+		idx = core.BNGetLowLevelILForInstruction(self.handle, arch.handle, addr)
+
+		if idx == len(self.llil):
+			return None
+
+		return self.llil[idx]
+
+	def get_llils_at(self, addr, arch=None):
+		"""
+		``get_llils_at`` gets the LowLevelILInstruction(s) corresponding to the given virtual address
+
+		:param int addr: virtual address of the function to be queried
+		:param Architecture arch: (optional) Architecture for the given function
+		:rtype: list(LowLevelILInstruction)
+		:Example:
+
+			>>> func = bv.functions[0]
+			>>> func.get_llils_at(func.start)
+			[<il: push(rbp)>]
+		"""
+		if arch is None:
+			arch = self.arch
+		count = ctypes.c_ulonglong()
+		instrs = core.BNGetLowLevelILInstructionsForAddress(self.handle, arch.handle, addr, count)
+		result = []
+		for i in range(0, count.value):
+			result.append(self.llil[instrs[i]])
+		core.BNFreeILInstructionList(instrs)
+		return result
+
 	def get_low_level_il_exits_at(self, addr, arch=None):
 		if arch is None:
 			arch = self.arch
-- 
cgit v1.3.1