#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/mman.h>
#include <fcntl.h>
#include <stdio.h>
#include <errno.h>
#include <elf.h>

#include "libdis.h"

Elf32_Ehdr elf_header;
Elf32_Phdr *prog_headers;

#define PAGESIZE 4096

//#define stdout stdout // don't ask.

struct {
	unsigned char *src;
	unsigned char *dst;
	int len;
	int flags;
	int type;
} things[4096];
int things_count;

void output(unsigned char *src, unsigned char *dst, x86_insn_t *insn)
{
	printf("src: 0x%08x, dst: 0x%08x\n", src, dst);
	things[things_count].src = src;
	things[things_count].dst = dst;
	things[things_count].len = insn->size;
	things[things_count].type = insn->type;
	things[things_count].flags = insn->flags_tested; 
	things_count++;
}

struct {
	unsigned long what;
	unsigned long len;
} memstuff[100];
int memcount;

void patch_binary()
{
	FILE *f;
	int i,j;
	
	f = fopen("int3_loader.h", "w+");
	fprintf(f, "struct entries{\n");
	fprintf(f, "\tunsigned long src;\n");
	fprintf(f, "\tunsigned long dst;\n");
	fprintf(f, "\tunsigned int len;\n");
	fprintf(f, "\tunsigned int type;\n");
	fprintf(f, "\tunsigned int flags;\n");
	fprintf(f, "} jumpme[] = {\n");
	for(i = 0; i < things_count; i++) {
		fprintf(f, "\t{ %p, %p, 0x%08x, 0x%08x, 0x%08x },\n", (unsigned long)(things[i].src)+1, (unsigned long)(things[i].dst), things[i].len - 1, things[i].type, things[i].flags);
	}
	fseek(f, -2, SEEK_CUR);
	fprintf(f, "\n};\n");
	fclose(f);

	for(i = 0; i < things_count; i++) {
		things[i].src[0] = 0xf1;
		for(j = 1; j < things[i].len; j++) things[i].src[j] = 0x90;
		things[i].src = things[i].dst = NULL;
		things[i].len = things[i].flags = 0;
	}
	things_count = 0;

	for(i = 0; i < memcount; i++) {
		msync(memstuff[i].what, memstuff[i].len, MS_SYNC);
		munmap(memstuff[i].what, memstuff[i].len);
	}
	memcount = 0;
}

