#include "../includes/woody.h"

int	get_load_segment64(t_elf_content *woody, int start, bool executable)
{
	t_elf64 *elf = woody->elf64;
	for (int i = start; i < elf->Ehdr->e_phnum; i++)
	{
		if (elf->Phdr[i].p_type == PT_LOAD)
		{
			if (executable)
			{
				if (elf->Phdr[i].p_flags & PF_X)
					return i;
			}
			else
				return i;
		}
	}
	return -1;
}

int	inject64(t_elf_content *woody)
{
	t_elf64 *elf = woody->elf64;
	t_payload *payload = get_payload();
	if (!payload)
		return EXIT_FAILURE;

	int i = get_load_segment64(woody, 0, true);
	int j = get_load_segment64(woody, i + 1, false);

	if (i == -1 || j != i + 1)
	{
		free(payload->payload);
		free(payload);
		return ft_put_error("PT_LOAD segment missing");
	}

	size_t code_cave_size = elf->Phdr[j].p_offset - (elf->Phdr[i].p_offset + elf->Phdr[i].p_filesz);
	size_t payload_position = elf->Phdr[i].p_offset + elf->Phdr[i].p_filesz;

	if (code_cave_size < (size_t)payload->len)
	{
		free(payload->payload);
		free(payload);
		return ft_put_error("Unable to insert payload, not enough space for code cave");
	}
	encrypt(woody->file, elf->Phdr[i].p_offset, elf->Phdr[i].p_memsz);

	if (insert_payload(woody, payload, payload_position, elf->Ehdr->e_entry, elf->Phdr[i].p_offset, elf->Phdr[i].p_memsz))
	{
		free(payload->payload);
		free(payload);
		return ft_put_error("Unable to insert payload, please regenerate it");
	}

	elf->Ehdr->e_entry = payload_position;
	elf->Phdr[i].p_filesz += payload->len;
	elf->Phdr[i].p_memsz += payload->len;
	elf->Phdr[i].p_flags = PF_X | PF_W | PF_R;
	free(payload->payload);
	free(payload);
	return EXIT_SUCCESS;
}

int get_elf_sections64(t_elf_content *woody)
{
	t_elf64 *elf = woody->elf64;
	
	elf->Ehdr = (Elf64_Ehdr *)fetch(woody->file, woody->file_size, 0, sizeof(Elf64_Ehdr));
	if (!elf->Ehdr)
		return EXIT_FAILURE;

	elf->Phdr = (Elf64_Phdr *)fetch(woody->file, woody->file_size, elf->Ehdr->e_phoff, sizeof(Elf64_Phdr));
	if (!elf->Phdr)
		return EXIT_FAILURE;

	elf->Shdr = (Elf64_Shdr *)fetch(woody->file, woody->file_size, elf->Ehdr->e_shoff, sizeof(Elf64_Shdr));
	if (!elf->Shdr || !fetch(woody->file, woody->file_size, elf->Ehdr->e_shoff, elf->Ehdr->e_shnum * sizeof(Elf64_Shdr)))
		return EXIT_FAILURE;

	if (!fetch(woody->file, woody->file_size, elf->Ehdr->e_shoff + (elf->Ehdr->e_shstrndx * sizeof(Elf64_Shdr)), sizeof(Elf64_Shdr)))
		return EXIT_FAILURE;

	char *Sshstrtab = (char *)fetch(woody->file, woody->file_size, elf->Shdr[elf->Ehdr->e_shstrndx].sh_offset, 0);
	if (Sshstrtab == NULL)
		return EXIT_FAILURE;

	for (int i = 0; i < elf->Ehdr->e_shnum;i++)
	{
    	if (elf->Shdr[i].sh_type == SHT_PROGBITS && elf->Shdr[i].sh_flags & SHF_EXECINSTR && elf->Shdr[i].sh_flags & SHF_ALLOC && elf->Shdr[i].sh_flags & SHF_EXECINSTR)
		{
			if (Sshstrtab + elf->Shdr[i].sh_name < (char *)woody->file + woody->file_size && !ft_strncmp(".text\0", Sshstrtab + elf->Shdr[i].sh_name, 6))
			{
				elf->text_section = &elf->Shdr[i];
				return EXIT_SUCCESS;
			}
		}
	}
	return EXIT_FAILURE;
}