#include <unicorn/unicorn.h>
#include <capstone/capstone.h>

#include <sys/types.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <stdbool.h>
#include <assert.h>

#define PAGE_SIZE 4 * 1024 * 1024
#define CODE_START 0x1000000
#define CODE_SIZE PAGE_SIZE
#define STACK_SIZE PAGE_SIZE
#define STACK_START CODE_START + CODE_SIZE + STACK_SIZE

#define EIP_START CODE_ADDRESS
#define ESP_START STACK_START
#define EBP_START ESP_START

static void dump_regs( uc_engine *uc )
{
	uint32_t eip;
	uint32_t esp;
	uint32_t ebp;
	uint32_t eax;
	uint32_t ebx;
	uint32_t ecx;
	uint32_t edx;
	uint32_t esi;
	uint32_t edi;
	
	uc_reg_read( uc, UC_X86_REG_EIP, &eip );
	uc_reg_read( uc, UC_X86_REG_ESP, &esp );
	uc_reg_read( uc, UC_X86_REG_EBP, &ebp );
	uc_reg_read( uc, UC_X86_REG_EAX, &eax );
	uc_reg_read( uc, UC_X86_REG_EBX, &ebx );
	uc_reg_read( uc, UC_X86_REG_ECX, &ecx );
	uc_reg_read( uc, UC_X86_REG_EDX, &edx );
	uc_reg_read( uc, UC_X86_REG_ESI, &esi );
	uc_reg_read( uc, UC_X86_REG_EDI, &edi );

	printf( "EIP: %08X\n", eip );
	printf( "ESP: %08X\n", esp );
	printf( "EBP: %08X\n", ebp );
	printf( "EAX: %08X\n", eax );
	printf( "EBX: %08X\n", ebx );
	printf( "ECX: %08X\n", ecx );
	printf( "EDX: %08X\n", edx );
	printf( "ESI: %08X\n", esi );
	printf( "EDI: %08X\n", edi );
}

static void dump_stack( uc_engine *uc )
{
	uint32_t esp;
	uint8_t mem[4];
		
	uc_reg_read( uc, UC_X86_REG_ESP, &esp );
	
	printf( "stack:\n" );
	for( int i = esp; i < STACK_START; i += 4 ) {
		uc_mem_read( uc, i, &mem, 4 );
		printf( "%08X: %02X%02X%02X%02X\n", i, mem[3], mem[2], mem[1], mem[0] ); 
	}
}

static void dump_memory( uc_engine *uc, uint64_t start_address, uint64_t end_address )
{
	uint8_t mem[4];

	printf( "data:\n" );
	for( uint64_t a = start_address; a < end_address; a += 4 ) {
		uc_mem_read( uc, a, &mem, 4 );
		printf( "%08X: %02X%02X%02X%02X\n", (uint32_t)a, mem[3], mem[2], mem[1], mem[0] ); 
	}
}

static uint32_t mul_hash( uint64_t x, int p )
{
	uint32_t v = x * INT32_C( 2654435761 );
	return v >> ( 32 - p );
}

static int compute_p( int size )
{
	int p = 0;
	while( size > 0 ) {
		size >>= 1;
		p++;
	}
	return p;
}

