#define _GNU_SOURCE

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

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

#include <sys/auxv.h>
#include <sys/stat.h>

#include "dynso.h"
#include "dynso_internal.h"

#define BREAK() asm volatile("int3":::"memory")

/*static void print_link_map(struct link_map* glob) {
	printf("[*] link map:\n");
	for (size_t i = 0; glob; glob = glob->l_next, ++i) {
		printf("- [%zu] %p -> '%s'\n", i, (void*)glob->l_addr, glob->l_name);
	}
}*/

static ptrdiff_t get_priv1_priv2_offset(void) {
	// the offset between the priv1 and priv2 parts of a link_map is a
	// constant but only determinted at runtime, as we don't know with which
	// glibc version we're running (nor do we want to hardcode stuff)

	// link_map of the main executable always comes first (... I hope)
	struct link_map_priv1* exe_map = (struct link_map_priv1*)_r_debug.r_map;
	ElfW(Ehdr)* exe_ehdr = (ElfW(Ehdr)*)exe_map->l_addr;

	// priv2 contains a field with the entrypoint of an ELF
	ElfW(Addr) entry = exe_ehdr->e_entry + exe_map->l_addr;

	// time for crimes
	const ElfW(Addr)* crimes = (const ElfW(Addr)*)&exe_map->l_info[0];

	for (size_t iter = 0;
			*crimes != entry && iter < 4096/sizeof(*crimes);
			++crimes, ++iter)
		;
	if (*crimes != entry) {
		//printf("[!] E: can't get priv2 offset, bailing out...\n");
		return -1;//exit(1);
	}
	// crimes is now &priv2->l_entry;

	ptrdiff_t r = (ptrdiff_t)crimes - (ptrdiff_t)exe_map - offsetof(struct link_map_priv2, l_entry);
	//printf("[-] offset of priv2 is %zd\n", r);
	return r;
}
static ptrdiff_t get_priv2_priv3_offset(void) {
	struct link_map_priv1* exe_map = (struct link_map_priv1*)_r_debug.r_map;
	struct link_map_priv2* exe2 = priv1_to_priv2(exe_map);

	// priv3 has a 'versyms' field, which we can easily fetch from the dyn
	// table for comparison. otherwise, same method as above
	ElfW(Versym)* versym_good = NULL;
	for (size_t i = 0; _DYNAMIC[i].d_tag != DT_NULL && i < 128; ++i) {
		if (_DYNAMIC[i].d_tag == DT_VERSYM) {
			versym_good = (ElfW(Versym*))_DYNAMIC[i].d_un.d_ptr;
			break;
		}
	}

	if (versym_good == NULL) {
		//printf("[!] E: can't find VERSYM, bailing out...\n");
		return -1;//exit(1);
	}

	const ElfW(Versym)** crimes = (const ElfW(Versym)**)&exe2->l_direct_opencount;
	for (size_t iter = 0;
			*crimes != versym_good && iter < 4096/sizeof(*crimes);
			++crimes, ++iter)
		;
	if (*crimes != versym_good) {
		//printf("[!] E: can't get priv3 offset, bailing out...\n");
		return -1;//exit(1);
	}
	// crimes is now &priv3->l_versyms;

	ptrdiff_t r = (ptrdiff_t)crimes - (ptrdiff_t)exe2 - offsetof(struct link_map_priv3, l_versyms);
	//printf("[-] offset of priv3 is %zd\n", r);
	return r;
}

static struct link_map_priv2* priv1_to_priv2(struct link_map_priv1* lm) {
	if (!lm) return NULL;

	static ptrdiff_t d = -1;
	if (d < 0) {
		d = get_priv1_priv2_offset();
		if (d < 0) return NULL;
	}

	return (struct link_map_priv2*)((ptrdiff_t)lm + d);
}
static struct link_map_priv3* priv2_to_priv3(struct link_map_priv2* l2) {
	if (!l2) return NULL;

	static ptrdiff_t d = -1;
	if (d < 0) {
		d = get_priv2_priv3_offset();
		if (d < 0) return NULL;
	}

	return (struct link_map_priv3*)((ptrdiff_t)l2 + d);
}

