#import <Foundation/Foundation.h>
#import <sys/stat.h>
#import <sys/param.h>
#import "kern_utils.h"
#import "helpers/kmem.h"
#import "helpers/patchfinder64.h"
#import "helpers/kexecute.h"
#import "helpers/offsetof.h"
#import "helpers/osobject.h"
#import "sandbox.h"

#define PROC_PIDPATHINFO_MAXSIZE  (4*MAXPATHLEN)
int proc_pidpath(pid_t pid, void *buffer, uint32_t buffersize);

#define TF_PLATFORM 0x400

#define	CS_VALID		            0x0000001	/* dynamically valid */
#define CS_ADHOC		            0x0000002	/* ad hoc signed */
#define CS_GET_TASK_ALLOW	        0x0000004	/* has get-task-allow entitlement */
#define CS_INSTALLER		        0x0000008	/* has installer entitlement */

#define	CS_HARD			            0x0000100	/* don't load invalid pages */
#define	CS_KILL			            0x0000200	/* kill process if it becomes invalid */
#define CS_CHECK_EXPIRATION	        0x0000400	/* force expiration checking */
#define CS_RESTRICT		            0x0000800	/* tell dyld to treat restricted */
#define CS_ENFORCEMENT		        0x0001000	/* require enforcement */
#define CS_REQUIRE_LV		        0x0002000	/* require library validation */
#define CS_ENTITLEMENTS_VALIDATED	0x0004000

#define	CS_ALLOWED_MACHO	        0x00ffffe

#define CS_EXEC_SET_HARD	        0x0100000	/* set CS_HARD on any exec'ed process */
#define CS_EXEC_SET_KILL	        0x0200000	/* set CS_KILL on any exec'ed process */
#define CS_EXEC_SET_ENFORCEMENT	    0x0400000	/* set CS_ENFORCEMENT on any exec'ed process */
#define CS_EXEC_SET_INSTALLER	    0x0800000	/* set CS_INSTALLER on any exec'ed process */

#define CS_KILLED		            0x1000000	/* was killed by kernel for invalidity */
#define CS_DYLD_PLATFORM	        0x2000000	/* dyld used to load this is a platform binary */
#define CS_PLATFORM_BINARY	        0x4000000	/* this is a platform binary */
#define CS_PLATFORM_PATH	        0x8000000	/* platform binary by the fact of path (osx only) */

#define CS_DEBUGGED                 0x10000000  /* process is currently or has previously been debugged and allowed to run with invalid pages */
#define CS_SIGNED                   0x20000000  /* process has a signature (may have gone invalid) */
#define CS_DEV_CODE                 0x40000000  /* code is dev signed, cannot be loaded into prod signed code */

uint64_t proc_find(int pd, int tries) {
    while (tries-- > 0) {
        uint64_t proc = rk64(kernprocaddr + 0x08);
        
        while (proc) {
            uint32_t proc_pid = rk32(proc + 0x10);
            
            if (proc_pid == pd) {
                return proc;
            }
        
            proc = rk64(proc + 0x08);
        }
    }
    
    return 0;
}

char *proc_name(int pd) {
    uint64_t proc = proc_find(pd, 1);
    if (proc == 0) {
        return NULL;
    }
    
    char *proc_name = (char *)calloc(40, sizeof(char));
    
    kread(proc + 0x26c, proc_name, 40);
    
    return proc_name;
}

CACHED_FIND(uint64_t, our_task_addr) {
    uint64_t our_proc = proc_find(getpid(), 1);

    if (our_proc == 0) {
        fprintf(stderr, "failed to find our_task_addr!\n");
        exit(EXIT_FAILURE);
    }

    return rk64(our_proc + offsetof_task);
}

uint64_t find_port(mach_port_name_t port) {
    uint64_t task_addr = our_task_addr();
  
    uint64_t itk_space = rk64(task_addr + offsetof_itk_space);
  
    uint64_t is_table = rk64(itk_space + offsetof_ipc_space_is_table);
  
    uint32_t port_index = port >> 8;
    const int sizeof_ipc_entry_t = 0x18;
  
    uint64_t port_addr = rk64(is_table + (port_index * sizeof_ipc_entry_t));
    return port_addr;
}

