#include "common.h"

#include "../modules/printk.h"

// this interface needs to be changed to support picking modules etc..

int get_module_size()
{
	return 4 * 4096;
}

int link_module(unsigned int mapped_at, unsigned char **code, int *len, unsigned int *entry_point)
{
	Elf32_Ehdr *ehdr;	// program header
	Elf32_Shdr *shdr;	// section header array
	Elf32_Shdr *textsh;	// .text section header
	Elf32_Sym *symtab;	// symbol table array
	int sym_cnt; 		// number of symbols
	Elf32_Rel *textrel;	// text relocations
	int trel_cnt;		// "" count
	
	int i;			// this isn't .edu :P~ only reason I'm commenting the variables 
				// is because its a rather large function, and will make it easier 
				// to reference later on.

	char *shstrtab;		// section names
	char *strtab;		// symbol string table;
	
	unsigned char *base;	// where we have allocated our buffer..
	unsigned char *text;	// pointer to the text section
	
	// shut the compiler up ;)
	
	textsh = NULL; 
	sym_cnt = 0;
	trel_cnt = 0;
	
	if((base = malloc(get_module_size())) == NULL) {
		printf("[!] Unable to allocate space for the temporary module linking area.\n");
		exit(EXIT_FAILURE);
	}
	
	ehdr = (Elf32_Ehdr *)(printk_o);
	shdr = (Elf32_Shdr *)((unsigned char *)(printk_o + ehdr->e_shoff));

	shstrtab = printk_o + shdr[ehdr->e_shstrndx].sh_offset;
	
	symtab = NULL;
	strtab = NULL;
	textrel = NULL;
	text = NULL;
	
	for(i = 0; i < ehdr->e_shnum; i++) {
		printf("[*] --> %s\n", shstrtab + shdr[i].sh_name);

		if(strcmp(shstrtab + shdr[i].sh_name, ".symtab") == 0) {
			symtab = (Elf32_Sym *)((unsigned char *)(printk_o + shdr[i].sh_offset));
			sym_cnt = shdr[i].sh_size / sizeof(Elf32_Sym);
			continue;
		} 

		if(strcmp(shstrtab + shdr[i].sh_name, ".strtab") == 0) {
			strtab = (char *)(printk_o + shdr[i].sh_offset);
			continue;
		}

		if(strcmp(shstrtab + shdr[i].sh_name, ".text") == 0) {
			text = (unsigned char *)(printk_o + shdr[i].sh_offset);
			textsh = (Elf32_Shdr *)(&shdr[i]);
			continue;
		}

		if(strcmp(shstrtab + shdr[i].sh_name, ".rel.text") == 0) {
			textrel = (Elf32_Rel *)((unsigned char *)(printk_o + shdr[i].sh_offset));
			trel_cnt = shdr[i].sh_size / sizeof(Elf32_Rel);
			continue;
		}
	}
	
	if(!symtab || !strtab || !textrel || !text) {
		printf("[!] Unable to find a required elf object.\n");
		printf("[!]	symtab: 0x%08x\n", (unsigned int)symtab);
		printf("[!]	strtab: 0x%08x\n", (unsigned int)strtab);
		printf("[!]	textrel: 0x%08x\n", (unsigned int) textrel);
		printf("[!]	text: 0x%08x\n", (unsigned int) text);
		exit(EXIT_FAILURE);
	}

	// process the text relocations
	//
	// These are just my notes for the time being..
	//
	// typedef struct
	// {
	//	Elf32_Addr    r_offset;               /* Address */
	//	Elf32_Word    r_info;                 /* Relocation type and symbol index */
	// } Elf32_Rel;
	//     
	// #define ELF32_R_SYM(val)                ((val) >> 8)
	// #define ELF32_R_TYPE(val)               ((val) & 0xff)
	//
	// #define R_386_32           1            /* Direct 32 bit  */
	// #define R_386_PC32         2            /* PC relative 32 bit */
	// 
	
	
	for(i = 0; i < trel_cnt; i++) {
		Elf32_Rel *r = textrel + i;			// reloc entry
		Elf32_Sym *s = symtab + ELF32_R_SYM(r->r_info);	// symbol entry
		Elf32_Shdr *S = &shdr[s->st_shndx];
		
		int type = ELF32_R_TYPE(r->r_info);
		unsigned int addend;		// added to the symbol value
		unsigned int val;		// value we're going to write.
		unsigned int tmp;	
		
		printf("Relocation entry %d\n", i);
		
		addend = *(int *)((unsigned char *)(text + r->r_offset));	// obtain the addend

		val = 0;
		
		switch(type) {
			case R_386_32:
				if(strcmp(strtab + s->st_name, "") != 0) {
					printf("We got a request for symbol %s in R_386_32.. unexpected.\n", strtab + s->st_name);
					exit(EXIT_FAILURE);
				}
				
				printf("Section header: %s\n", shstrtab + S->sh_name);
				val = mapped_at + S->sh_offset + addend;

				printf("Got 0x%08x for value to write in.\n", val);
				*(int *)((unsigned char *)(text + r->r_offset)) = val;
				
				
				break;
				
			case R_386_PC32:
				
				if(strcmp(strtab + s->st_name, "") == 0) {
					printf("We got a request for a symbol in R_386_PC32.. unexpected.\n");
					exit(EXIT_FAILURE);
				}
				
				if(ELF32_ST_TYPE(s->st_info) == STT_NOTYPE) {
					if((val = locate_symbol(strtab + s->st_name)) == 0) {
						printf("[!] Unable to locate symbol %s\n", strtab + s->st_name);
						exit(EXIT_FAILURE);
					}
				} else {
					val = mapped_at + S->sh_offset + s->st_value;
				}
	
				printf("%s is at 0x%08x\n", strtab + s->st_name, val);
				
				printf("name: %s\n", shstrtab + S->sh_name);
				
				tmp = mapped_at + textsh->sh_offset + r->r_offset - addend;
				
				printf("reloc is at 0x%08x\n", tmp);
				tmp = val - tmp;
				printf("relative jump is 0x%08x\n", tmp);

				*(int *)((unsigned char *)(text + r->r_offset)) = tmp;
				
				break;
			default:
				printf("Got a relocation type %d that this doesn't handle.. bailing.\n", type);
				
		}
		
	}

	*code = printk_o;
	*len = printk_o_len;
	
	*entry_point = 0;
	
	for(i = 0; i < sym_cnt; i++) {
		Elf32_Sym *s = &symtab[i];
		
		if(strcmp(strtab + s->st_name, "init_module") == 0) {
			*entry_point = mapped_at + textsh->sh_offset + s->st_value;
			break;
		}
	}

	return 0;
	
}