2 files changed, 165 insertions, 160 deletions

diff --git a/view/kernelcache/api/python/kernelcache.py b/view/kernelcache/api/python/kernelcache.py
index 9425ef9c..25360213 100644
--- a/view/kernelcache/api/python/kernelcache.py
+++ b/view/kernelcache/api/python/kernelcache.py
@@ -1,229 +1,210 @@
-import os
 import ctypes
 import dataclasses
-import traceback
+from typing import Optional
 
 import binaryninja
+from binaryninja import BinaryView
 from binaryninja._binaryninjacore import BNFreeStringList, BNAllocString, BNFreeString
 
 from . import _kernelcachecore as kccore
 from .kernelcache_enums import *
 
-
 @dataclasses.dataclass
-class KCImageMemoryMapping:
+class CacheImage:
 	name: str
-	vmAddress: int
-	rawViewOffset: int
-	size: int
+	header_virtual_address: int
+	header_file_address: int
 
 	def __str__(self):
 		return repr(self)
 
 	def __repr__(self):
-		return (f"<KCImageMemoryMapping '{self.name}' {self.vmAddress:x}+{self.size:x} "
-				f"raw<{self.rawViewOffset:x}>>")
-
+		return f"<CacheImage '{self.name}': 0x{self.header_virtual_address:x} (@0x{self.header_file_address:x})>"
 
 @dataclasses.dataclass
-class KCImage:
+class CacheSymbol:
+	symbol_type: kccore.SymbolTypeEnum
+	address: int
 	name: str
-	headerFileAddress: int
-	mappings: list[KCImageMemoryMapping]
 
 	def __str__(self):
 		return repr(self)
 
 	def __repr__(self):
-		return f"<KCImage {self.name} @ {self.headerFileAddress:x}>"
+		return f"<CacheSymbol '{self.name}': 0x{self.address:x}>"
 
+def image_from_api(image: kccore.BNKernelCacheImage) -> CacheImage:
+	return CacheImage(
+		name=image.name,
+		header_file_address=image.headerFileAddress,
+		header_virtual_address=image.headerVirtualAddress
+	)
 
-@dataclasses.dataclass
-class KCSymbol:
-	name: str
-	image: str
-	address: int
+def image_to_api(image: CacheImage) -> kccore.BNKernelCacheImage:
+	return kccore.BNKernelCacheImage(
+		_name=BNAllocString(image.name),
+		headerFileAddress=image.header_file_address,
+		headerVirtualAddress=image.header_virtual_address,
+	)
 
-	def __str__(self):
-		return repr(self)
+def symbol_from_api(symbol: kccore.BNKernelCacheSymbol) -> CacheSymbol:
+	return CacheSymbol(
+		symbol_type=symbol.symbolType,
+		address=symbol.address,
+		name=symbol.name
+	)
 
-	def __repr__(self):
-		return f"<KCSymbol {self.name} @ {self.address:x} ({self.image})>"
+def symbol_to_api(symbol: CacheSymbol) -> kccore.BNKernelCacheSymbol:
+	return kccore.BNKernelCacheSymbol(
+		symbolType=symbol.symbol_type,
+		address=symbol.address,
+		_name=BNAllocString(symbol.name)
+	)
 
+class KernelCacheController:
+	def __init__(self, view: BinaryView):
+		"""
+		Retrieve the shared cache controller for a given view.
+		Call `is_valid` to check if the controller is valid.
+		"""
+		self.handle = kccore.BNGetKernelCacheController(view.handle)
 
-@dataclasses.dataclass
-class KernelCacheMachOHeader:
-	header: str
+	def __del__(self):
+		if self.handle is not None:
+			kccore.BNFreeKernelCacheControllerReference(self.handle)
 
-	@classmethod
-	def LoadFromString(cls, s: str):
-		return cls(header=s)
+	def __str__(self):
+		return repr(self)
 
 	def __repr__(self):
-		return f"<KernelCacheMachOHeader {self.header!r}>"
+		return f"<KernelCacheController: {len(self.images)} images>"
 