static enum dynso_err get_rtld_symtab(struct rtldinfo* ret) {
	ElfW(Addr) ldso_base = (ElfW(Addr))getauxval(AT_BASE);
	//printf("[-] getauxval(AT_BASE)         = %p\n", (void*)ldso_base);
	//printf("[-] getauxval(AT_SYSINFO_EHDR) = %p\n", (void*)getauxval(AT_SYSINFO_EHDR));
	struct link_map* glob = _r_debug.r_map;
	//print_link_map(glob);
	for (size_t i = 0; glob; glob = glob->l_next, ++i) {
		if (glob->l_addr != ldso_base) continue;

		//printf("[*] [%zu] %p:%s is ld.so!\n", i, (void*)glob->l_addr, glob->l_name);

		ElfW(Sym)* symtab = NULL;
		char* strtab = NULL;
		size_t strsz = 0;
		size_t shsz = ~(size_t)0;
		ptrdiff_t syment = 0;
		bool sym_malloc = false, str_malloc = false;

		// gives us .dynsym, not .symtab -> kinda useless
		ElfW(Dyn)* dyn = glob->l_ld;

		for (size_t j = 0; dyn->d_tag != DT_NULL; ++dyn, ++j) {
			if (dyn->d_tag == DT_STRTAB) {
				strtab = (char*)(/*glob->l_addr +*/ dyn->d_un.d_ptr);
				//printf(" - [%zu] strtab\n", j);
			}
			if (dyn->d_tag == DT_SYMTAB) {
				symtab = (ElfW(Sym) *)(/*glob->l_addr +*/ dyn->d_un.d_ptr);
				//printf(" - [%zu] symtab\n", j);
			}
			if (dyn->d_tag == DT_STRSZ) {
				strsz = dyn->d_un.d_val;
				//printf(" - [%zu] strsz\n", j);
			}
			if (dyn->d_tag == DT_SYMENT) {
				syment = dyn->d_un.d_val;
				//printf(" - [%zu] syment\n", j);
			}
		}

		if (!memcmp((const void*)ldso_base, ELFMAG, SELFMAG)) {
			//printf("[-] ldso is ELF\n");

			FILE* fff = fopen(glob->l_name, "rb"); // TODO: use l_origin instead?
			if (!fff) goto nah;

			ElfW(Ehdr)* ldso = (ElfW(Ehdr)*)ldso_base;

			if (ldso->e_shoff && ldso->e_shnum && ldso->e_shentsize) {
				//printf("     - shoff=0x%zx shnum=0x%zx\n", ldso->e_shoff, ldso->e_shnum);

				enum dynso_err errr = dynso_ok;
				if (fseek(fff, ldso->e_shoff, SEEK_SET) < 0) errr = dynso_ldso_ioerr;

				uint8_t shdr_blob[ldso->e_shentsize];
				ElfW(Shdr)* shdr = (ElfW(Shdr)*)shdr_blob;

				if (fseek(fff, ldso->e_shoff + ldso->e_shstrndx*ldso->e_shentsize, SEEK_SET) < 0) errr = dynso_ldso_ioerr;
				if (fread(shdr_blob, 1, sizeof(shdr_blob), fff) < 0) errr = dynso_ldso_ioerr;
				char* shstrtab = (char*)malloc(shdr->sh_size);
				if (shstrtab == NULL) errr = dynso_oom;
				if (fseek(fff, shdr->sh_offset, SEEK_SET) < 0) errr = dynso_ldso_ioerr;
				if (fread(shstrtab, 1, shdr->sh_size, fff) < 0) errr = dynso_ldso_ioerr;

				for (size_t j = 0; j < ldso->e_shnum; ++j) {
					if (fseek(fff, ldso->e_shoff + j*ldso->e_shentsize, SEEK_SET) < 0) errr = dynso_ldso_ioerr;
					if (fread(shdr_blob, 1, sizeof(shdr_blob), fff) < 0) errr = dynso_ldso_ioerr;

					//printf("    [-] shdr [%zx] -> %p, type=0x%zx\n", j, shdr, shdr->sh_type);

					if (shdr->sh_type == SHT_SYMTAB) {
						//printf("  [-] using .symtab instead of .dynsym at offset 0x%zx\n", shdr->sh_offset);
						////symtab = (ElfW(Sym)*)(ldso_base + shdr->sh_offset); // not mapped in phdrs...
						sym_malloc = true;
						symtab = (ElfW(Sym)*)malloc(shdr->sh_size);
						if (symtab == NULL) errr = dynso_oom;
						if (fseek(fff, shdr->sh_offset, SEEK_SET) < 0) errr = dynso_ldso_ioerr;
						if (fread(symtab, 1, shdr->sh_size, fff) < 0) errr = dynso_ldso_ioerr;
						shsz = shdr->sh_size;
						syment = shdr->sh_entsize; // not *technically* needed, but doing this for completeness' sake
					}
					if (shdr->sh_type == SHT_STRTAB && !strcmp(&shstrtab[shdr->sh_name], ".strtab")) {
						//printf("  [-] using .strtab instead of .dynstr at offset 0x%zx\n", shdr->sh_offset);
						str_malloc = true;
						strtab = (char*)malloc(shdr->sh_size);
						if (strtab == NULL) errr = dynso_oom;
						if (fseek(fff, shdr->sh_offset, SEEK_SET) < 0) errr = dynso_ldso_ioerr;
						if (fread(strtab, 1, shdr->sh_size, fff) < 0) errr = dynso_ldso_ioerr;
						strsz = shdr->sh_size;
					}
				}

				free(shstrtab);
				if (errr != dynso_ok) {
					if (fff) fclose(fff);
					if (str_malloc && strtab) free(strtab);
					if (sym_malloc && symtab) free(symtab);
					return errr;
				}
			} //else printf("wut\n");
		}

	nah:
		if (symtab == NULL || strtab == NULL) {
			return dynso_ldso_stripped;
		}

		//printf("  [-] symtab=%p\tstrtab=%p\tstrsz=0x%zx\tsyment=0x%zx\n",
		//		symtab, strtab, strsz, syment);
		//printf("[-] resolving hidden needed symbols\n");

		ret->base = ldso_base;
		ret->symtab = symtab;
		ret->strtab = strtab;
		ret->strsz = strsz;
		ret->shsz = shsz;
		ret->syment = syment;
		ret->sym_malloc = sym_malloc;
		ret->str_malloc = str_malloc;
		return dynso_ok;
	}