void fixupsetuid(int pid){
    char pathbuf[PROC_PIDPATHINFO_MAXSIZE];
    bzero(pathbuf, sizeof(pathbuf));
    
    int ret = proc_pidpath(pid, pathbuf, sizeof(pathbuf));
    if (ret < 0){
        fprintf(stderr,"Unable to get path for PID %d\n", pid);
        return;
    }
    struct stat file_st;
    if (lstat(pathbuf, &file_st) == -1){
#ifdef JAILBREAKDDEBUG
        fprintf(stderr,"Unable to get stat for file %s\n", pathbuf);
#endif
        return;
    }
    if (file_st.st_mode & S_ISUID){
        uid_t fileUID = file_st.st_uid;
#ifdef JAILBREAKDDEBUG
        fprintf(stderr,"Fixing up setuid for file owned by %u\n", fileUID);
#endif
        
        uint64_t proc = proc_find(pid, 3);
        if (proc != 0) {
            uint64_t ucred = rk64(proc + offsetof_p_ucred);
            
            uid_t cr_svuid = rk32(ucred + offsetof_ucred_cr_svuid);
#ifdef JAILBREAKDDEBUG
            fprintf(stderr,"Original sv_uid: %u\n", cr_svuid);
#endif
            wk32(ucred + offsetof_ucred_cr_svuid, fileUID);
#ifdef JAILBREAKDDEBUG
            fprintf(stderr,"New sv_uid: %u\n", fileUID);
#endif
        }
    } else {
#ifdef JAILBREAKDDEBUG
        fprintf(stderr,"File %s is not setuid!\n", pathbuf);
#endif
        return;
    }
}

void set_csflags(uint64_t proc) {
    uint32_t csflags = rk32(proc + offsetof_p_csflags);
    csflags = (csflags | CS_PLATFORM_BINARY | CS_INSTALLER | CS_GET_TASK_ALLOW | CS_DEBUGGED) & ~(CS_RESTRICT | CS_HARD | CS_KILL);
    wk32(proc + offsetof_p_csflags, csflags);
}

void set_tfplatform(uint64_t proc) {
    // task.t_flags & TF_PLATFORM
    uint64_t task = rk64(proc + offsetof_task);
    uint32_t t_flags = rk32(task + offsetof_t_flags);
    fprintf(stderr,"Old t_flags: 0x%x\n", t_flags);
    t_flags |= TF_PLATFORM;
    wk32(task+offsetof_t_flags, t_flags);
    fprintf(stderr,"New t_flags: 0x%x\n", t_flags);
}

void set_csblob(uint64_t proc) {
    uint64_t textvp = rk64(proc + offsetof_p_textvp); // vnode of executable
    off_t textoff = rk64(proc + offsetof_p_textoff);
    
    if (textvp != 0){
      uint32_t vnode_type_tag = rk32(textvp + offsetof_v_type);
      uint16_t vnode_type = vnode_type_tag & 0xffff;
      uint16_t vnode_tag = (vnode_type_tag >> 16);
      
      if (vnode_type == 1) {
          uint64_t ubcinfo = rk64(textvp + offsetof_v_ubcinfo);
          
          uint64_t csblobs = rk64(ubcinfo + offsetof_ubcinfo_csblobs);
          while (csblobs != 0) {
              cpu_type_t csblob_cputype = rk32(csblobs + offsetof_csb_cputype);
              unsigned int csblob_flags = rk32(csblobs + offsetof_csb_flags);
              off_t csb_base_offset = rk64(csblobs + offsetof_csb_base_offset);
              uint64_t csb_entitlements = rk64(csblobs + offsetof_csb_entitlements_offset);
              unsigned int csb_signer_type = rk32(csblobs + offsetof_csb_signer_type);
              unsigned int csb_platform_binary = rk32(csblobs + offsetof_csb_platform_binary);
              unsigned int csb_platform_path = rk32(csblobs + offsetof_csb_platform_path);

              wk32(csblobs + offsetof_csb_platform_binary, 1);

              csb_platform_binary = rk32(csblobs + offsetof_csb_platform_binary);
              
              csblobs = rk64(csblobs);
          }
      }
    }
}