+	def is_valid(self) -> bool:
+		return self.handle is not None
 
-class KernelCache:
-	def __init__(self, view):
-		# Create a KernelCache object from a BinaryView.
-		self.handle = kccore.BNGetKernelCache(view.handle)
+	def apply_image(self, view: BinaryView, image: CacheImage) -> bool:
+		api_image: kccore.BNKernelCacheImage = image_to_api(image)
+		result = kccore.BNKernelCacheControllerApplyImage(self.handle, view.handle, api_image)
+		kccore.BNKernelCacheFreeImage(api_image)
+		return result
 
-	@staticmethod
-	def get_load_progress(view) -> int:
-		"""
-		Returns the current load progress.
-		Note: In the C++ FFI this function takes the BinaryView's file session ID.
-		"""
-		return KCViewLoadProgress(kccore.BNKCViewGetLoadProgress(view.file.session_id))
+	def is_image_loaded(self, image: CacheImage) -> bool:
+		api_image: kccore.BNKernelCacheImage = image_to_api(image)
+		result = kccore.BNKernelCacheControllerIsImageLoaded(self.handle, api_image)
+		kccore.BNKernelCacheFreeImage(api_image)
+		return result
 
-	@staticmethod
-	def fast_get_image_count(view) -> int:
-		"""
-		Quickly returns the number of images in the kernel cache.
-		"""
-		return kccore.BNKCViewFastGetImageCount(view.handle)
+	def get_image_at(self, address: int) -> Optional[CacheImage]:
+		api_image = kccore.BNKernelCacheImage()
+		if not kccore.BNKernelCacheControllerGetImageAt(self.handle, address, api_image):
+			return None
+		image = image_from_api(api_image)
+		kccore.BNKernelCacheFreeImage(api_image)
+		return image
 
-	def load_image_with_install_name(self, install_name: str) -> bool:
-		"""
-		Load a kernel cache image by its install name.
-		"""
-		return kccore.BNKCViewLoadImageWithInstallName(self.handle, install_name)
+	def get_image_containing(self, address: int) -> Optional[CacheImage]:
+		api_image = kccore.BNKernelCacheImage()
+		if not kccore.BNKernelCacheControllerGetImageContaining(self.handle, address, api_image):
+			return None
+		image = image_from_api(api_image)
+		kccore.BNKernelCacheFreeImage(api_image)
+		return image
 
-	def load_image_containing_address(self, addr: int) -> bool:
-		"""
-		Load the image that contains the given address.
-		"""
-		return kccore.BNKCViewLoadImageContainingAddress(self.handle, addr)
+	def get_image_with_name(self, name: str) -> Optional[CacheImage]:
+		api_image = kccore.BNKernelCacheImage()
+		if not kccore.BNKernelCacheControllerGetImageWithName(self.handle, name, api_image):
+			return None
+		image = image_from_api(api_image)
+		kccore.BNKernelCacheFreeImage(api_image)
+		return image
 
-	@property
-	def image_names(self) -> list[str]:
+	def get_image_dependencies(self, image: CacheImage) -> [str]:
 		"""
-		Return a list of available kernel cache image install names.
+		Returns a list of image names that this image depends on.
 		"""
 		count = ctypes.c_ulonglong()
-		value = kccore.BNKCViewGetInstallNames(self.handle, count)
+		api_image: kccore.BNKernelCacheImage = image_to_api(image)
+		value = kccore.BNKernelCacheControllerGetImageDependencies(self.handle, api_image, count)
+		kccore.BNKernelCacheFreeImage(api_image)
 		if value is None:
 			return []
 		result = []
 		for i in range(count.value):
-			result.append(value[i].decode('utf-8'))
+			result.append(value[i].decode("utf-8"))
 		BNFreeStringList(value, count)
 		return result
 
