#include <iostream>
#include <windows.h>
#include "unjunk.h"

extern "C" {
	#include "libdis.h"
	#include "ia32_list.h"
}

static struct DIS_INFO {
	/* file streams */
	FILE *in, *out, *err;
	/* size of input file */
	unsigned long size;
	/* flag for --entry option */
	int entry;
	/* pagesize for buffering STDIN */
	unsigned long pagesize;
	/* output syntax */
	enum x86_asm_format syntax;
} info = { NULL, NULL, NULL, 0, 0, 1, /*raw_syntax*/ intel_syntax };


bool force_quit;
long _offset_start;
long _offset_new;
long _map_old;
long _map_new;
long _current;
long _start;
long _img_base;
int op_number;



struct RVALIST {
	unsigned long rva;
	struct RVALIST *next;
} rva_list_head = {0},rva_list_all_head ={0};

void rva_clear(void)
{
	struct RVALIST *rl,*rl_new;
	rl=rva_list_head.next;
	while(rl)
	{
		rl_new=rl;
		rl=rl->next;
		free(rl_new);
	}
	rva_list_head.next=0;
}
int rva_list_add( unsigned long rva ) {
	struct RVALIST *rl, *rl_new;

	for ( rl = &rva_list_head; rl; rl = rl->next ) {
		/* first rva is always 0 -- the list head */
		if ( rva > rl->rva ) {
			if ( ! rl->next || rva < rl->next->rva ) {
				/* we use exit() to free this, btw */
				rl_new = (RVALIST*)calloc(sizeof(struct RVALIST), 1);
				rl_new->rva = rva;
				rl_new->next = rl->next;
				rl->next = rl_new;
				return(1);
			}
		} else if ( rva == rl->rva ) {
			printf("already seen %X\n",rva);
			return(0);	/* already seen this rva */
		}
	}
	printf("Return 0\n");
	return(0);
}

int rva_list_add_full( unsigned long rva ) {
	struct RVALIST *rl, *rl_new;

	if( rva < 0x00402FFF)
		return 0;

	for ( rl = &rva_list_all_head; rl; rl = rl->next ) {
		/* first rva is always 0 -- the list head */
		if ( rva > rl->rva ) {
			if ( ! rl->next || rva < rl->next->rva ) {
				/* we use exit() to free this, btw */
				rl_new = (RVALIST*)calloc(sizeof(struct RVALIST), 1);
				rl_new->rva = rva;
				rl_new->next = rl->next;
				rl->next = rl_new;
				return(1);
			}
		} else if ( rva == rl->rva ) {
			printf("already seen in global list%X\n",rva);
			return(0);	/* already seen this rva */
		}
	}
	printf("Return 0\n");
	return(0);
}

void x86dis_manual_print( x86_insn_t *insn, void *arg ) {
	char line[4096];
	int i;
	
	//if ( x86_format_insn(insn, line, 4096, info.syntax) <= 0 ) {
	//	return;
	//}

	//insn->offset-3000 = position dans le fichier!

	//if(insn->type == insn_jcc)
	//{
	//	//rva_list_add_full(insn->operands[0].op.data.relative_far + insn->addr);
	//	//printf("%d\n",insn->operands[0].op.type);
	//
	//	if(insn->operands[0].op.type == op_relative_far)
	//	{
	//		//printf("jcc found [%X] %X\n",insn->addr,insn->operands[0].op.data.relative_far + insn->addr + insn->size);
	//		disasm_list_add(insn->operands[0].op.data.relative_far + insn->addr + insn->size -0x5000 - 0x400000);
	//	}
	//		
	//	
	//	if(insn->operands[0].op.type == op_relative_near)
	//	{
	//		//printf("jcc found [%X] %X\n",insn->addr,insn->operands[0].op.data.relative_near + insn->addr + insn->size);
 //           disasm_list_add(insn->operands[0].op.data.relative_near + insn->addr + insn->size -0x5000 - 0x400000);
	//	}
	//		
	//}

	if(insn->type == insn_return )
	{
		//printf("force_quit = true\n");
		force_quit = true;
	}

	if(insn->type == insn_jmp)
	{
		line[0]='\0';
		return;
	}
	else
	{
		memcpy((void *)_current,insn->bytes,insn->size);
		int old_cur = (int)_current;
		_current=_current + insn->size;

		if(insn->group == insn_controlflow && (insn->type == insn_call ||insn->type == insn_jcc ) && insn->operands[0].op.type == op_relative_far)
		{
			int old_addr =insn->addr + insn->operands[0].op.data.relative_far + insn->size;
			int new_addr = (int)(_current) - _start + _offset_new + _img_base ;
			*( (long *)( _current-4) ) = old_addr - new_addr;
		}
	}

	//if ( info.syntax == att_syntax || info.syntax == intel_syntax ) {
	//	/* print an address and hex bytes, since libdisasm doesn't */
	//	printf("%08lX", insn->addr );
	//	for ( i = 0; i < 10; i++ ) {
	//		if ( i < insn->size ) {
	//			printf(" %02X", insn->bytes[i]);
	//		} else {
	//			printf("   ");
	//		}
	//	}
	//	printf("\t");
	//}

	//printf("%s\n", line);
}


