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


int elf_magic_numbers(char *str)
{
	return (!ft_strncmp(str, ELFMAG, SELFMAG));
}

int		save_elf(char *path, char *file, unsigned long int size)
{
	int fd = open(path, O_CREAT | O_WRONLY | O_TRUNC, 0755);
    if (fd == -1) {
		ft_printf("Error: Failed to create new file \'%s\'\n", path);
		return EXIT_FAILURE;
    }

    if (write(fd, file, size) == -1) {
        close(fd);
        ft_printf("Error: Failed to write new file \'%s\'\n", path);
        return EXIT_FAILURE;
    }

    if (close(fd) == -1) {
		ft_printf("Error: Failed to close new file \'%s\'\n", path);
        return EXIT_FAILURE;
    }
	return EXIT_SUCCESS;
}

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

void offset_sections(t_efl_content *woody, unsigned int from, unsigned int offset_ammount)
{
	for (int i = 0; i < woody->Ehdr->e_phnum; i++)
	{
		if (woody->Phdr[i].p_offset > from)
			woody->Phdr[i].p_offset += offset_ammount;
	}
	for (int i = 0; i < woody->Ehdr->e_shnum; i++)
	{
		if (woody->Shdr[i].sh_offset > from)
			woody->Shdr[i].sh_offset += offset_ammount;
	}
}

size_t	create_codecave(t_efl_content *woody, Elf64_Phdr *load_segment)
{
	const unsigned int page_size = 4096;// getpagesize(); not authorized
	unsigned int padding_size = ((sizeof(PAYLOAD) / page_size) + 1) * page_size;
	unsigned int codecave_start = load_segment->p_offset + load_segment->p_filesz;
	offset_sections(woody, codecave_start, padding_size);
	char *new_woody = malloc(woody->file_size + padding_size);
	if (!new_woody)
		return 0;
	ft_memcpy(new_woody, woody->file, codecave_start);
	ft_bzero(new_woody + codecave_start, padding_size);
	ft_memcpy(new_woody + codecave_start + padding_size,  woody->file + codecave_start, woody->file_size - codecave_start);
	munmap(woody->file, woody->file_size);
	woody->file = new_woody;
	woody->file_size += padding_size;
	woody->Ehdr = (Elf64_Ehdr *)new_woody;
	woody->Phdr = (Elf64_Phdr *)fetch(woody->file, woody->file_size, woody->Ehdr->e_phoff, sizeof(Elf64_Phdr));
	woody->Shdr = (Elf64_Shdr *)fetch(woody->file, woody->file_size, woody->Ehdr->e_shoff, sizeof(Elf64_Shdr));
	return codecave_start;
}


int	insert_payload(t_efl_content *woody, size_t payload_position)
{
	char payload[] = PAYLOAD;
    size_t len_payload = sizeof(PAYLOAD) - 1;
	char *ptr = ft_strnstr_nullterminated(payload, JUMP, len_payload);

	if (ptr)
	{
		int32_t jmp_index = ptr - payload; 
		int32_t jump_value = ((payload_position + len_payload) - woody->Ehdr->e_entry) * -1;
		ft_memcpy(&payload[jmp_index + 1], &jump_value, sizeof(jump_value));
		ft_memcpy(woody->file + payload_position, payload, len_payload);

		printf("Old entry : %ld (%lx)\n", woody->Ehdr->e_entry, woody->Ehdr->e_entry);
		printf("Code_cave_start = %ld (%lx)\n", payload_position, payload_position);
		printf("Payload size = %ld (%lx)\n", len_payload, len_payload);
		printf("Backward offset = %d (%x)\n", jump_value, jump_value);
		
		return EXIT_SUCCESS;
	}
	return EXIT_FAILURE;
}