const char* abs_path_exceptions[] = {
    "/meridian",
    "/Library",
    "/private/var/mobile/Library",
    "/private/var/mnt",
    NULL
};

uint64_t get_exception_osarray(void) {
    static uint64_t cached = 0;
    
    if (cached == 0) {
        cached = OSUnserializeXML(
            "<array>"
            "<string>/meridian/</string>"
            "<string>/Library/</string>"
            "<string>/private/var/mobile/Library/</string>"
            "<string>/private/var/mnt/</string>"
            "</array>"
        );
    }

    return cached;
}

static const char *exc_key = "com.apple.security.exception.files.absolute-path.read-only";

void set_sandbox_extensions(uint64_t proc) {
    uint64_t proc_ucred = rk64(proc + 0x100);
    uint64_t sandbox = rk64(rk64(proc_ucred + 0x78) + 0x8 + 0x8);

    if (sandbox == 0) {
        fprintf(stderr, "no sandbox, skipping \n");
        return;
    }

    if (has_file_extension(sandbox, abs_path_exceptions[0])) {
        fprintf(stderr, "already has '%s', skipping \n", abs_path_exceptions[0]);
        return;
    }

    uint64_t ext = 0;
    const char** path = abs_path_exceptions;
    while (*path != NULL) {
        ext = extension_create_file(*path, ext);
        if (ext == 0) {
            fprintf(stderr, "extension_create_file(%s) failed, panic! \n", *path);
            NSLog(@"extension_create_file(%s) failed, panic!", *path);
        }
        ++path;
    }
    
    if (ext != 0) {
        extension_add(ext, sandbox, exc_key);
    }
}

void set_amfi_entitlements(uint64_t proc) {
    uint64_t proc_ucred = rk64(proc + 0x100);
    uint64_t amfi_entitlements = rk64(rk64(proc_ucred + 0x78) + 0x8);

    OSDictionary_SetItem(amfi_entitlements, "get-task-allow", find_OSBoolean_True());
    OSDictionary_SetItem(amfi_entitlements, "com.apple.private.skip-library-validation", find_OSBoolean_True());

    uint64_t present = OSDictionary_GetItem(amfi_entitlements, exc_key);

    int rv = 0;

    if (present == 0) {
        rv = OSDictionary_SetItem(amfi_entitlements, exc_key, get_exception_osarray());
    } else if (present != get_exception_osarray()) {
        unsigned int itemCount = OSArray_ItemCount(present);
        
        BOOL foundEntitlements = NO;
        
        uint64_t itemBuffer = OSArray_ItemBuffer(present);
        
        for (int i = 0; i < itemCount; i++) {
            uint64_t item = rk64(itemBuffer + (i * sizeof(void *)));
            char *entitlementString = OSString_CopyString(item);
            if (strcmp(entitlementString, "/meridian/") == 0){
                foundEntitlements = YES;
                free(entitlementString);
                break;
            }
            free(entitlementString);
        }
        
        if (!foundEntitlements){
            rv = OSArray_Merge(present, get_exception_osarray());
        } else {
            rv = 1;
        }
    } else {
        rv = 1;
    }

    if (rv != 1) {
        NSLog(@"Setting exc FAILED! amfi_entitlements: 0x%llx present: 0x%llx\n", amfi_entitlements, present);
    }
 
}

int setcsflagsandplatformize(int pid) {
    uint64_t proc = proc_find(pid, 3);
    if (proc == 0) {
        NSLog(@"Unable to find pid %d to entitle!", pid);
        return 1;
    }
    fprintf(stderr,"setcsflagsandplatformize start on PID %d\n", pid);
    char name[40] = {0};
    kread(proc+0x268, name, 20);
    fprintf(stderr,"PID %d name is %s\n", pid, name);
    set_csflags(proc);
    set_tfplatform(proc);
    set_amfi_entitlements(proc);
    set_sandbox_extensions(proc);
    set_csblob(proc);
    return 0;
}