# Copyright (c) 2015-2025 Vector 35 Inc
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.

from typing import Optional, Union
from dataclasses import dataclass
import ctypes

import binaryninja
from . import _binaryninjacore as core
from .log import log_error_for_exception
from . import types
from . import highlevelil
from . import binaryview


class _CustomStringTypeMetaClass(type):
	def __iter__(self):
		binaryninja._init_plugins()
		count = ctypes.c_ulonglong()
		types = core.BNGetCustomStringTypeList(count)
		assert types is not None, "core.BNGetCustomStringTypeList returned None"
		try:
			for i in range(0, count.value):
				yield CustomStringType(handle=types[i])
		finally:
			core.BNFreeCustomStringTypeList(types)

	def __getitem__(cls, value):
		binaryninja._init_plugins()
		string_type = core.BNGetCustomStringTypeByName(str(value))
		if string_type is None:
			raise KeyError("'%s' is not a valid type" % str(value))
		return CustomStringType(handle=string_type)


class CustomStringType(metaclass=_CustomStringTypeMetaClass):
	def __init__(self, handle):
		self.handle = core.handle_of_type(handle, core.BNCustomStringType)

	def __str__(self):
		return self.name

	def __repr__(self):
		return f"<{self.__class__.__name__}: {self.name}>"

	def __eq__(self, other):
		if not isinstance(other, self.__class__):
			return NotImplemented
		return ctypes.addressof(self.handle.contents) == ctypes.addressof(other.handle.contents)

	def __ne__(self, other):
		if not isinstance(other, self.__class__):
			return NotImplemented
		return not (self == other)

	def __hash__(self):
		return hash(ctypes.addressof(self.handle.contents))

	@staticmethod
	def register(name: str, string_prefix="", string_postfix="") -> 'CustomStringType':
		info = core.BNCustomStringTypeInfo()
		info.name = name
		info.stringPrefix = string_prefix
		info.stringPostfix = string_postfix
		handle = core.BNRegisterCustomStringType(info)
		return CustomStringType(handle)

	@property
	def name(self) -> str:
		return core.BNGetCustomStringTypeName(self.handle)

	@property
	def string_prefix(self) -> str:
		return core.BNGetCustomStringTypePrefix(self.handle)

	@property
	def string_postfix(self) -> str:
		return core.BNGetCustomStringTypePostfix(self.handle)


class _StringRecognizerMetaClass(type):
	def __iter__(self):
		binaryninja._init_plugins()
		count = ctypes.c_ulonglong()
		recognizers = core.BNGetStringRecognizerList(count)
		assert recognizers is not None, "core.BNGetStringRecognizerList returned None"
		try:
			for i in range(0, count.value):
				yield CoreStringRecognizer(handle=recognizers[i])
		finally:
			core.BNFreeStringRecognizerList(recognizers)

	def __getitem__(cls, value):
		binaryninja._init_plugins()
		recognizer = core.BNGetStringRecognizerByName(str(value))
		if recognizer is None:
			raise KeyError("'%s' is not a valid recognizer" % str(value))
		return CoreStringRecognizer(handle=recognizer)