void load_elf(int fd)
{	
	unsigned long min, offset, i, prot;

	
	if(read(fd, &elf_header, sizeof(Elf32_Ehdr)) != sizeof(Elf32_Ehdr)) {
		fprintf(stdout, "\t[-] Unable to read %d bytes from file, %s, most likely not an ELF file\n", sizeof(Elf32_Ehdr), strerror(errno));
		exit(EXIT_FAILURE);
	}

	if(memcmp(elf_header.e_ident, ELFMAG, SELFMAG) != 0) {
		fprintf(stdout, "\t[-] File doesn't appear to be an ELF binary, bailing\n");
		exit(EXIT_FAILURE);
	}

	if(elf_header.e_machine != EM_386 || elf_header.e_type != ET_EXEC || elf_header.e_version != EV_CURRENT ||
		elf_header.e_ehsize != sizeof(Elf32_Ehdr) || elf_header.e_phentsize != sizeof(Elf32_Phdr) || 
		elf_header.e_phnum > 8) {
		fprintf(stdout, "\t[-] Binary layout isn't what we expected, bailing\n");
		exit(EXIT_FAILURE);
	}
	
	prog_headers = malloc(elf_header.e_phnum * sizeof(Elf32_Phdr));
	if(prog_headers == NULL) {
		fprintf(stdout, "\t[-] Unable to malloc memory needed for program headers, bailing\n");
		exit(EXIT_FAILURE);
	}

	if(lseek(fd, elf_header.e_phoff, SEEK_SET) != elf_header.e_phoff) {
		fprintf(stdout, "\t[-] Unable to lseek to offset, %s\n", strerror(errno));
		exit(EXIT_FAILURE);
	}

	if(read(fd, prog_headers, elf_header.e_phnum * sizeof(Elf32_Phdr)) != elf_header.e_phnum * sizeof(Elf32_Phdr)) {
		fprintf(stdout, "\t[-] Unable to read in program headers: %s\n\t[-] bailing\n", strerror(errno));
		exit(EXIT_FAILURE);
	}
	fprintf(stdout, "\t[*] %d program headers\n", elf_header.e_phnum);
	
	
	fprintf(stdout, "[*] Scanning through program headers\n");

	min = prog_headers[0].p_vaddr;
	offset = 0;
	
	for(i = 0; i < elf_header.e_phnum; i++) {
		fprintf(stdout, "\t[*] Entry %d: %sload entry\n", i, prog_headers[i].p_type == PT_LOAD?"":"non-");
		if(prog_headers[i].p_type != PT_LOAD) continue;
		
		prot = 0;
		if(prog_headers[i].p_flags & PF_R) prot |= PROT_READ;
		if(prog_headers[i].p_flags & PF_W) prot |= PROT_WRITE;
		if(prog_headers[i].p_flags & PF_X) prot |= PROT_EXEC;

		if(mmap(min, prog_headers[i].p_filesz, PROT_READ|PROT_WRITE, MAP_SHARED|MAP_FIXED, fd, 0) != min) {
			printf("Unable to map program @ %p\n", min);
			exit(EXIT_FAILURE);
		}
	
		memstuff[memcount].what = min;
		memstuff[memcount++].len = prog_headers[i].p_filesz;
		
		min += prog_headers[i].p_filesz;
		min += PAGESIZE;
		min &= ~(PAGESIZE - 1);

		offset += prog_headers[i].p_filesz;
	}
	
}

void quikdis_reporter( enum x86_report_codes code, void *arg ) {
	char * str = NULL;

	/* here would could examine the error and do something useful;
	* instead we just print that an error occurred */
	switch ( code ) {
		case report_disasm_bounds:
			str = "Attempt to disassemble RVA beyond end of buffer";
			break;
		case report_insn_bounds:
			str = "Instruction at RVA extends beyond buffer";
			break;
		case report_invalid_insn:
			str = "Invalid opcode at RVA";
			break;
		case report_unknown:
			str = "Unknown Error";
			break;
	}

	fprintf(stdout, "QUIKDIS: ERROR \'%s:\' %p\n", str, arg);
}


unsigned char *stack[4096];
int stack_count = 0;

void stack_push(unsigned char *wot)
{
	if(stack_count > 4095) {
		printf("stack_push(): stack exhausted :< \n");
		exit(EXIT_FAILURE);
	}
	stack[stack_count++] = wot;
}

unsigned char* stack_pop()
{
	if(stack_count == 0) {
		printf("stack_pop(): stack exhausted :< \n");
		exit(EXIT_FAILURE);
	}
	return stack[--stack_count];
}

unsigned char *jumps[4096];
int jump_count = 0;

void add_jump(unsigned char *where)
{
	if(jump_count > 4095) {
		printf("known eip change locations exhausted\n");
		exit(EXIT_FAILURE);
	}
	jumps[jump_count++] = where;
}

int seen_jump(unsigned char *where)
{
	int i;
	for(i = 0; i < jump_count; i++) {
		if(jumps[i] == where) return 1;
	}

	return 0;
}