+	def get_symbol_at(self, address: int) -> Optional[CacheSymbol]:
+		api_symbol = kccore.BNKernelCacheSymbol()
+		if not kccore.BNKernelCacheControllerGetSymbolAt(self.handle, address, api_symbol):
+			return None
+		symbol = symbol_from_api(api_symbol)
+		kccore.BNKernelCacheFreeSymbol(api_symbol)
+		return symbol
+
+	def get_symbol_with_name(self, name: str) -> Optional[CacheSymbol]:
+		api_symbol = kccore.BNKernelCacheSymbol()
+		if not kccore.BNKernelCacheControllerGetSymbolWithName(self.handle, name, api_symbol):
+			return None
+		symbol = symbol_from_api(api_symbol)
+		kccore.BNKernelCacheFreeSymbol(api_symbol)
+		return symbol
+
 	@property
-	def images(self) -> list[KCImage]:
-		"""
-		Return all kernel cache images.
-		"""
+	def images(self) -> [CacheImage]:
 		count = ctypes.c_ulonglong()
-		value = kccore.BNKCViewGetAllImages(self.handle, count)
+		value = kccore.BNKernelCacheControllerGetImages(self.handle, count)
 		if value is None:
 			return []
 		result = []
 		for i in range(count.value):
-			mappings = []
-			for j in range(value[i].mappingCount):
-				mapping = KCImageMemoryMapping(
-					name=value[i].mappings[j].name,
-					vmAddress=value[i].mappings[j].vmAddress,
-					rawViewOffset=value[i].mappings[j].rawViewOffset,
-					size=value[i].mappings[j].size,
-				)
-				mappings.append(mapping)
-			result.append(KCImage(
-				name=value[i].name,
-				headerFileAddress=value[i].headerFileAddress,
-				mappings=mappings
-			))
-		kccore.BNKCViewFreeAllImages(value, count)
+			result.append(image_from_api(value[i]))
+		kccore.BNKernelCacheFreeImageList(value, count)
 		return result
 
 	@property
-	def loaded_images(self) -> list[KCImage]:
+	def loaded_images(self) -> [CacheImage]:
 		"""
-		Return the kernel cache images that are currently loaded.
+		Get a list of images that are currently loaded in the view.
 		"""
 		count = ctypes.c_ulonglong()
-		images = kccore.BNKCViewGetLoadedImages(self.handle, count)
-		if images is None:
+		value = kccore.BNKernelCacheControllerGetLoadedImages(self.handle, count)
+		if value is None:
 			return []
 		result = []
 		for i in range(count.value):
-			mappings = []
-			for j in range(images[i].mappingCount):
-				mapping = KCImageMemoryMapping(
-					name=images[i].mappings[j].name,
-					vmAddress=images[i].mappings[j].vmAddress,
-					rawViewOffset=images[i].mappings[j].rawViewOffset,
-					size=images[i].mappings[j].size,
-				)
-				mappings.append(mapping)
-			result.append(KCImage(
-				name=images[i].name,
-				headerFileAddress=images[i].headerFileAddress,
-				mappings=mappings
-			))
-		kccore.BNKCViewFreeAllImages(images, count)
+			result.append(image_from_api(value[i]))
+		kccore.BNKernelCacheFreeImageList(value, count)
 		return result
 
-	def load_all_symbols_and_wait(self) -> list[KCSymbol]:
-		"""
-		Load all symbols from the kernel cache and wait for completion.
-		"""
+	@property
+	def symbols(self) -> [CacheSymbol]:
 		count = ctypes.c_ulonglong()
-		value = kccore.BNKCViewLoadAllSymbolsAndWait(self.handle, count)
+		value = kccore.BNKernelCacheControllerGetSymbols(self.handle, count)
 		if value is None:
 			return []
 		result = []
 		for i in range(count.value):
-			sym = KCSymbol(
-				name=value[i].name,
-				image=value[i].image,
-				address=value[i].address
-			)
-			result.append(sym)
-		kccore.BNKCViewFreeSymbols(value, count)
+			result.append(symbol_from_api(value[i]))
+		kccore.BNKernelCacheFreeSymbolList(value, count)
 		return result
 
