from secml_malware.attack.whitebox.c_discretized_bytes_evasion import CDiscreteBytesEvasion


class CPaddingEvasion(CDiscreteBytesEvasion):
	"""
	Constructs an attack object that append one byte at time.
	"""
	def __init__(
		self,
		end2end_model,
		how_many,
		iterations=100,
		is_debug=False,
		random_init=False,
		threshold=0.5,
		penalty_regularizer=0
	):
		"""Creates the padding evasion object

		Arguments:
			malconv {CClassifierMalConvPyTorch} -- The classifier to evade
			surrogate_data {CDataset} -- data used for the surrogate (if any)

		Keyword Arguments:
			how_many {int} -- How many index to perturb. None means use all indexes (default: {None})
			index_to_perturb {list of int} -- Indexes that will be used for the attack (default: [2, ..., 58])
			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})
		"""
		super(CPaddingEvasion, self).__init__(
			end2end_model,
			index_to_perturb=[],
			iterations=iterations,
			is_debug=is_debug,
			random_init=random_init,
			threshold=threshold,
			penalty_regularizer=penalty_regularizer
		)
		self.how_many_padding_bytes = how_many

	def _run(self, x0, y0, x_init=None):
		invalid_value = 256 if self.invalid_pos == -1 else self.invalid_pos
		padding_positions = x0.find(x0 == invalid_value)
		if not padding_positions:
			self.indexes_to_perturb = []
		else:
			self.indexes_to_perturb = list(
				range(
					padding_positions[0],
					min(x0.size, padding_positions[0] + self.how_many_padding_bytes),
				)
			)
		return super(CPaddingEvasion, self)._run(x0, y0, x_init=x_init)
