#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) {
		printf("Error: Failed to create new file \'%s\'\n", path);
		return EXIT_FAILURE;
    }

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

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

int	get_load_segment(t_elf_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_elf_content *woody, unsigned int from, unsigned int offset_ammount)
{
	for (int i = 0; i < woody->Ehdr->e_shnum; i++)
	{
		if (woody->Shdr[i].sh_type == SHT_REL)
		{
			int num_entries = woody->Shdr[i].sh_size / woody->Shdr[i].sh_entsize;
			Elf64_Rel *rel = (Elf64_Rel *)fetch(woody->file, woody->file_size, woody->Shdr[i].sh_offset, woody->Shdr[i].sh_size);
            for (int j = 0; j < num_entries; j++)
			{
				rel[j].r_offset += offset_ammount;
                printf("  Offset: 0x%lx, Info: 0x%lx, Type: %u, Symbol: %u\n",
                       (unsigned long)rel[j].r_offset, (unsigned long)rel[j].r_info,
                       (unsigned int)ELF64_R_TYPE(rel[j].r_info), (unsigned int)ELF64_R_SYM(rel[j].r_info));
            }
		}
		else if (woody->Shdr[i].sh_type == SHT_RELA)
		{
			int num_entries = woody->Shdr[i].sh_size / woody->Shdr[i].sh_entsize;
			Elf64_Rela *rela  = (Elf64_Rela *)fetch(woody->file, woody->file_size, woody->Shdr[i].sh_offset , woody->Shdr[i].sh_size);
			for (int j = 0; j < num_entries; j++)
			{
				rela[j].r_offset += offset_ammount;
				if (rela[j].r_addend >= from)
					rela[j].r_addend += offset_ammount;

				    // printf("  Offset: 0x%lx, Info: 0x%lx, Type: %u, Symbol: %u  adden: 0x%lx\n",
                    //    (unsigned long)rela[j].r_offset, (unsigned long)rela[j].r_info,
                    //    (unsigned int)ELF64_R_TYPE(rela[j].r_info), (unsigned int)ELF64_R_SYM(rela[j].r_info), (unsigned long)rela[j].r_addend);


            }
		} else if (woody->Shdr[i].sh_type == SHT_DYNAMIC)
		{
            int num_dyn_entries = woody->Shdr[i].sh_size / woody->Shdr[i].sh_entsize;
			Elf64_Dyn *dyn_entries  = (Elf64_Dyn *)fetch(woody->file, woody->file_size, woody->Shdr[i].sh_offset , woody->Shdr[i].sh_size);
            for (int j = 0; j < num_dyn_entries; j++)
			{
				// printf("dyn_entries[%d].d_tag = %lx    ptr = %lx    value = %lx\n",j, dyn_entries[j].d_tag, dyn_entries[j].d_un.d_ptr, dyn_entries[j].d_un.d_val);
    			if (dyn_entries[j].d_tag == DT_PLTGOT || dyn_entries[j].d_tag == DT_FINI || dyn_entries[j].d_tag == DT_INIT_ARRAY || dyn_entries[j].d_tag == DT_FINI_ARRAY)
				{
					dyn_entries[j].d_un.d_ptr += offset_ammount;
				}
            }
		}
		char *section_name = get_section_name(woody, i);
		if (section_name && !ft_strcmp(".got", section_name))
		{
			int num_entries = woody->Shdr[i].sh_size / woody->Shdr[i].sh_entsize;
			Elf64_Addr *got_entri = (Elf64_Addr *)fetch(woody->file, woody->file_size, woody->Shdr[i].sh_offset, woody->Shdr[i].sh_size);
            for (int j = 0; j < num_entries; j++)
			{
				if (got_entri[j])
				{
					got_entri[j] += 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;
			woody->Phdr[i].p_vaddr += offset_ammount;
			woody->Phdr[i].p_paddr += 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;
			woody->Shdr[i].sh_addr += offset_ammount;
		}
	}
	for (int i = 1; i < woody->num_symbols; i++)
	{
		if (woody->symbols[i].st_value >= from)
		{
			woody->symbols[i].st_value += offset_ammount;
		}
	}

		
}

void	create_codecave(t_elf_content *woody, t_payload *payload, size_t enlarge_position)
{
	const unsigned int page_size = 4096; // getpagesize(); not authorized
	unsigned int padding_size = ((payload->len / page_size) + 1) * page_size;
	printf("padding size = %d\n", padding_size);
	offset_sections(woody, enlarge_position, padding_size);
	char *new_woody = malloc(woody->file_size + padding_size);
	ft_bzero(new_woody, woody->file_size + padding_size);
	if (!new_woody)
		return ;
	ft_memcpy(new_woody, woody->file, enlarge_position);
	ft_memcpy(new_woody + enlarge_position + padding_size,  woody->file + enlarge_position, woody->file_size - enlarge_position);
	munmap(woody->file, woody->file_size);
	woody->file = new_woody;
	woody->file_size += padding_size;
	woody->Ehdr = (Elf64_Ehdr *)new_woody;
	woody->Ehdr->e_shoff += padding_size;
	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));
}

t_payload *get_payload()
{
	t_payload *payload = malloc(sizeof(t_payload));
	if (!payload)
		return NULL;
	char buffer[1024];
	int fd = open("payload", O_RDONLY);
	if (fd == -1) {
		perror("error opening payload");
		free(payload);
		exit(1);
	}
	payload->len = read(fd, buffer, 1024);
	payload->payload = malloc(sizeof(char) * payload->len);
	ft_memcpy(payload->payload, buffer, payload->len);
	return payload;
}