int main( int argc, char *argv[] )
{
	uc_engine *uc;
	uc_err uerr;
	char *code;
	size_t code_size;
	size_t file_size;
	FILE *f;
	int res;
	csh cs;
	cs_err cerr;
	size_t nof_instrs;
	cs_insn *instrs;
	uint64_t data_start = 0;

	uerr = uc_open( UC_ARCH_X86, UC_MODE_32, &uc );
	if( uerr != UC_ERR_OK ) {
		fprintf( stderr, "ERROR: failed to call uc_open( ): %s\n", uc_strerror( uerr ) );
		exit( EXIT_FAILURE );
	}
	
	cerr = cs_open( CS_ARCH_X86, CS_MODE_32, &cs );
	if( cerr != CS_ERR_OK ) {
		uc_close( uc );
		fprintf( stderr, "ERROR: failed to call uc_open( ): %s\n", uc_strerror( uerr ) );
		exit( EXIT_FAILURE );
	}
	
	uerr = uc_mem_map( uc, CODE_START, CODE_SIZE, UC_PROT_ALL );
	if( uerr != UC_ERR_OK ) {
		fprintf( stderr, "ERROR: failed to call uc_mem_map( ) for code memory: %s\n", uc_strerror( uerr ) );		
		uc_close( uc );
		exit( EXIT_FAILURE );
	}

	uerr = uc_mem_map( uc, STACK_START - STACK_SIZE, STACK_SIZE, UC_PROT_ALL );
	if( uerr != UC_ERR_OK ) {
		fprintf( stderr, "ERROR: failed to call uc_mem_map( ) for stack memory: %s\n", uc_strerror( uerr ) );		
		uc_close( uc );
		exit( EXIT_FAILURE );
	}
	
	f = fopen( argv[1], "r" );
	if( f == NULL ) {
		fprintf( stderr, "ERROR: unable to read file '%s': %s\n", argv[1], strerror( errno ) );
		uc_close( uc );
		exit( EXIT_FAILURE );
	}
	
	res = fseek( f, 0, SEEK_END );
	if( res != 0 ) {
		fprintf( stderr, "ERROR: unable to seek to end of file '%s': %s\n", argv[1], strerror( errno ) );
		uc_close( uc );
		exit( EXIT_FAILURE );
	}

	file_size = ftell( f );
	if( file_size > CODE_SIZE ) {
		fprintf( stderr, "ERROR: we should really not run the emulator with a code segment bigger than %d\n", CODE_SIZE );
		uc_close( uc );
		exit( EXIT_FAILURE );
	}
	
	code = (char *)malloc( file_size );
	if( code == NULL ) {
		fprintf( stderr, "ERROR: out of memory\n" );
		uc_close( uc );
		exit( EXIT_FAILURE );
	}
		
	fseek( f, 0, SEEK_SET );
	if( res != 0 ) {
		fprintf( stderr, "ERROR: unable to seek to start of file '%s': %s\n", argv[1], strerror( errno ) );
		uc_close( uc );
		exit( EXIT_FAILURE );
	}

	code_size = fread( code, 1, file_size, f );
	if( code_size != file_size ) {
		fprintf( stderr, "ERROR: could not read complete code file, read only %zu bytes instead of %zu bytes\n",
			code_size, file_size );
		uc_close( uc );
		exit( EXIT_FAILURE );
	}
	printf( "Read %zu bytes of code..\n", code_size );

	// disassemble the whole code block	
	nof_instrs = cs_disasm( cs, (const uint8_t *)code, code_size, CODE_START, 0, &instrs );
	if( nof_instrs == 0 ) {
		fprintf( stderr, "ERROR: failed to call cs_disasm( ): %s\n", cs_strerror( cs_errno( cs ) ) );		
		cs_close( &cs );
		uc_close( uc );
		exit( EXIT_FAILURE );
	}

	// print all the disassembled code
	for( int i = 0; i < nof_instrs; i++ ) {
		printf( "%04X: ", (uint32_t)instrs[i].address );
		for( int j = 0; j < instrs[i].size; j++ ) {
			printf( "%02X", instrs[i].bytes[j] );
		}
		for( int j = ( 16 - instrs[i].size ) * 2; j > 0; j-- ) {
			printf( " " );
		}
		if( instrs[i].size == 2 && instrs[i].bytes[0] == 0 && instrs[i].bytes[1] == 0 ) {
			if( data_start == 0 ) {
				data_start = instrs[i].address;
			}
			printf( "data\n" );
		} else {
			printf( "%s %s\n", instrs[i].mnemonic, instrs[i].op_str );
		}
	}
	
	// remember address to instrs indexes so we can get the current
	// opcode when reaching a certain EIP address
	int N = nof_instrs * 2;
	int p = compute_p( N );
	N = ( 1 << p );
	int *instrs_map = calloc( N, sizeof( int ) * 2 );
	for( int i = 0; i < nof_instrs; i++ ) {
		int n = mul_hash( instrs[i].address, p ) * 2;
		assert( n < 2 * N );
		while( instrs_map[n+1] != 0 ) {
			n += 2;
			if( n >= 2 * N ) {
				n = 0;
			}
		}
		instrs_map[n] = i;
		instrs_map[n+1] = instrs[i].address;
	}

	// verify the EIP to instr index map has been constructed correctly
	for( int i = 0; i < nof_instrs; i++ ) {
		uint64_t a = instrs[i].address;
		int n = mul_hash( a, p ) * 2;
		while( instrs_map[n+1] != a ) {
			n += 2;
			if( n >= 2 * N ) {
				n = 0;
			}
		}
		n = instrs_map[n];
		assert( n == i );
	}
		
	// write executable code to emulator
	uerr = uc_mem_write( uc, CODE_START, code, code_size );
	if( uerr != UC_ERR_OK ) {
		fprintf( stderr, "ERROR: failed to call uc_mem_write( ): %s\n", uc_strerror( uerr ) );		
		uc_close( uc );
		exit( EXIT_FAILURE );
	}

	int addr = CODE_START;
	int esp = STACK_START;
	uc_reg_write( uc, UC_X86_REG_ESP, &esp );
	
	uint64_t address = CODE_START;
	bool terminate = false;
	int iteration = 1;

	printf( "Single step execution:\n" );
	while( !terminate ) {

		printf( "-- iteration %d\n", iteration );
		iteration++;

		int n = mul_hash( address, p ) * 2;
		while( instrs_map[n+1] != address ) {
			n += 2;
			if( n >= 2 * N ) {
				n = 0;
			}
		}
		n = instrs_map[n];

		printf( "%04X: ", (unsigned int)address );

		for( int i = 0; i < instrs[n].size; i++ ) {
			printf( "%02X", instrs[n].bytes[i] );
		}
		for( int i = ( 16 - instrs[n].size ) * 2; i > 0; i-- ) {
			printf( " " );
		}
				
		printf( "%s %s\n", instrs[n].mnemonic, instrs[n].op_str );
		
		uerr = uc_emu_start( uc, addr, CODE_START + CODE_SIZE, 0, 1 );
		if( uerr != UC_ERR_OK ) {
			fprintf( stderr, "ERROR: failed to call uc_emu_start( ): %s\n", uc_strerror( uerr ) );		
			cs_close( &cs );
			uc_close( uc );
			exit( EXIT_FAILURE );
		}

		int eip;
		uc_reg_read( uc, UC_X86_REG_EIP, &eip );
		addr = eip;
		address = eip;
		
		dump_regs( uc );
		dump_stack( uc );
		//dump_memory( uc, data_start, CODE_START + code_size );
		
		if( strcmp( instrs[n].mnemonic, "hlt" ) == 0 ) {
			terminate = true;
		}
	}

	free( instrs_map );
	cs_free( instrs, nof_instrs );
	cs_close( &cs );
	uc_close( uc );
	
	exit( EXIT_SUCCESS );
}