#include "woody.h"

int wdy_get_file(char *path, t_map *file) {
	int fd = open(path, O_RDONLY);
	if (fd == -1) {
		return wdy_perror(path);
	}

	file->size = lseek(fd, 0, SEEK_END);
	if (file->size == -1) {
		close(fd);
		return wdy_perror(path);
	}

	file->data = mmap(
			NULL,
			file->size,
			PROT_READ | PROT_WRITE,
			MAP_PRIVATE,
			fd,
			0);
	close(fd);
	if (file->data == MAP_FAILED) {
		return wdy_perror(path);
	}

	return RET_OK;
}

int pack_elf32(t_map file) {
	(void)file;
	return wdy_error("ELF needs to be in 64 bits format");
}

int get_load_segment64(
		Elf64_Ehdr elf_header,
		Elf64_Phdr *program_headers,
		Elf64_Phdr **load_segment) {
	for (int i = 0; i < elf_header.e_phnum; i++) {
		Elf64_Phdr *p_hdr = &program_headers[i];
		if (p_hdr->p_type == PT_LOAD
				&& p_hdr->p_flags & PF_X
				&& p_hdr->p_vaddr <= elf_header.e_entry
				&& p_hdr->p_vaddr + p_hdr->p_filesz > elf_header.e_entry) {
			*load_segment = p_hdr;
			return RET_OK;
		}
	}
	return RET_ERR;
}

// TODO refacto
t_map get_code_cave(t_map file, Elf64_Phdr load_segment) {
	size_t page_size = load_segment.p_align;
	size_t len = load_segment.p_filesz + page_size - load_segment.p_filesz % page_size; 
	size_t offset = load_segment.p_offset;
	unsigned char *seg = fetch(file, offset, len);
	if (!seg) {
		printf("fallback to weird code cave finder\n");
		seg = fetch(file, offset, file.size - offset);
		len = file.size - offset;
		if (!seg) {
			ft_printf("unreachable !!!\n");
		}
	}

	size_t longest = 0;
	size_t longest_i = 0;
	for (size_t i = 0; i < len; i++) {
		if (seg[i] == 0) {
			size_t j = 0;
			while (i + j < len && seg[i + j] == 0) {
				j++;
			}

			if (j > longest) {
				longest_i = i;
				longest = j;
			}
			i += j;

		}
	}

	t_map code_cave;
	code_cave.data = fetch(file, longest_i + offset, longest);
	code_cave.size = longest;
	if (!code_cave.data) {
		ft_printf("unreachable !!!\n");
	}
	return code_cave;
}

int encrypt_default(
		__attribute__((unused)) t_map file,
		__attribute__((unused)) Elf64_Phdr load_segment) {
	return RET_OK;
}

int encrypt_xor(t_map file, Elf64_Phdr load_segment) {
	uint64_t key = 0xabcdefabcdefabcd;
	printf("xor begin\n");
	size_t i = 0;
	while (i < load_segment.p_filesz >> 3) {
		((uint64_t *)file.data)[i + (load_segment.p_offset >> 3)] ^= key;
		i++;
	}
	key <<= (8 - (load_segment.p_filesz % 8)) * 8;
	key >>= (8 - (load_segment.p_filesz % 8)) * 8;
	((uint64_t *)file.data)[i + (load_segment.p_offset >> 3)] ^= key;
	// TODO end encrypt and fix asm
	printf("xor ok\n");
	return RET_OK;
}

t_payload64 get_xor_payload64(void) {
	t_payload64 payload;

	size_t diff = 30;

	payload.len = 191 - diff;
	payload.jump_offset = 149 - diff;
	payload.load_ptr_offset = 167 - diff;
	payload.load_size_offset = 175 - diff;
	payload.encrypt = &encrypt_xor;
	payload.data = malloc(payload.len * sizeof(unsigned char));
	ft_memcpy(payload.data,
			"\x55\x54\x53\x50\x51\x52\x56\x57\x41\x50\x48\x8d\x35\x6a\x00\x00"
			"\x00\x48\x89\xf3\x48\x2b\x1d\x6e\x00\x00\x00\x4c\x8b\x05\x6f\x00"
			"\x00\x00\x48\x8b\x15\x70\x00\x00\x00\x4c\x89\xc1\x48\xf7\xd1\x48"
			"\x83\xe1\x07\x48\xff\xc1\x48\xc1\xe1\x03\x48\x89\xd0\x48\xd3\xe0"
			"\x48\xd3\xe8\x49\x83\xe0\xf8\x4a\x31\x04\x03\x49\x83\xe8\x08\x49"
			"\x83\xf8\xf8\x74\x06\x4a\x31\x14\x03\xeb\xf0\xbf\x01\x00\x00\x00"
			"\xba\x0e\x00\x00\x00\xb8\x01\x00\x00\x00\x0f\x05\x41\x58\x5f\x5e"
			"\x5a\x59\x58\x5b\x5c\x5d\xe9\x5f\xda\xda\xda\x2e\x2e\x2e\x2e\x57"
			"\x4f\x4f\x44\x59\x2e\x2e\x2e\x2e\x0a\xba\xba\xba\xba\xba\xba\xba"
			"\xba\xca\xca\xca\xca\xca\xca\xca\xca\xcd\xab\xef\xcd\xab\xef\xcd"
			"\xab",
			payload.len);
	return payload;
}