int main(int argc, char **argv)
{
	unsigned char *p;
	struct stat statbuf;
	int fd,i, sz;
	
	x86_insn_t insn;
	char instruction[1024];
	char reloffset;
	
	setvbuf(stdout, NULL, _IONBF, 0);
	
	printf("analyse.c - andrewg@felinemenace.org\n");
	
	if(argc == 1) {
		printf("Please specify a filename to open\n");
		exit(EXIT_FAILURE);
	}
	
	disassemble_init(0, INTEL_SYNTAX);

	argv++;
	while(argv[0]) {
		printf("- Opening %s\n", argv[0]);
		
		fd = open(argv[0], O_RDWR);
		if(fd == -1) {
			perror("open");
			exit(EXIT_FAILURE);
		}
		
		if(fstat(fd, &statbuf) == -1) {
			perror("fstat");
			exit(EXIT_FAILURE);
		}
	
		load_elf(fd);
		
		printf("\t- Entry point: %p\n", elf_header.e_entry);
	
		p = (unsigned char *)(elf_header.e_entry);
		while(1) {	
			memset(&insn, 0, sizeof(x86_insn_t));
			//printf("Disassembling @ address: %p\n", p);
			
			sz = x86_disasm(prog_headers[0].p_vaddr, statbuf.st_size, prog_headers[0].p_vaddr, p - prog_headers[0].p_vaddr, &insn);
			if(sz) {
				x86_format_insn(&insn, instruction, sizeof(instruction), intel_syntax);

				if(insn.group == insn_controlflow && (insn.type >> 12) == insn_controlflow) {
					printf("%p\t%s\n", p, instruction);

					if(seen_jump(p)) {
						printf("\tI've reached this piece of code, heading backwards..\n");
						p = stack_pop();
						continue;
					}
					
					add_jump(p);
					
					/*for(i = 0; i < 3; i++) {
						if(insn.operands[i].type == op_unused) {
							//printf("%d - Unused operand.\n", i);
							continue;
						}
						if(insn.operands[i].type == op_register) {
							printf("%d - register %d\n", i, insn.operands[i].data.reg.name);
							continue;
						}*/
					i = 0;

					if(insn.type == insn_return) {
						printf("\tReturn: eip currently %p\n", p);
						p = stack_pop();
						printf("\tEIP now: %p\n", p);
						continue;
					}

					if(insn.type != insn_jmp) stack_push(p + sz);
					
					printf("insn.flags_tested: %08x\n", insn.flags_tested);
					//printf("insn.flags_set: %08x\n", insn.flags_set);
					
					if(insn.operands[0].type == op_register) { 
						if(insn.type == insn_jmp) {
							printf("\tjumping to a register, %s\n", insn.operands[0].data.reg.name);
							p = stack_pop();
						} else {
							printf("\tCalling some register, continuing after call\n");
							p += sz;
							continue;
						}
					}	
					
					if(insn.operands[0].type == op_expression) {
						printf("\tJumping to an expression\n");
						if(insn.type == insn_jmp) {
							printf("\treturning to previous way of getting here\n");
							p = stack_pop();
							continue;
						} else {
							printf("\tcontining afterwards\n");
							p += sz;
							continue;
						}
					}
					
					if(insn.operands[0].type == op_relative) {
						printf("%d - relative offset\n", i);
						if(insn.operands[0].datatype == op_byte) {
							printf("\t- single byte offset\n");
							printf("\t- modified ptr: %p\n", (p+sz) + insn.operands[0].data.sbyte);
							output(p, (p+sz)+insn.operands[0].data.sbyte, &insn);
							p = (p+sz) + insn.operands[0].data.sbyte;
							continue;
						} else if(insn.operands[0].datatype == op_dword) {
							printf("\t- long offset\n");
							printf("\t- modified ptr: %p\n", (p+sz) + (long) insn.operands[0].data.sdword);
							output(p, (p+sz)+insn.operands[0].data.sdword, &insn);
							p = (p+sz) + insn.operands[0].data.sdword;
							continue;
						}
						
						printf("Unknown offset: %d :<\n", insn.operands[0].datatype);
						p += sz;
						continue;
					}
		
					printf("UNHANDLED: operand type: %d\n", insn.operands[0].type);
					p += sz;
					printf("!!");
					sleep(1);
					/*
					switch(ins n.type) {
						case insn_return: p = stack_pop(); continue; break;
					*/
				} 
				if(insn.type == insn_halt) {
					printf("Got halt statement, this usually means we've reached the end..\n");
					break;
				}
				p += sz;
			} else {
				printf("Invalid opcode :<\n");
				exit(EXIT_FAILURE);
				p++;
			}
		}
		patch_binary();
		close(fd);
		argv++;
	}

	
}