import struct
import numpy as np
import torch
from secml.settings import SECML_PYTORCH_USE_CUDA

from secml_malware.attack.whitebox.c_discretized_bytes_evasion import CDiscreteBytesEvasion

use_cuda = torch.cuda.is_available() and SECML_PYTORCH_USE_CUDA


class CHeaderEvasion(CDiscreteBytesEvasion):
	"""Creates the attack that perturbs the header of a Windows PE malware.
	More details in the reference paper.
	"Explaining Vulnerabilities of Deep Learning to Adversarial Malware Binaries"
	https://arxiv.org/abs/1901.03583
	Demetrio, Biggio, Lagorio, Roli, Armando
	
	Arguments:
		CMalConvEvasion {[type]} -- [description]
	"""

	def __init__(
			self,
			end2end_model,
			index_to_perturb=None,
			iterations=100,
			is_debug=False,
			random_init=False,
			optimize_all_dos=False,
			threshold=0.5,
			penalty_regularizer=0
	):
		"""Creates the evasion object

		Arguments:
			end2end_model {CClassifierMalConvPyTorch} -- The classifier to evade

		Keyword Arguments:
			index_to_perturb {list of int} -- Indexes that will be used for the attack (default: [2, ..., 58])
			how_many {int} -- How many index to perturb. None means use all indexes (default: {None})
			surrogate_classifier {CClassifierMalConvPyTorch} -- Surrogate classifier, if any (default: {None})
			use_surrogate {bool} -- Specify if should use the surrogate instead of real classifier (default: {False})
			iterations {float} -- How many iterations for the attack (default: {100.0})
			is_debug {bool} -- If true, it will print on console additional information duringthe attack (default: {False})
			random_init {bool} -- Randomize the bytes located at the specified indexes before starting the attack(default: {False})
			conclude_iterations {bool} -- should conclude iterations or stop below the given threshold
		"""

		if index_to_perturb is None:
			index_to_perturb = [i for i in range(2, 0x3C)]
		super(CHeaderEvasion, self).__init__(
			end2end_model,
			index_to_perturb,
			iterations,
			is_debug,
			random_init,
			threshold,
			penalty_regularizer
		)
		self.optimize_all_dos = optimize_all_dos

	def _set_dos_indexes(self, x_init):
		if self.optimize_all_dos:
			pe_position = x_init[0x3C:0x40].tondarray().astype(np.uint16)[0]
			if self.shift_values:
				pe_position = np.array([p - 1 for p in pe_position])
			pe_position = struct.unpack("<I", bytes(pe_position.astype(np.uint8)))[0]
			self.indexes_to_perturb = [i for i in range(2, 0x3C)] + [
				i for i in range(0x40, pe_position)
			]
			self._how_many = len(self.indexes_to_perturb)
			if self.is_debug:
				print(f"PE POSITION: {pe_position}, perturbing {self._how_many}")

	def _run(self, x0, y0, x_init=None):
		self._set_dos_indexes(x_init)
		return super(CHeaderEvasion, self)._run(x0, y0, x_init)
