#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <elf.h>
#include <link.h>

/*#if sizeof(ElfW(Ehdr)) != sizeof(Elf64_Ehdr)
#error "Need 64-bit ELF for now, sorry"
#endif*/

static bool manip_elf(ElfW(Ehdr)* elf, const char* symname) {
	if (elf->e_ident[0] != ELFMAG0 || elf->e_ident[1] != ELFMAG1
		 	|| elf->e_ident[2] != ELFMAG2 || elf->e_ident[3] != ELFMAG3) {
		printf("bad elf magic\n");
		return true;
	}

	// TODO: port to other architectures
	//       what would be needed for a port:
	//       * fix code looking for 'lea rdi, [rel main]'
	//       * fix relocation type
	//       * thats it i think??? (well also fix this validation stuff)
	if (elf->e_ident[EI_CLASS] != ELFCLASS64 || elf->e_ident[EI_DATA] != ELFDATA2LSB) {
		printf("not a 64-bit little-endian elf\n");
		return true;
	}
	if (elf->e_type != ET_DYN || elf->e_machine != EM_X86_64) {
		printf("need dynamic x86_64 elf\n");
		return true;
	}

	ElfW(Phdr)* phdr = (ElfW(Phdr)*)(elf->e_phoff + (size_t)elf), *ph_loadexec;
	ElfW(Dyn)* dyn = NULL;
	bool textrel = false;
	enum {
		libc_unk, libc_g /* glibc */, libc_musl, libc_uc /* uclibc */
	} libc = libc_unk;
	// first, try to determine the libc used
	for (size_t i = 0; i < elf->e_phnum;
			++i, phdr = (ElfW(Phdr)*)((size_t)phdr + elf->e_phentsize)) {
		if (phdr->p_type == PT_INTERP) {
			const char* str = (char*)elf + phdr->p_offset;
			if (!strcmp(str, "/lib64/ld-linux-x86-64.so.2")) {
				libc = libc_g;
				printf("\tglibc\n");
			} else if (!strcmp(str, "/lib/ld-musl-x86_64.so.1")) {
				libc = libc_musl;
				printf("\tmusl\n");
			} else {
				printf("\tunknown libc: '%s'\n", str);
			}
		}

		if (phdr->p_type == PT_DYNAMIC) {
			dyn = (ElfW(Dyn)*)((size_t)elf + phdr->p_offset);
			printf("dyn=%zu\n", phdr->p_offset);
		}
	}

	if (libc == libc_unk) {
		ElfW(Dyn)* dd = dyn;
		const char* ddstr = NULL;
		for (size_t i = 0; dd->d_tag != DT_NULL; ++i, ++dd) {
			if (dd->d_tag == DT_STRTAB)
				ddstr = (const char*)((size_t)elf + dd->d_un.d_ptr);
		}
		dd = dyn;
		for (size_t i = 0; dd->d_tag != DT_NULL; ++i, ++dd) {
			if (dd->d_tag == DT_NEEDED) {
				const char* s = ddstr + dd->d_un.d_val;

				if (strstr(s, "libc.so") != s) continue;

				if (!strcmp(s, "libc.so.6")) {
					libc = libc_g;
					printf("\tglibc\n");
					break;
				} else if (!strcmp(s, "libc.so.0")) {
					libc = libc_uc;
					printf("\tuclibc\n");
					break;
				} else if (!strcmp(s, "libc.so")) {
					libc = libc_musl;
					printf("\tmusl\n");
					break;
				} else {
					printf("\tunknown libc: '%s'\n", s);
				}
			}
		}
	}

	phdr = (ElfW(Phdr)*)(elf->e_phoff + (size_t)elf);
	for (size_t i = 0; i < elf->e_phnum;
			++i, phdr = (ElfW(Phdr)*)((size_t)phdr + elf->e_phentsize)) {
		if (phdr->p_type == PT_LOAD) {
			// need this at runtime for the relocs to work (yeah, sorry)
			if (phdr->p_flags & PF_X) {
				ph_loadexec = phdr;
				if (libc == libc_g || libc == libc_uc) {
					textrel = true;
				} else {
					phdr->p_flags |= PF_W;
				}
			}
		}
	}

	if (dyn == NULL) {
		fprintf(stderr, "No DYNAMIC table! bailing out...\n");
		return true;
	}

	ElfW(Addr)* relasz = NULL, *relaent = NULL, *jmprel = NULL;
	ElfW(Rela)* relatab = NULL;
	ElfW(Sym)* dynsym = NULL;
	const char* dynstr = NULL;
	size_t strsz = 0, syment = 0;

	ElfW(Addr) relaoff = 0;

	for (size_t i = 0; dyn->d_tag != DT_NULL; ++i, ++dyn) {
		if (dyn->d_tag == DT_TEXTREL) textrel = false;
		if (dyn->d_tag == DT_RELA) {
			relatab = (ElfW(Rela)*)((size_t)elf + dyn->d_un.d_ptr);
			relaoff = dyn->d_un.d_ptr;
		}
		if (dyn->d_tag == DT_RELASZ) relasz = (ElfW(Addr)*)(&dyn->d_un.d_ptr);
		if (dyn->d_tag == DT_RELAENT) relaent = (ElfW(Addr)*)(&dyn->d_un.d_ptr);
		if (dyn->d_tag == DT_JMPREL) jmprel = (ElfW(Addr)*)(&dyn->d_un.d_ptr);

		if (dyn->d_tag == DT_REL) {
			printf("WARN: REL table present in DYN\n");
		}

		if (dyn->d_tag == DT_SYMTAB) dynsym = (ElfW(Sym)*)((size_t)elf + dyn->d_un.d_ptr);
		if (dyn->d_tag == DT_STRTAB) dynstr = (const char*)((size_t)elf + dyn->d_un.d_ptr);
		if (dyn->d_tag == DT_SYMENT) syment = dyn->d_un.d_val;
		if (dyn->d_tag == DT_STRSZ ) strsz  = dyn->d_un.d_val;
	}
	if (textrel) {
		dyn->d_tag = DT_TEXTREL; // need this one
		dyn[1].d_tag = DT_NULL;
	}

	if (!relasz || !relaent || !relatab) {
		fprintf(stderr, "No RELA/RELASZ/RELAENT in DYN!\n");
		return true;
	}
	if (!dynsym || !dynstr || !strsz || !syment) {
		fprintf(stderr, "no SYMTAB/STRTAB/STRSZ/SYMENT in DYN!\n");
		return true;
	}

	// now we have everything to start patching stuff

	static const uint8_t magic_code[] = {0x48,0x8d,0x3d}; // lea rdi, [rel <pcrel32>]

	// assumption: the first matching instruction loads 'main' into rdi
	// both glibc and musl load main this way, they could do it in a different
	// way, but apparently it's not the case
	ElfW(Addr) reloctarget = 0;
	ElfW(Addr) entry_vma = elf->e_entry;
	ElfW(Addr) entry_lma = elf->e_entry - ph_loadexec->p_vaddr + ph_loadexec->p_offset;
	printf("entry_vma=0x%zx, entry_lma=0x%zx\n", entry_vma, entry_lma);
	void* entry_mem = (void*)(entry_lma + (size_t)elf);
	for (size_t i = 0; entry_lma+i - ph_loadexec->p_offset < ph_loadexec->p_filesz; ++i) {
		const uint8_t* tocmp = (const uint8_t*)((size_t)entry_mem + i);

		for (size_t j = 0; j < sizeof(magic_code)/sizeof(*magic_code); ++j) {
			if (tocmp[j] != magic_code[j]) goto next;
		}

		// found it!
		reloctarget = entry_vma + i + sizeof(magic_code)/sizeof(*magic_code);
		printf("target = 0x%zx\n", reloctarget);
		break;

	next:;
	}
	if (!reloctarget) {
		printf("couldn't find relocation target...\n");
		return true;
	}

	ElfW(Word) sym = 0;

	for (size_t i = 0; (size_t)dynsym < (size_t)dynstr; ++i, dynsym = (ElfW(Sym)*)((size_t)dynsym + syment)) {
		if (i && dynsym->st_name != 0 && dynsym->st_name < strsz
				&& !strcmp(symname, dynstr + dynsym->st_name)) {
			sym = i;
			break;
		}
	}

	if (!sym) {
		printf("replacement symbol '%s' not found!\n", symname);
		return true;
	}
	printf("symbol '%s' -> index %zu\n", symname, sym);

	ElfW(Rela) extrarel;
	extrarel.r_offset = reloctarget;
	extrarel.r_info = ELF64_R_INFO(sym, R_X86_64_PC32); // TODO: make portable
	extrarel.r_addend = -4; // TODO: make portable

	printf("extrarel: *0x%016zx = 0x%016zx + %d\n", extrarel.r_offset, extrarel.r_info, extrarel.r_addend);

	// resize the phdr containing the relocations
	phdr = (ElfW(Phdr)*)(elf->e_phoff + (size_t)elf);
	for (size_t i = 0; i < elf->e_phnum;
			++i, phdr = (ElfW(Phdr)*)((size_t)phdr + elf->e_phentsize)) {
		if (phdr->p_offset >= relaoff && phdr->p_offset + phdr->p_filesz < relaoff) {
			printf("rela phdr index = %zu\n", i);
			phdr->p_filesz += *relaent;
			phdr->p_memsz += *relaent;

			break;
		}
	}
	// resize the shdr containing the relocations
	// FIXME: ASSUMPTION: the relocation tables end up at the very end of this
	//                    phdr, first .rela.dyn and then .rela.plt, and no others
	//        things will break BADLY if this isn't true, and it's hard to detect
	//        if this will happen, so watch out!
	ElfW(Shdr)* shdr = (ElfW(Shdr)*)(elf->e_shoff + (size_t)elf);
	for (size_t i = 0; i < elf->e_shnum;
			++i, shdr = (ElfW(Shdr)*)((size_t)shdr + elf->e_shentsize)) {
		if (shdr->sh_type == SHT_REL) {
			printf("WARN: REL table present\n");
		} else if (shdr->sh_type == SHT_RELA) {
			if (shdr->sh_offset == relaoff) { // .rela.dyn
				//shdr->sh_size += shdr->sh_entsize; // one extra entry
			} else { // .rela.plt
				if (jmprel && *jmprel != shdr->sh_offset) {
					printf("WARN: RELA tables look fishy\n");
				}

				// move one entry below to make room
				memmove((void*)((size_t)elf + shdr->sh_offset + shdr->sh_entsize),
						(void*)((size_t)elf + shdr->sh_offset),
						shdr->sh_size);
				shdr->sh_offset += shdr->sh_entsize;
				if (jmprel) *jmprel += shdr->sh_entsize;
			}
		}
	}

	// FINALLY write back the actually relevant stuff
	printf("relasz = %zu, index = %zu\n", *relasz, *relasz / *relaent);
	relatab[*relasz / *relaent] = extrarel;
	*relasz += *relaent;

	printf("relasz now %zu\n", *relasz);

	return false;
}