"""
.. module:: CGammaAPIEvasionProblem
	:synopsis: Class performs black-box evasion attacks, applying Gamma on APIs.

.. moduleauthor:: Luca Demetrio <luca.demetrio@dibris.unige.it>

"""

import copy

import lief
from secml.array import CArray

from secml_malware.attack.blackbox.c_blackbox_problem import CBlackBoxProblem
from secml_malware.attack.blackbox.c_wrapper_phi import CWrapperPhi


class CGammaAPIEvasionProblem(CBlackBoxProblem):
	"""
	Manipulate the sample by adding APIs, using the Gamma approach.

	Parameters
	----------
	api_list : list
		A list of couples [dll_name, function_name] of each API that can be included during the optimization
	model_wrapper : CWrapperPhi
		The model under attack
	population_size : int
		How many samples will be generated by the genetic strategy at each round
	penalty_regularizer : float
		The value of the regularization parameter
	iterations: int
		Specifies how many generations will be created by the genetic algorithm.
		The overall number of queries will be (iterations + 1) * population_size
	seed : int
		Specify a seed for init the random engine
	is_debug : bool
		Prints additional information during the optimization

	Returns
	----------
	An evasion object ready to run
	"""
	def __init__(self,
				 api_list: list,
				 model_wrapper: CWrapperPhi,
				 population_size: int,
				 penalty_regularizer: float,
				 iterations: int,
				 seed: int = None,
				 is_debug: bool = False):
		super(CGammaAPIEvasionProblem, self).__init__(model_wrapper,
												   len(api_list),
												   population_size,
												   penalty_regularizer,
												   iterations,
												   seed, is_debug)

		self.api_list = api_list
		self._cache_lib = None
		self._names = None

	def compute_penalty_term(self, original_x: CArray, adv_x: CArray, par: float) -> float:
		return par * abs(adv_x.shape[-1] - original_x.shape[-1])

	def get_nix(self):
		return self.latent_space_size

	def _add_import(self, liefpe, dll, func_name):
		lib = [l for l in liefpe.imports if l.name.lower() == dll.lower()]
		if lib == []:
			lib = liefpe.add_library(dll)
		lib = lief.add_library(dll) if lib == [] else lib[0]
		names = set([e.name for e in lib.entries])
		if not func_name in names:
			lib.add_entry(func_name)
		# if self._cache_lib is None:
		# 	lib = [l for l in lief.imports if l.name.lower() == dll.lower()]
		# 	if lib == []:
		# 		lib = lief.add_library(dll)
		# 	lib = lief.add_library(dll) if lib == [] else lib[0]
		# 	self._cache_lib = lib
		# names = set([e.name for e in self._cache_lib.entries])
		# if not func_name in names:
		# 	self._cache_lib.add_entry(func_name)

	def apply_feasible_manipulations(self, t, x: CArray) -> CArray:
		x_adv = copy.deepcopy(x)
		x_adv = x_adv.tolist()
		lief_adv : lief.PE.Binary = lief.PE.parse(x_adv[0])
		for i in range(t.shape[-1]):
			dll_name, function_name = self.api_list[i]
			self._add_import(lief_adv, dll_name, function_name)
		builder = lief.PE.Builder(lief_adv)
		builder.build_imports(True).patch_imports(True)
		builder.build()
		x_adv = CArray(builder.get_build())
		x_adv = x_adv.reshape((1, x_adv.shape[-1]))
		return x_adv