void	inject(t_efl_content *woody)
{
    size_t len_payload = sizeof(PAYLOAD) - 1;

	int i = get_load_segment(woody, 0, true);	
	int j = get_load_segment(woody, i + 1, false);
	
	
	size_t code_cave_size = woody->Phdr[j].p_offset - (woody->Phdr[i].p_offset + woody->Phdr[i].p_filesz);
	size_t payload_position;

	if (code_cave_size > len_payload)
	{
		payload_position = woody->Phdr[i].p_offset + woody->Phdr[i].p_memsz;
		printf("Code_cave_size = %ld (%lx)\n", code_cave_size, code_cave_size);
	}
	else
	{
		payload_position = create_codecave(woody, &woody->Phdr[i]);
	}
	insert_payload(woody, payload_position);

	woody->Ehdr->e_entry = payload_position;
	woody->Phdr[i].p_filesz += len_payload;
	woody->Phdr[i].p_memsz += len_payload;

	printf("New entry = %ld (%lx)\n", woody->Ehdr->e_entry, woody->Ehdr->e_entry);
}

int get_elf_sections(t_efl_content *woody)
{
	woody->Ehdr = (Elf64_Ehdr *)fetch(woody->file, woody->file_size, 0, sizeof(Elf64_Ehdr));
	if (!woody->Ehdr || !elf_magic_numbers(woody->file) || woody->Ehdr->e_ident[EI_CLASS] != ELFCLASS64)
	{
		ft_printf("Error: \'%s\' is not a valid 64-bit ELF file\n", woody->file_path);
		return EXIT_FAILURE;
	}
	woody->Phdr = (Elf64_Phdr *)fetch(woody->file, woody->file_size, woody->Ehdr->e_phoff, sizeof(Elf64_Phdr));

	woody->Shdr = (Elf64_Shdr *)fetch(woody->file, woody->file_size, woody->Ehdr->e_shoff, sizeof(Elf64_Shdr));
	if (woody->Shdr == NULL || !fetch(woody->file, woody->file_size, woody->Ehdr->e_shoff, woody->Ehdr->e_shnum * sizeof(Elf64_Shdr)))
	{
		return ft_put_error("Corrupted file");
	}



	// Elf64_Shdr *symbols_table = NULL;
    // for (int i = 0; i < woody->Ehdr->e_shnum; i++) {
	// 	if (Shdr[i].sh_type == SHT_SYMTAB) {
    //         symbols_table = fetch(woody->file, woody->file_size, woody->Ehdr->e_shoff + (i * sizeof(Elf64_Shdr)), sizeof(Elf64_Shdr));
    //     }
    // }
	// if (symbols_table == NULL)
	// 	return ft_put_error("No symbols");

	// if (!fetch(woody->file, woody->file_size, woody->Ehdr->e_shoff + (woody->Ehdr->e_shstrndx * sizeof(Elf64_Shdr)), sizeof(Elf64_Shdr)))
	// 	return ft_put_error("Corrupted file");

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

    // Elf64_Shdr *strtab_header = (Elf64_Shdr *)fetch(woody->file, woody->file_size, woody->Ehdr->e_shoff + (symbols_table->sh_link * woody->Ehdr->e_shentsize), sizeof(Elf64_Shdr)); 
	// if (!strtab_header)
	// 	return ft_put_error("Corrupted file");

    // char *strtab = fetch(woody->file, woody->file_size, strtab_header->sh_offset, 0);
	// if (strtab == NULL)
	// 	return ft_put_error("Corrupted file");
    // Elf64_Sym *symbols = (Elf64_Sym *)fetch(woody->file, woody->file_size, symbols_table->sh_offset, sizeof(Elf64_Sym));
	// if (symbols == NULL)
	// 	return ft_put_error("Corrupted file");

	return EXIT_SUCCESS;
}

int prepare_injection(t_efl_content *woody)
{
	int elf_statut = get_elf_sections(woody);
	if (elf_statut)
		return elf_statut;

	inject(woody);

	char *woody_file;
	if (!(woody_file = malloc(woody->file_size)))
		return ft_put_error("Allocation error");

	ft_memcpy(woody_file, woody->file, woody->file_size);

	// encrypt(woody_file, strtab_header->sh_offset , strtab_header->sh_size);
	munmap(woody_file, woody->file_size);
	save_elf("woody", woody_file, woody->file_size);
	free(woody_file);
	return EXIT_SUCCESS;
}