int	insert_payload(t_elf_content *woody, t_payload *payload, size_t payload_position, int load_segment_index)
{
	char *ptr_jmp_value = ft_strnstr_nullterminated(payload->payload, JUMP_VALUE, payload->len);
	char *ptr_woody = ft_strnstr_nullterminated(payload->payload, WOODY, payload->len);
	char *ptr_text_section = ft_strnstr_nullterminated(payload->payload, TEXT_OFFSET, payload->len);
	char *ptr_section_size = ft_strnstr_nullterminated(payload->payload, SECTION_SIZE, payload->len);
	if (ptr_jmp_value && ptr_woody && ptr_text_section && ptr_section_size)
	{
		int32_t woody_index = ptr_woody - payload->payload;
		int32_t jmp_index = ptr_jmp_value - sizeof(JUMP) - payload->payload;
		int32_t jump_value = ((payload_position + jmp_index + 5) - woody->Ehdr->e_entry) * -1; // 5 = JUMP SIZE (OPCODE + 4 bytes operand)
		ft_memcpy(&payload->payload[jmp_index + 1], &jump_value, sizeof(jump_value));

		int64_t text_index = ptr_text_section - payload->payload;
		int64_t text_value = payload_position - woody->Phdr[load_segment_index].p_offset + woody_index;
		ft_memcpy(&payload->payload[text_index], &text_value, sizeof(text_value));

		int64_t section_index = ptr_section_size - payload->payload;
		int64_t section_value = woody->Phdr[load_segment_index].p_memsz; //woody->text_section->sh_size;
		ft_memcpy(&payload->payload[section_index], &section_value, sizeof(section_value));

		ft_memcpy(woody->file + payload_position, payload->payload, payload->len);
		
		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", payload->len, payload->len);
		printf("Backward offset = %d (%x)(%x)\n", jump_value, jump_value, -jump_value);
		return EXIT_SUCCESS;
	}
	return EXIT_FAILURE;
}

void	inject(t_elf_content *woody)
{
	t_payload *payload = get_payload();
	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 = woody->Phdr[i].p_offset + woody->Phdr[i].p_filesz;

	if (code_cave_size < payload->len) // inverse here to test the other technique
	{
		printf("create code_Cave %ld\n", code_cave_size);
		create_codecave(woody, payload, woody->Phdr[j].p_offset);
	}

	woody->Phdr[i].p_filesz += payload->len;
	woody->Phdr[i].p_memsz += payload->len;
	printf("%ld\n", woody->Phdr[i].p_align);
	(void)payload_position;
	insert_payload(woody, payload, payload_position, i);
	woody->Ehdr->e_entry = payload_position;
	printf("New entry = %ld (%lx)\n", woody->Ehdr->e_entry, woody->Ehdr->e_entry);
	printf("Load segment p_filesz = %lx\n", woody->Phdr[i].p_filesz);
	printf("Load segment p_memsz = %lx\n", woody->Phdr[i].p_memsz);
	woody->Phdr[i].p_flags = PF_X | PF_W | PF_R;
}

int get_elf_sections(t_elf_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)
	{
		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|| !fetch(woody->file, woody->file_size, woody->Ehdr->e_shoff, woody->Ehdr->e_shnum * sizeof(Elf64_Shdr)))
		return EXIT_FAILURE;

    for (int i = 0; i < woody->Ehdr->e_shnum; i++) {
		if (woody->Shdr[i].sh_type == SHT_SYMTAB) {
            woody->symbols_table = fetch(woody->file, woody->file_size, woody->Ehdr->e_shoff + (i * sizeof(Elf64_Shdr)), sizeof(Elf64_Shdr));
        }
		if (woody->Shdr[i].sh_type == SHT_PROGBITS && (woody->Shdr[i].sh_flags & SHF_EXECINSTR) && (woody->Shdr[i].sh_flags & SHF_ALLOC)) {
			woody->text_section = fetch(woody->file, woody->file_size, woody->Ehdr->e_shoff + (i * sizeof(Elf64_Shdr)), sizeof(Elf64_Shdr));
		}

    }
	if (woody->symbols_table == NULL || woody->text_section == NULL)
		return EXIT_FAILURE; //Not sure about this
	woody->symbols = (Elf64_Sym *)fetch(woody->file, woody->file_size, woody->symbols_table->sh_offset, sizeof(Elf64_Sym));
	if (woody->symbols == NULL)
		return EXIT_FAILURE;//Not sure about this
	woody->Sshstrtab = (char *)fetch(woody->file, woody->file_size, woody->Shdr[woody->Ehdr->e_shstrndx].sh_offset, 0);
	if (woody->Sshstrtab == NULL)
		return EXIT_FAILURE; //Not sure about this
    woody->num_symbols = get_symbols_count(woody->symbols_table->sh_size, woody->symbols_table->sh_entsize);
	if (!fetch(woody->file, woody->file_size, woody->symbols_table->sh_offset + (sizeof(Elf64_Sym) * (woody->num_symbols)), 0))
		return EXIT_FAILURE;

	return EXIT_SUCCESS;
}

int prepare_injection(t_elf_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);
	munmap(woody_file, woody->file_size);
	save_elf("woody", woody_file, woody->file_size);
	free(woody_file);
	return EXIT_SUCCESS;
}