import math
import struct
import lief
import numpy as np
from secml.array import CArray
from secml_malware.models import CClassifierEnd2EndMalware


def shift_pointer_to_section_content(liefpe, raw_code, entry_index, amount, pe_shifted_by = 0):
	pe_position = liefpe.dos_header.addressof_new_exeheader + pe_shifted_by
	optional_header_size = liefpe.header.sizeof_optional_header
	coff_header_size = 24
	section_entry_length = 40
	size_of_raw_data_pointer = 20
	shift_position = (
		pe_position
		+ coff_header_size
		+ optional_header_size
		+ (entry_index * section_entry_length)
		+ size_of_raw_data_pointer
	)
	old_value = struct.unpack("<I", raw_code[shift_position : shift_position + 4])[0]
	new_value = old_value + amount
	new_value = struct.pack("<I", new_value)
	raw_code[shift_position : shift_position + 4] = new_value

	return raw_code


def shift_pe_header(liefpe, raw_code, amount):
	pe_position = liefpe.dos_header.addressof_new_exeheader
	raw_code[0x3C:0x40] = struct.pack("<I", pe_position + amount)

	[raw_code.insert(pe_position, 0) for _ in range(amount)]

	return raw_code


def apply_shift(file_name, new_file_name=None, amount=0x200):
	file_path = file_name
	with open(file_path, "rb") as f:
		code = bytearray(f.read())
	return apply_shift_to_raw_code(amount, code, new_file_name)


def shift_pe_header_by(x : list, preferable_extension_amount : int):
	if not preferable_extension_amount:
		return x, []
	liefpe = lief.PE.parse(x)
	section_file_alignment = liefpe.optional_header.file_alignment
	if section_file_alignment == 0:
		return x, []
	first_content_offset = liefpe.dos_header.addressof_new_exeheader
	extension_amount = int(math.ceil(preferable_extension_amount / section_file_alignment)) * section_file_alignment
	index_to_perturb = list(range(2, 0x3c)) + list(range(0x40, first_content_offset + extension_amount))
	x = shift_pe_header(liefpe, x, extension_amount)
	for i, _ in enumerate(liefpe.sections):
		x = shift_pointer_to_section_content(liefpe, bytearray(x), i, extension_amount, extension_amount)
	return x, index_to_perturb


def shift_section_by(x : list, preferable_extension_amount : int, pe_shifted_by : int = 0):
	if not preferable_extension_amount:
		return x, []
	liefpe = lief.PE.parse(x)
	section_file_alignment = liefpe.optional_header.file_alignment
	if section_file_alignment == 0:
		return x, []
	first_content_offset = liefpe.sections[0].offset
	extension_amount = int(math.ceil(preferable_extension_amount / section_file_alignment)) * section_file_alignment
	index_to_perturb = list(range(first_content_offset, first_content_offset + extension_amount))
	for i, _ in enumerate(liefpe.sections):
		x = shift_pointer_to_section_content(liefpe, x, i, extension_amount, pe_shifted_by)
	x = x[:first_content_offset] + b'\x00'*extension_amount + x[first_content_offset:]
	return x, index_to_perturb

def shift_section_by_using_lief(x : list, liefpe : lief.PE.Binary, preferable_extension_amount : int, pe_shifted_by : int = 0):
	if not preferable_extension_amount:
		return x, []
	section_file_alignment = liefpe.optional_header.file_alignment
	first_content_offset = liefpe.sections[0].offset
	extension_amount = int(math.ceil(preferable_extension_amount / section_file_alignment)) * section_file_alignment
	index_to_perturb = list(range(first_content_offset, first_content_offset + extension_amount))

	# shift offset of each section entry by specified amount
	for i, _ in enumerate(liefpe.sections):
		x = shift_pointer_to_section_content(liefpe, x, i, extension_amount, pe_shifted_by)

	# code_to_replicate = x[first_content_offset : first_content_offset + extension_amount]
	# # shift content by amount, by adding the old content
	# [x.insert(c, first_content_offset) for c in range(code_to_replicate[::-1])]

	# shift content by amount
	# [x.insert(first_content_offset, 0) for _ in range(extension_amount)]
	x = x[:first_content_offset] + b'\x00'*extension_amount + x[first_content_offset:] #[x.insert(first_content_offset, 0) for _ in range(extension_amount)]
	return x, index_to_perturb

def apply_shift_to_raw_code(amount, code, new_file_name):
	parse_pe = lief.PE.parse(list(code))
	amount = (
		parse_pe.optional_header.file_alignment if amount is None else amount
	)
	for i, s in enumerate(parse_pe.sections):
		print(f"Shifting {s.name}")
		code = shift_pointer_to_section_content(parse_pe, code, i, amount)
	code = shift_pe_header(parse_pe, code, amount)
	if new_file_name is not None:
		with open(new_file_name, "wb") as f:
			f.write(code)
		print(f"Written {new_file_name}")
	return code


def create_int_list_from_x_adv(x_adv: CArray, embedding_value, is_shifting_values):
	invalid_value = 256 if embedding_value == -1 else embedding_value
	padding_positions = x_adv.find(x_adv == invalid_value)
	if padding_positions:
		x_adv = x_adv[:padding_positions[0]]
	if is_shifting_values:
		x_adv = x_adv - 1
	x_adv_edit = x_adv[0, :].astype(np.uint8).tolist()[0]
	return bytearray(x_adv_edit)