	//printf("wtf\n");
	return dynso_ldso_not_found;//exit(1);
}


typedef void* (*gnu_ifunc_f)(void);

static void* sym2addr(ElfW(Addr) ldso_base, ElfW(Sym)* sym) {
	void* r = (void*)(ldso_base + sym->st_value);

	if (ELF64_ST_TYPE(sym->st_info) == STT_GNU_IFUNC) { // sigh
		//printf("[-] W: IFUNC. not tested.\n");
		return ((gnu_ifunc_f)(r))();
	} else return r;
}
static enum dynso_err link_private_rtld_syms(struct dynso_lib* l) {
	struct rtldinfo rtldinf;

	enum dynso_err r = get_rtld_symtab(&rtldinf);
	if (r != dynso_ok) return r;

	ElfW(Sym)* sym = rtldinf.symtab;

	for (size_t j = 0; j*rtldinf.syment < rtldinf.shsz;
			++j, sym = (ElfW(Sym*))((ptrdiff_t)sym + rtldinf.syment)) {
		if (ELF64_ST_BIND(sym->st_info) != STB_LOCAL) continue;
		if (sym->st_value == 0) continue;
		if (sym->st_name >= rtldinf.strsz) continue;
		const char* name = &rtldinf.strtab[sym->st_name];
		if (!*name) continue;

		if (!strcmp(name, "_dl_new_object")) {
			//printf("- [%zu] -> info=0x%x other=0x%x value=0x%zx\t'%s'\n", j,
			//	sym->st_info, sym->st_other, sym->st_value, name);

			l->_dl_new_object = (dl_new_object_f)sym2addr(rtldinf.base, sym);
		}
		if (!strcmp(name, "_dl_add_to_namespace_list")) {
			//printf("- [%zu] -> info=0x%x other=0x%x value=0x%zx\t'%s'\n", j,
			//	sym->st_info, sym->st_other, sym->st_value, name);

			l->_dl_add_to_namespace_list = (dl_add_to_namespace_list_f)sym2addr(rtldinf.base, sym);
		}
		if (!strcmp(name, "_dl_setup_hash")) {
			//printf("- [%zu] -> info=0x%x other=0x%x value=0x%zx\t'%s'\n", j,
			//	sym->st_info, sym->st_other, sym->st_value, name);

			l->_dl_setup_hash = (dl_setup_hash_f)sym2addr(rtldinf.base, sym);
		}
	}

	if (rtldinf.sym_malloc) free(rtldinf.symtab);
	if (rtldinf.str_malloc) free(rtldinf.strtab);

	//printf("[*] _dl_new_object=%p\n", _dl_new_object);
	//printf("[*] _dl_add_to_namespace_list=%p\n", _dl_add_to_namespace_list);
	//printf("[*] _dl_setup_hash=%p\n", _dl_setup_hash);

	if (l->_dl_new_object && l->_dl_add_to_namespace_list && l->_dl_setup_hash)
		return dynso_ok;
	else return dynso_no_ldso_syms;
}