long x86dis_resolver( x86_op_t *op, x86_insn_t *insn ) {
	long retval = -1;

	if (! rva_list_add(insn->addr) ) {
		/* we have seen this one already; return -1 */
		return(-1);
	}
	if(force_quit)
	{
		return -1;
	}

	/* this part is a flat ripoff of internal_resolver in libdis.c */
	/* we don't do any register or stack resolving */
	if ( op->type == op_absolute || op->type == op_offset ) {
		retval = op->data.sdword; /* no need to cast the void* */
	} else if (op->type == op_relative_near ){
		retval = insn->addr + insn->size + op->data.sbyte;
	} else if (op->type == op_relative_far ){
		if ( op->datatype == op_word ) {
			retval = insn->addr + insn->size + op->data.sword;
		} else if ( op->datatype == op_dword ) {
			retval = insn->addr + insn->size + op->data.sdword;
		}
	}

	return( retval );
}



int replace_func(long func_addr,long func_new) /*0x08C3000*/
{
	op_number=0;
	_offset_start	= func_addr - _img_base;
	_offset_new		= func_new - _img_base;
	force_quit = false;
	long _old_current;

	rva_clear();
	if(!rva_list_add_full(func_addr))
		return 0;

	if(!_current)
	{
		//		         
		_current = _map_new + _offset_new -0x3600;
	}
	_start = _current;
	_old_current= _current;
	

	/*
	*
	*	Put Jump to new function
	*/
	*(char *)( _map_new + _offset_start - 0x2000) = 0xE9;
	*(int *)(_map_new + 1 + _offset_start - 0x2000) = _offset_new - _offset_start -5;

	x86_disasm_forward( (unsigned char *) _map_old + 0x3000, /*BUF*/
						0x454200, /*Size*/
						(unsigned long) 0x0405000, /*RVA of section*/
						_offset_start - 0x5000, /*Adress in section*/
						x86dis_manual_print,
						0,
						x86dis_resolver
						);

	*((long *)_current)=0x90909090;
	_current+=4;
    return _current - _old_current;

}

void init(long map_old,long map_new,long img_base)
{
	_map_old = map_old;
	_map_new = map_new;
	_img_base = img_base;
}




int main(int argc,char **argv)
{
	x86_init(opt_none, NULL);
	//Dont overwrite if exists!
	CopyFile("Meloquynthe_puls.exe","Meloquynthe_new2.exe",true);
	HANDLE source,dest,source_m,dest_m;	
	LPVOID psource,pdest;

	source=CreateFile("Meloquynthe.exe",
		GENERIC_READ | GENERIC_WRITE,
		FILE_SHARE_READ | FILE_SHARE_WRITE,
		0,
		OPEN_EXISTING,
		0,
		0);

	dest=CreateFile("Meloquynthe_new2.exe",
		GENERIC_READ | GENERIC_WRITE,
		FILE_SHARE_READ | FILE_SHARE_WRITE,
		0,
		OPEN_EXISTING,
		0,
		0);

	source_m=CreateFileMapping(source,0,
		PAGE_READONLY,
		0,
		0,
		0);

	dest_m=CreateFileMapping(dest,0,
		PAGE_READWRITE,
		0,
		0,
		0);

	psource = MapViewOfFile(source_m,
		FILE_MAP_READ,
		0,0,0);

	pdest = MapViewOfFile(dest_m,
		FILE_MAP_ALL_ACCESS,
		0,0,0);


	init((long)psource,(long)pdest,0x400000);

	long number;
	long whereto=0x08C3000;

	while(1)
	{
		std::cout << "Entre la fonction à remplacer" << std::endl;
		std::cin >> std::hex >> number;
		std::cout << "On remplace la function " << std::hex << number << std::endl;
		if(!number)
			break;
        whereto+=replace_func(number,whereto);

		int num=0;
		while(num=disasm_list_pop())
		{
			printf("Now Replacing %X \n",num+0x5000 + 0x400000);
			whereto+=replace_func(num+0x5000 + 0x400000,whereto);
		}
	}
	


	x86_cleanup();
	UnmapViewOfFile(psource);
	UnmapViewOfFile(pdest);
	CloseHandle(source_m);
	CloseHandle(dest_m);
	CloseHandle(source);
	CloseHandle(dest);
}