#import "kernel_utils.h"
#import "patchfinder64.h"

static mach_port_t tfpzero;
static uint64_t kernel_base;
static uint64_t KASLR_Slide;

#import <Foundation/Foundation.h>
#define LOG(str, args...) do { NSLog(@"[*] " str "\n", ##args); } while(0)

void init_kernel_utils(mach_port_t tfp0, uint64_t kbase) {
    tfpzero = tfp0;
    kernel_base = kbase;
    KASLR_Slide = (uint32_t)(kernel_base - 0xFFFFFFF007004000); // slid kernel base - kernel base = kaslr slide
}

uint64_t get_kernel_slide() {
    return KASLR_Slide;
}

uint64_t Kernel_alloc(vm_size_t size) {
    mach_vm_address_t address = 0;
    mach_vm_allocate(tfpzero, (mach_vm_address_t *)&address, size, VM_FLAGS_ANYWHERE);
    return address;
}

void Kernel_free(mach_vm_address_t address, vm_size_t size) {
    mach_vm_deallocate(tfpzero, address, size);
}

int Kernel_strcmp(uint64_t kstr, const char* str) {
    // XXX be safer, dont just assume you wont cause any
    // page faults by this
    size_t len = strlen(str) + 1;
    char *local = malloc(len + 1);
    local[len] = '\0';
    
    int ret = 1;
    
    if (KernelRead(kstr, local, len) == len) {
        ret = strcmp(local, str);
    }
    
    free(local);
    
    return ret;
}

size_t KernelRead(uint64_t where, void *p, size_t size) {
    int rv;
    size_t offset = 0;
    while (offset < size) {
        mach_vm_size_t sz, chunk = 2048;
        if (chunk > size - offset) {
            chunk = size - offset;
        }
        rv = mach_vm_read_overwrite(tfpzero, where + offset, chunk, (mach_vm_address_t)p + offset, &sz);
        if (rv || sz == 0) {
            printf("[-] error on KernelRead(0x%016llx)\n", where);
            break;
        }
        offset += sz;
    }
    return offset;
}

uint32_t KernelRead_32bits(uint64_t where) {
    uint32_t out;
    KernelRead(where, &out, sizeof(uint32_t));
    return out;
}

uint64_t KernelRead_64bits(uint64_t where) {
    uint64_t out;
    KernelRead(where, &out, sizeof(uint64_t));
    return out;
}

size_t KernelWrite(uint64_t where, const void *p, size_t size) {
    int rv;
    size_t offset = 0;
    while (offset < size) {
        size_t chunk = 2048;
        if (chunk > size - offset) {
            chunk = size - offset;
        }
        rv = mach_vm_write(tfpzero, where + offset, (mach_vm_offset_t)p + offset, chunk);
        if (rv) {
            printf("[-] error on KernelWrite(0x%016llx)\n", where);
            break;
        }
        offset += chunk;
    }
    return offset;
}

void KernelWrite_32bits(uint64_t where, uint32_t what) {
    uint32_t _what = what;
    KernelWrite(where, &_what, sizeof(uint32_t));
}


void KernelWrite_64bits(uint64_t where, uint64_t what) {
    uint64_t _what = what;
    KernelWrite(where, &_what, sizeof(uint64_t));
}

const uint64_t kernel_address_space_base = 0xffff000000000000;
void Kernel_memcpy(uint64_t dest, uint64_t src, uint32_t length) {
    if (dest >= kernel_address_space_base) {
        // copy to kernel:
        KernelWrite(dest, (void*) src, length);
    } else {
        // copy from kernel
        KernelRead(src, (void*)dest, length);
    }
}

uint64_t proc_of_pid(pid_t pid) {
    uint64_t allproc = Find_allproc();
    uint64_t proc = KernelRead_64bits(allproc), pd;

    while (proc) { //iterate over all processes till we find the one we're looking for
        pd = KernelRead_32bits(proc + 0x10);
        if (pd == pid) return proc;
        proc = KernelRead_64bits(proc);
    }
    
    return 0;
}

uint64_t ZmFixAddr(uint64_t addr) {
    static kmap_hdr_t zm_hdr = {0, 0, 0, 0};
    
    if (zm_hdr.start == 0) {
        // xxx rk64(0) ?!
        uint64_t zone_map = KernelRead_64bits(Find_zone_map_ref());
        // hdr is at offset 0x10, mutexes at start
        size_t r = KernelRead(zone_map + 0x10, &zm_hdr, sizeof(zm_hdr));
        //printf("zm_range: 0x%llx - 0x%llx (read 0x%zx, exp 0x%zx)\n", zm_hdr.start, zm_hdr.end, r, sizeof(zm_hdr));
        
        if (r != sizeof(zm_hdr) || zm_hdr.start == 0 || zm_hdr.end == 0) {
            printf("[-] KernelRead of zone_map failed!\n");
            return 1;
        }
        
        if (zm_hdr.end - zm_hdr.start > 0x100000000) {
            printf("[-] zone_map is too big, sorry.\n");
            return 1;
        }
    }
    
    uint64_t zm_tmp = (zm_hdr.start & 0xffffffff00000000) | ((addr) & 0xffffffff);
    
    return zm_tmp < zm_hdr.start ? zm_tmp + 0x100000000 : zm_tmp;
}