static enum dynso_err inject_in_scope_of(struct inject_undo_info* r,
		struct link_map_priv3* vic3, struct link_map_priv1* add) {
	r->alloced = NULL;

	struct r_scope_elem* oldscope = vic3->l_scope[0];

	struct r_scope_elem* scope = NULL;
	scope = (struct r_scope_elem*)calloc(1, sizeof(*scope)
			+ sizeof(struct link_map_priv1*)*(oldscope->r_nlist+1));
	if (scope == NULL) return dynso_oom;
	r->alloced = scope;

	struct link_map_priv1** rlist = (struct link_map_priv1**)(
			(size_t)scope + sizeof(*scope));

	for (size_t i = 0; i < oldscope->r_nlist; ++i)
		rlist[i] = oldscope->r_list[i];
	rlist[oldscope->r_nlist] = add;

	scope->r_list = rlist;
	scope->r_nlist = oldscope->r_nlist + 1;

	// now apply the new scope
	vic3->l_scope[0] = scope;

	r->dest = &vic3->l_scope[0];
	r->old = oldscope;
	//r->rlist = rlist;
	return dynso_ok;
}
static void inject_undo(const struct inject_undo_info* inf) {
	*inf->dest = inf->old;
	if (inf->alloced) free(inf->alloced);
	//free(inf->rlist);
}

enum dynso_err dynso_create(struct dynso_lib** ll, size_t address,
		char* realname, const char* libname, struct link_map* calling, Lmid_t ns) {
	if (!ll) return dynso_argument;

	if (calling == NULL) calling = _r_debug.r_map;

	struct dynso_lib* r = (struct dynso_lib*)calloc(1, sizeof(struct dynso_lib)
			+ 16*sizeof(struct dynso_sym));
	if (r == NULL) return dynso_oom;

	enum dynso_err e = link_private_rtld_syms(r);
	if (e != dynso_ok) { free(r); return e; }

	struct link_map_priv1* lm =
		r->_dl_new_object(realname, libname, lt_loaded, calling, 0, ns);
	if (!lm) { free(r); return dynso_no_newobj; }

	struct link_map_priv2* l2 = priv1_to_priv2(lm);
	struct link_map_priv3* l3 = priv2_to_priv3(l2);
	if (!l2) { free(r); return dynso_no_priv2; }
	if (!l3) { free(r); return dynso_no_priv3; }

	l2->l_faked = 1;
	l2->l_nbuckets = 0; // disables symbol lookup
	l2->l_gnu_bitmask = NULL; // disable GNU hashing, use SysV hashing
	l3->l_initfini = NULL;
	lm->l_addr = address;

	r->ns = ns;
	r->caller = (struct link_map_priv1*)calling;
	r->fake1 = lm;
	r->fake2 = l2;
	r->fake3 = l3;
	r->address = address;

	r->syms = (struct dynso_sym*)((ptrdiff_t)r + sizeof(*r));
	r->nsyms = 0;
	r->syms_cap = 16;

	*ll = r;

	return dynso_ok;
}

enum dynso_err dynso_add_sym_ex(struct dynso_lib* f, const char* name,
		void* value, int type, size_t size) {
	if (!f || f->added || f->removed) return dynso_argument;

	if (f->nsyms == f->syms_cap) {
		f->syms_cap <<= 1;

		if (f->alloc_syms) {
			f->syms = (struct dynso_sym*)realloc(
					f->syms, f->syms_cap * sizeof(struct dynso_sym));
			if (f->syms == NULL) return dynso_oom;
		} else {
			struct dynso_sym* nsym = (struct dynso_sym*)calloc(f->syms_cap, sizeof(struct dynso_sym));
			if (nsym == NULL) return dynso_oom;

			memcpy(nsym, f->syms, sizeof(struct dynso_sym)*f->nsyms);
			f->syms = nsym;
			f->alloc_syms = true;
		}
	}

	f->syms[f->nsyms].name  = name ;
	f->syms[f->nsyms].value = value;
	f->syms[f->nsyms].type  = type;
	f->syms[f->nsyms].size  = size;
	++f->nsyms;

	return dynso_ok;
}

enum dynso_err dynso_bind(struct dynso_lib* l) {
	if (!l || l->removed || l->added) return dynso_argument;

	struct link_map_priv1* lm = l->fake1;
	struct link_map_priv2* l2 = l->fake2;
	struct link_map_priv3* l3 = l->fake3;

	lm->l_addr = l->address;

	size_t strsz = 2; // beginning and ending null terminators
	for (size_t i = 0; i < l->nsyms; ++i) {
		strsz += strlen(l->syms[i].name) + 1; // include null terminator
	}