t_payload64 get_debug_payload64(void) {
	t_payload64 payload;

	payload.len = 55;
	payload.jump_offset = 37;
	payload.encrypt = &encrypt_default;
	payload.data = malloc(payload.len * sizeof(unsigned char));
	ft_memcpy(payload.data,
			"\x55\x54\x50\x52\x56\x57\xbf\x01\x00\x00\x00\x48\x8d\x35\x17\x00"
			"\x00\x00\xba\x0e\x00\x00\x00\xb8\x01\x00\x00\x00\x0f\x05\x5f\x5e"
			"\x5a\x58\x5c\x5d\xe9\xb1\xda\xda\xda\x2e\x2e\x2e\x2e\x57\x4f\x4f"
			"\x44\x59\x2e\x2e\x2e\x2e\x0a",
			payload.len);
	return payload;
}

int pack_elf64(t_map file) {
	Elf64_Ehdr *elf_header = (Elf64_Ehdr *)fetch(file, 0, sizeof(Elf64_Ehdr));
	if (!elf_header) {
		return wdy_error("cannot fetch elf header");
	}

	size_t phdrs_len = elf_header->e_phnum * elf_header->e_phentsize;
	Elf64_Phdr *program_headers = fetch(file, elf_header->e_phoff, phdrs_len);
	if (!program_headers) {
		return wdy_error("cannot fetch program headers table");
	}

	Elf64_Phdr *load_segment;
	if (get_load_segment64(*elf_header, program_headers, &load_segment) == RET_ERR) {
		return wdy_error("cannot get load segment");
	}

	// TODO the whole section below till the end of the function needs 
	// cleaning and refactoring, better naming also.

	// This should go in a separate function regarding encryption if encryption has been requested.
	/*
	   size_t encryption_block_size = 31; // 32 bits and padding...
	   t_map to_encrypt;

	   to_encrypt.data = fetch(file, load_segment.p_offset, 0);
	// size is a multiple of 31...
	to_encrypt.size = load_segment.p_filesz + encryption_block_size - load_segment.p_filesz % encryption_block_size;
	 */

	t_map code_cave = get_code_cave(file, *load_segment);

	t_payload64 payload = get_xor_payload64();

	// This should fallback to compression algorithm, or smaller payload (eg rsa->xor)
	if (payload.len > (size_t)code_cave.size) {
		printf("code cave size: %ld (0x%lx) bytes\n", code_cave.size, code_cave.size);
		return wdy_error("payload length exceed code cave size");
	}

	// e_entry because relative to this (where we gonna go)
	// code cave start because thats the start of the code
	// jump_offset is the index of jump from code cave start
	// 4 because jump has a 4 byte operand)
	size_t code_cave_start = code_cave.data - file.data;
	int jump_value = elf_header->e_entry - code_cave_start - payload.jump_offset - 4;

	ft_memcpy(payload.data + payload.jump_offset, &jump_value, sizeof(jump_value));

	uint64_t load_ptr_value = code_cave_start - load_segment->p_offset + payload.jump_offset + 4;
	ft_memcpy(payload.data + payload.load_ptr_offset, &load_ptr_value, sizeof(load_ptr_value));

	ft_memcpy(payload.data + payload.load_size_offset, &load_segment->p_memsz, sizeof(load_segment->p_memsz));
	payload.encrypt(file, *load_segment);
	elf_header->e_entry = code_cave.data - file.data;
	load_segment->p_filesz += payload.len;
	load_segment->p_memsz += payload.len;
	load_segment->p_flags |= PF_W | PF_R;
	ft_memcpy(code_cave.data, payload.data, payload.len);

	int fd = open("woody", O_WRONLY | O_CREAT, 0755);
	if (fd == -1) {
		return wdy_perror("woody");
	}
	write(fd, file.data, file.size);

	return RET_OK;
}

int pack_elf(t_map file) {	
	// at this point arch is known to be either ELFCLASS32 or ELFCLASS64
	unsigned char arch = ((unsigned char *)file.data)[4];
	if (arch == ELFCLASS64) {
		return pack_elf64(file);
	}
	return pack_elf32(file);
}

int woody_woodpacker(char *path) {
	t_map file;

	if (wdy_get_file(path, &file) == RET_ERR) {
		return RET_ERR;
	}
	if (check_ident(file) == RET_ERR) {
		return RET_ERR;
	}
	return pack_elf(file);
}