-	def get_name_for_address(self, address: int) -> str:
-		"""
-		Return the symbol name for a given address.
-		"""
-		name = kccore.BNKCViewGetNameForAddress(self.handle, address)
-		if name is None:
-			return ""
-		return name
 
-	def get_image_name_for_address(self, address: int) -> str:
-		"""
-		Return the image name for a given address.
-		"""
-		name = kccore.BNKCViewGetImageNameForAddress(self.handle, address)
-		if name is None:
-			return ""
-		return name
+def _get_kernel_cache(instance: binaryninja.PythonScriptingInstance):
+	if instance.interpreter.active_view is None:
+		return None
+	controller = KernelCacheController(instance.interpreter.active_view)
+	if not controller.is_valid():
+		return None
+	return controller
 
-	def get_macho_header_for_image(self, name: str):
-		"""
-		Return a KernelCacheMachOHeader for the image with the given install name.
-		"""
-		outputStr = kccore.BNKCViewGetImageHeaderForName(self.handle, name)
-		if outputStr is None:
-			return None
-		return KernelCacheMachOHeader.LoadFromString(outputStr)
 
-	def get_macho_header_for_address(self, address: int):
-		"""
-		Return a KernelCacheMachOHeader for the image containing the given address.
-		"""
-		outputStr = kccore.BNKCViewGetImageHeaderForAddress(self.handle, address)
-		if outputStr is None:
-			return None
-		return KernelCacheMachOHeader.LoadFromString(outputStr)
+binaryninja.PythonScriptingProvider.register_magic_variable(
+	"kc",
+	_get_kernel_cache
+)
 
-	@property
-	def state(self):
-		"""
-		Return the current state of the kernel cache.
-		"""
-		return KCViewState(kccore.BNKCViewGetState(self.handle))
+binaryninja.PythonScriptingProvider.register_magic_variable(
+	"kernel_cache",
+	_get_kernel_cache
+)
diff --git a/view/kernelcache/api/python/kernelcache_enums.py b/view/kernelcache/api/python/kernelcache_enums.py
index e610c527..83cbf027 100644
--- a/view/kernelcache/api/python/kernelcache_enums.py
+++ b/view/kernelcache/api/python/kernelcache_enums.py
@@ -1,14 +1,38 @@
 import enum
 
 
-class KCViewLoadProgress(enum.IntEnum):
-	LoadProgressNotStarted = 0
-	LoadProgressLoadingCaches = 1
-	LoadProgressLoadingImages = 2
-	LoadProgressFinished = 3
+class KernelCacheEntryType(enum.IntEnum):
+	KernelCacheEntryTypePrimary = 0
+	KernelCacheEntryTypeSecondary = 1
+	KernelCacheEntryTypeSymbols = 2
+	KernelCacheEntryTypeDyldData = 3
+	KernelCacheEntryTypeStub = 4
 
 
-class KCViewState(enum.IntEnum):
-	Unloaded = 0
-	Loaded = 1
-	LoadedWithImages = 2
+class KernelCacheRegionType(enum.IntEnum):
+	KernelCacheRegionTypeImage = 0
+	KernelCacheRegionTypeStubIsland = 1
+	KernelCacheRegionTypeDyldData = 2
+	KernelCacheRegionTypeNonImage = 3
+
+
+class SegmentFlag(enum.IntEnum):
+	SegmentExecutable = 1
+	SegmentWritable = 2
+	SegmentReadable = 4
+	SegmentContainsData = 8
+	SegmentContainsCode = 16
+	SegmentDenyWrite = 32
+	SegmentDenyExecute = 64
+
+
+class SymbolType(enum.IntEnum):
+	FunctionSymbol = 0
+	ImportAddressSymbol = 1
+	ImportedFunctionSymbol = 2
+	DataSymbol = 3
+	ImportedDataSymbol = 4
+	ExternalSymbol = 5
+	LibraryFunctionSymbol = 6
+	SymbolicFunctionSymbol = 7
+	LocalLabelSymbol = 8