#include <os.h>
#include "dirlist.h"
#include "browse.h"
#include "types.h"
#include "keys.h"
#include "screen.h"
#include "charmaps.h"
#include "config.h"

#define PAYLOAD_SIZE 0x40000

#define NOP_SIZE 0x1F04 
/* 
To find this:
Feed boot1.5 an image of all As with base address 0x11200000. 
It will crash with "undefined instruction 41414141". 
This value is (crashed_PC - 0x11200000).	
*/

#define TRAMPOLINE_SIZE 8


#define SWAPENDIAN(n)((((n)&0x000000ff)<<24)|(((n)&0x0000ff00)<<8)|(((n)&0x00ff0000)>>8)|(((n)&0xff000000)>>24))

uint16_t* offscreen;
uint16_t* screen;

void clearScroll(uint16_t* buf) {
	setCurColorRGB(255,255,255);
	drawBufFullBox(buf,SCREEN_WIDTH-1-SCROLL_WIDTH,HEAD_HEIGHT2,SCREEN_WIDTH-1,SCREEN_HEIGHT);
	setCurColorRGB(0,0,0);
}

int main(int argc, char** argv) {
	uint8_t *buffer=0, *ptr;
	FILE *boot2=0, *nano=0, *out=0;
	char path_nano[PATH_SIZE], path_boot2[PATH_SIZE], path_init[PATH_SIZE], path_out[2*PATH_SIZE], msg[2*PATH_SIZE], tmp[64];
	uint32_t boot2_size=0, boot2_rounded_size=0, nano_size=0, total_size=0, t,t2,t3;
	char *rootpath =get_documents_dir();
	uint16_t i,r=0;
	*msg=0;
	initScreen();
	startScreen();
	strcpy(path_init,argv[0]);
	*(strrchr(path_init,'/')+1)=0;
	uint16_t* baseoffscreen = (uint16_t*) malloc(SCREEN_SIZE);
	if(!baseoffscreen)
		strcpy(msg,"Error allocating screen buffer");
	if(!*msg) {
		offscreen=baseoffscreen;
		memset(getScreen(),0xFF,SCREEN_SIZE);
		
		setCurColorRGB(0xFF,0,0xFF);
		strcpy(tmp,"PG-1");
		drwBufStr(getScreen(),SCREEN_WIDTH-2-strlen(tmp)*CHAR_WIDTH,0,"PG-1",0,1);
		setCurColorRGB(0xFF,0,0);
		drwBufStr(getScreen(),SCREEN_WIDTH-1-strlen(tmp)*CHAR_WIDTH,0,"PG-1",0,1);
		
		setCurColorRGB(0xFF,0,0xFF);
		strcpy(tmp,"(1/2) Choose payload for Boot1.5 exploit :");
		drwBufStr(getScreen(),1,0,tmp,0,1);
		setCurColorRGB(0,0,0xFF);
		drwBufStr(getScreen(),0,0,tmp,0,1);
		setCurColorRGB(0,0,0);
		r=chooseFile(path_nano,path_init,rootpath,NULL);
	}
	if(r && !*msg) {
		nano=fopen(path_nano,"rb");
		if(!nano) {
			strcpy(msg,"Error opening payload image for reading:\n");
			strcat(msg,path_nano);
		}
	}
	if(r && !*msg) {
		fseek(nano,0,SEEK_END);
		nano_size=ftell(nano);
		fseek(nano,0,SEEK_SET);
		if(nano_size>PAYLOAD_SIZE)
			strcpy(msg,"Error payload is too big");
	}
	if(r && !*msg) {
		strcpy(path_init,path_nano);
		*(strrchr(path_init,'/')+1)=0;
		
		// clear prompt
		setCurColorRGB(0,0,0);
		memset(tmp,' ', strlen(tmp));
		drwBufStr(getScreen(),0,0,tmp,0,0);

		setCurColorRGB(0xFF,0,0xFF);
		strcpy(tmp,"(2/2) Choose Boot2 image to use it with :");
		drwBufStr(getScreen(),1,0,tmp,0,1);
		setCurColorRGB(0,0,0xFF);
		drwBufStr(getScreen(),0,0,tmp,0,1);
		setCurColorRGB(0,0,0);
		memset(((uint16_t*)getScreen())+HEAD_HEIGHT1*SCREEN_WIDTH,0xFF,SCREEN_SIZE-HEAD_HEIGHT1*SCREEN_WIDTH*2);
		r=chooseFile(path_boot2,path_init,rootpath,NULL);
	}
	if(r && !*msg) {
		boot2=fopen(path_boot2,"rb");
		if(!boot2) {
			strcpy(msg,"Error opening boot2 image for reading:\n");
			strcat(msg,path_boot2);
		}
	}
	if(r && !*msg) {
		fseek(boot2,0,SEEK_END);
		boot2_size=ftell(boot2);
		fseek(boot2,0,SEEK_SET);
		strcpy(path_out,argv[0]);
		*(strrchr(path_out,'/')+1)=0;
		t=strlen(path_out);
		strcat(path_out,strrchr(path_nano,'/')+1);
		ptr=strchr(path_out+t,'.');
		if(!ptr) ptr=path_out+strlen(path_out);
		*ptr=0;
		strcat(path_out,"+");
		t=strlen(path_out);
		strcat(path_out,strrchr(path_boot2,'/')+1);
		ptr=strchr(path_out+t,'.');
		*(ptr+1)=0;
		if(!ptr) ptr=path_out+strlen(path_out);
		strcat(path_out,"img.tns");
		out=fopen(path_out,"wb");
		if(!out) {
			strcpy(msg,"Error opening image for writing:\n");
			strcat(msg,path_out);
		}
	}
	if(r && !*msg) {
		boot2_rounded_size=(boot2_size+3)/4*4;
		total_size=boot2_rounded_size+PAYLOAD_SIZE+NOP_SIZE+TRAMPOLINE_SIZE;
		buffer=malloc(total_size);
		if(!buffer) strcpy(msg,"Error allocating image buffer");
	}
	if(r && !*msg) {
		ptr=buffer;
		memset(ptr,0,total_size);
		fread(ptr,boot2_size,1,boot2);
		ptr+=boot2_rounded_size;
		fread(ptr,nano_size,1,nano);
		ptr+=PAYLOAD_SIZE;
		t2=0x111c0000-boot2_rounded_size;
		*(uint32_t*)(ptr-8)=t2;
		*(uint32_t*)(ptr-4)=boot2_rounded_size;
		for(t=0;t<NOP_SIZE/4;t++) {
			memcpy(ptr,"\x00\x00\xA0\xE1",4);
			ptr+=4;
		}
		/* ldr pc, [pc, #-4] / .word 0x11C00000 */
		memcpy(ptr,"\x04\xF0\x1F\xE5\x00\x00\x1C\x11",TRAMPOLINE_SIZE);
		t=SWAPENDIAN(total_size+32);
		t2=SWAPENDIAN(t2);
		t3=SWAPENDIAN(total_size);
		if(	!fwrite("\x80\x0F",2,1,out) ||
			!fwrite(&t,4,1,out) ||
			/* compatible '0F' & '10'; base address 11200000 */
			!fwrite("\x80\x00\x80\x44\x6E\x4C\x64\x72\x80\x13\x35\x30\x46\x80\x13\x35\x31\x30\x80\x84",20,1,out) ||
			!fwrite(&t2,4,1,out) ||
			!fwrite("\x80\x7F",2,1,out) ||
			!fwrite(&t3,4,1,out) ||
			!fwrite(buffer,total_size,1,out) ||
			!fwrite("\x02\xF0\xFF\xF0",4,1,out)) {
			strcpy(msg,"Error writing image:\n");
			strcat(msg,path_out);
		}
	}
	if(out) fclose(out);
	if(nano) fclose(nano);
	if(boot2) fclose(boot2);
	stopScreen();
	free(baseoffscreen);
	if(buffer) free(buffer);
	if(r) {
		char* tmp2="Error";
		if(!*msg) {
			strcpy(msg,"Success creating image:\n");
			strcat(msg,path_out);
			strcat(msg,"\nPress [home][2] to refresh.");
			tmp2="Success";
		}
		show_msgbox(tmp2, msg);
	}
	return 0;
}