class StringRecognizer(metaclass=_StringRecognizerMetaClass):
	_registered_recognizers = []
	recognizer_name = None

	def __init__(self, handle=None):
		if handle is not None:
			self.handle = core.handle_of_type(handle, core.BNStringRecognizer)

	def register(self):
		if self.__class__.recognizer_name is None:
			raise ValueError("Recognizer name is missing")
		self._cb = core.BNCustomStringRecognizer()
		self._cb.context = 0
		if self.is_valid_for_type.__func__ != StringRecognizer.is_valid_for_type:
			self._cb.isValidForType = self._cb.isValidForType.__class__(self._is_valid_for_type)
		if self.recognize_constant.__func__ != StringRecognizer.recognize_constant:
			self._cb.recognizeConstant = self._cb.recognizeConstant.__class__(self._recognize_constant)
		if self.recognize_constant_pointer.__func__ != StringRecognizer.recognize_constant_pointer:
			self._cb.recognizeConstantPointer = self._cb.recognizeConstantPointer.__class__(
				self._recognize_constant_pointer)
		if self.recognize_extern_pointer.__func__ != StringRecognizer.recognize_extern_pointer:
			self._cb.recognizeExternPointer = self._cb.recognizeExternPointer.__class__(self._recognize_extern_pointer)
		if self.recognize_import.__func__ != StringRecognizer.recognize_import:
			self._cb.recognizeImport = self._cb.recognizeImport.__class__(self._recognize_import)
		self.handle = core.BNRegisterStringRecognizer(self.__class__.recognizer_name, self._cb)
		self.__class__._registered_recognizers.append(self)

	def _is_valid_for_type(self, ctxt, hlil, type):
		try:
			hlil = highlevelil.HighLevelILFunction(handle=core.BNNewHighLevelILFunctionReference(hlil))
			type = types.Type.create(handle=core.BNNewTypeReference(type))
			return self.is_valid_for_type(hlil, type)
		except:
			log_error_for_exception("Unhandled Python exception in StringRecognizer._is_valid_for_type")
			return False

	def _recognize_constant(self, ctxt, hlil, expr, type, val, result):
		try:
			hlil = highlevelil.HighLevelILFunction(handle=core.BNNewHighLevelILFunctionReference(hlil))
			type = types.Type.create(handle=core.BNNewTypeReference(type))
			instr = hlil.get_expr(highlevelil.ExpressionIndex(expr))
			ref = self.recognize_constant(instr, type, val)
			if ref is None:
				return False
			result[0] = ref._to_core_struct(True)
			return True
		except:
			log_error_for_exception("Unhandled Python exception in StringRecognizer._recognize_constant")
			return False

	def _recognize_constant_pointer(self, ctxt, hlil, expr, type, val, result):
		try:
			hlil = highlevelil.HighLevelILFunction(handle=core.BNNewHighLevelILFunctionReference(hlil))
			type = types.Type.create(handle=core.BNNewTypeReference(type))
			instr = hlil.get_expr(highlevelil.ExpressionIndex(expr))
			ref = self.recognize_constant_pointer(instr, type, val)
			if ref is None:
				return False
			result[0] = ref._to_core_struct(True)
			return True
		except:
			log_error_for_exception("Unhandled Python exception in StringRecognizer._recognize_constant_pointer")
			return False

	def _recognize_extern_pointer(self, ctxt, hlil, expr, type, val, offset, result):
		try:
			hlil = highlevelil.HighLevelILFunction(handle=core.BNNewHighLevelILFunctionReference(hlil))
			type = types.Type.create(handle=core.BNNewTypeReference(type))
			instr = hlil.get_expr(highlevelil.ExpressionIndex(expr))
			ref = self.recognize_extern_pointer(instr, type, val, offset)
			if ref is None:
				return False
			result[0] = ref._to_core_struct(True)
			return True
		except:
			log_error_for_exception("Unhandled Python exception in StringRecognizer._recognize_extern_pointer")
			return False

	def _recognize_import(self, ctxt, hlil, expr, type, val, result):
		try:
			hlil = highlevelil.HighLevelILFunction(handle=core.BNNewHighLevelILFunctionReference(hlil))
			type = types.Type.create(handle=core.BNNewTypeReference(type))
			instr = hlil.get_expr(highlevelil.ExpressionIndex(expr))
			ref = self.recognize_import(instr, type, val)
			if ref is None:
				return False
			result[0] = ref._to_core_struct(True)
			return True
		except:
			log_error_for_exception("Unhandled Python exception in StringRecognizer._recognize_import")
			return False

	@property
	def name(self) -> str:
		if hasattr(self, 'handle'):
			return core.BNGetStringRecognizerName(self.handle)
		return self.__class__.recognizer_name

	def is_valid_for_type(self, func: 'highlevelil.HighLevelILFunction', type: 'types.Type') -> bool:
		return True

	def recognize_constant(
		self, instr: 'highlevelil.HighLevelILInstruction', type: 'types.Type', val: int
	) -> Optional['binaryview.DerivedString']:
		return None

	def recognize_constant_pointer(
		self, instr: 'highlevelil.HighLevelILInstruction', type: 'types.Type', val: int
	) -> Optional['binaryview.DerivedString']:
		return None

	def recognize_extern_pointer(
		self, instr: 'highlevelil.HighLevelILInstruction', type: 'types.Type', val: int, offset: int
	) -> Optional['binaryview.DerivedString']:
		return None

	def recognize_import(
		self, instr: 'highlevelil.HighLevelILInstruction', type: 'types.Type', val: int
	) -> Optional['binaryview.DerivedString']:
		return None


_recognizer_cache = {}


class CoreStringRecognizer(StringRecognizer):
	def __init__(self, handle: core.BNStringRecognizer):
		super(CoreStringRecognizer, self).__init__(handle=handle)
		if type(self) is CoreStringRecognizer:
			global _recognizer_cache
			_recognizer_cache[ctypes.addressof(handle.contents)] = self

	@classmethod
	def _from_cache(cls, handle) -> 'StringRecognizer':
		"""
		Look up a recognizer from a given BNStringRecognizer handle
		:param handle: BNStringRecognizer pointer
		:return: Recognizer instance responsible for this handle
		"""
		global _recognizer_cache
		return _recognizer_cache.get(ctypes.addressof(handle.contents)) or cls(handle)

	def is_valid_for_type(self, func: 'highlevelil.HighLevelILFunction', type: 'types.Type') -> bool:
		return core.BNIsStringRecognizerValidForType(self.handle, func.handle, type.handle)

	def recognize_constant(
		self, instr: 'highlevelil.HighLevelILInstruction', type: 'types.Type', val: int
	) -> Optional['binaryview.DerivedString']:
		string = core.BNDerivedString()
		if not core.BNStringRecognizerRecognizeConstant(self.handle, instr.function.handle, instr.expr_index, type.handle, val, string):
			return None
		return binaryview.DerivedString._from_core_struct(string, True)

	def recognize_constant_pointer(
		self, instr: 'highlevelil.HighLevelILInstruction', type: 'types.Type', val: int
	) -> Optional['binaryview.DerivedString']:
		string = core.BNDerivedString()
		if not core.BNStringRecognizerRecognizeConstantPointer(self.handle, instr.function.handle, instr.expr_index, type.handle, val, string):
			return None
		return binaryview.DerivedString._from_core_struct(string, True)

	def recognize_extern_pointer(
		self, instr: 'highlevelil.HighLevelILInstruction', type: 'types.Type', val: int, offset: int
	) -> Optional['binaryview.DerivedString']:
		string = core.BNDerivedString()
		if not core.BNStringRecognizerRecognizeExternPointer(self.handle, instr.function.handle, instr.expr_index, type.handle, val, offset, string):
			return None
		return binaryview.DerivedString._from_core_struct(string, True)

	def recognize_import(
		self, instr: 'highlevelil.HighLevelILInstruction', type: 'types.Type', val: int
	) -> Optional['binaryview.DerivedString']:
		string = core.BNDerivedString()
		if not core.BNStringRecognizerRecognizeImport(self.handle, instr.function.handle, instr.expr_index, type.handle, val, string):
			return None
		return binaryview.DerivedString._from_core_struct(string, True)