	ElfW(Dyn)* dyn = (ElfW(Dyn)*)calloc(5, sizeof(ElfW(Dyn)));
	if (!dyn) return dynso_oom;
	Elf_Symndx* idx = (Elf_Symndx*)calloc(l->nsyms + 2, sizeof(Elf_Symndx));
	if (!idx) { free(dyn); return dynso_oom; }
	ElfW(Sym)* sym = (ElfW(Sym)*)calloc(l->nsyms + 2, sizeof(ElfW(Sym)));
	if (!sym) { free(dyn); free(idx); return dynso_oom; }
	char* str = (char*)calloc(strsz, 1);
	if (!str) { free(dyn); free(idx); free(sym); return dynso_oom; }

	// NOTE: should we give it a phdr? seems to work without for now...

	str[0] = 0;
	//BREAK();
	//sym[0] = ((ElfW(Sym)){.st_name=0, .st_info=0, .st_other=0, .st_shndx=0, .st_value=0, .st_size=0};
	for (size_t i = 0, strind = 1; i < l->nsyms; ++i) {
		int typ = l->syms[i].type & 0xf;
		if (typ == 0) typ = STT_FUNC;

		sym[i+1].st_name = strind;
		sym[i+1].st_info = typ | (STB_GLOBAL << 4);
		sym[i+1].st_other = STV_DEFAULT;
		sym[i+1].st_shndx = 0;
		sym[i+1].st_value = (ElfW(Addr))l->syms[i].value - l->address;
		sym[i+1].st_size = 0;

		strcpy(&str[strind], l->syms[i].name);
		strind += strlen(l->syms[i].name) + 1;
	}
	//sym[l->nsyms+1] = ((ElfW(Sym)){.st_name=0, .st_info=0, .st_other=0, .st_shndx=0, .st_value=0, .st_size=0};
	str[strsz-1] = 0;

	for (size_t i = 0; i < l->nsyms; ++i) idx[i] = i + 1;
	idx[l->nsyms+0] = STN_UNDEF;
	idx[l->nsyms+1] = SHN_UNDEF;

	dyn[0].d_tag = DT_SYMTAB;
	dyn[0].d_un.d_ptr = (ElfW(Addr))sym;
	dyn[1].d_tag = DT_STRTAB;
	dyn[1].d_un.d_ptr = (ElfW(Addr))str;
	// these two aren't strictly needed, but adding them for completeness' sake
	dyn[2].d_tag = DT_STRSZ;
	dyn[2].d_un.d_val = strsz;
	dyn[3].d_tag = DT_SYMENT;
	dyn[3].d_un.d_val = sizeof(ElfW(Sym));

	dyn[4].d_tag = DT_NULL;

	lm->l_ld = dyn;
	size_t ldnum = 0;
	for (size_t i = 0; dyn[i].d_tag != DT_NULL; ++i, ++ldnum)
		lm->l_info[dyn[i].d_tag] = &dyn[i];
	l2->l_ldnum = ldnum;

	// lm->l_origin = "something"; // maybe? eh.

	l->_dl_setup_hash(lm);

	l2->l_relocated = 1;
	l3->l_used = 1;

	l2->l_global = 1;
	l2->l_faked = 1;
	//l2->l_type = lt_loaded;
	l2->l_nbuckets = 1;
	l2->l_buckets = idx; // TODO
	l2->l_chain = idx; // TODO
	l2->l_gnu_bitmask = NULL; // disable GNU hashing, use SysV legacy hashing (was easier to 'implement' on my end)
	l3->l_initfini = NULL;

	l->_dl_add_to_namespace_list(l->fake1, l->ns);

	enum dynso_err r = inject_in_scope_of(&l->undoinf, priv1_to_priv3(l->caller), lm);
	if (r == dynso_ok) {
		l->dyn = dyn;
		l->sym = sym;
		l->idx = idx;
		l->str = str;

		l->added = true;

		return dynso_ok;
	}

	free(dyn); free(idx); free(sym); free(str);

	return r;
}

void dynso_remove(struct dynso_lib* l) {
	if (!l) return;

	if (l->added) {
		inject_undo(&l->undoinf);

		for (size_t i = 0; l->dyn[i].d_tag != DT_NULL; ++i)
			l->fake1->l_info[l->dyn[i].d_tag] = NULL;

		l->fake1->l_ld = NULL;
		l->fake2->l_ldnum = 0;

		free(l->dyn);
		free(l->sym);
		free(l->idx);
		free(l->str);
	}
	if (l->alloc_syms) free(l->syms);

	// TODO: free dyn, symtab, strtab, idxtab

	l->fake2->l_removed = 1;
	l->fake3->l_used = 0;
	l->fake2->l_global = 0;
	l->fake2->l_nbuckets = 0;

	l->removed = true;
	free(l);
}