//
//  electra.swift
//  Chimera13
//
//  Created by CoolStar on 3/1/20.
//  Copyright © 2020 coolstar. All rights reserved.
//

import Foundation

enum JAILBREAK_RETURN_STATUS {
    case ERR_NOERR
    case ERR_VERSION
    case ERR_EXPLOIT
    case ERR_UNSUPPORED
    case ERR_TFP0
    case ERR_ALREADY_JAILBROKEN
    case ERR_ROOTFS_RESTORE
    case ERR_REMOUNT
    case ERR_SNAPSHOT
    case ERR_JAILBREAK
    case ERR_CONFLICT
}

class Electra {
    private let tfp0: mach_port_t
    #if MAINAPP
    private let any_proc: UInt64
    private let enable_tweaks: Bool
    private let restore_rootfs: Bool
    private let nonce: String
    #endif
    
    private let offsets = Offsets.shared
    private let consts = Consts.shared
    
    #if MAINAPP
    private var kernel_slide: UInt64 = 0
    #endif
    
    private var all_proc: UInt64 = 0
    
    private(set) var our_proc: UInt64 = 0
    private var launchd_proc: UInt64 = 0
    private var kernel_proc: UInt64 = 0
    
    private(set) var amfid_pid: UInt32 = 0
    
    private var our_task_addr: UInt64 = 0
    #if MAINAPP
    private var our_label: UInt64 = 0
    
    private var root_vnode: UInt64 = 0
    
    public init(tfp0: mach_port_t, any_proc: UInt64, enable_tweaks: Bool, restore_rootfs: Bool, nonce: String) {
        self.tfp0 = tfp0
        self.any_proc = any_proc
        self.enable_tweaks = enable_tweaks
        self.restore_rootfs = restore_rootfs
        self.nonce = nonce
    }
    #else
    public init(tfp0: mach_port_t, all_proc: UInt64) {
        self.tfp0 = tfp0
        self.all_proc = all_proc
    }
    #endif
    
    #if MAINAPP
    private func find_allproc() {
        var proc = any_proc
        while proc != 0 {
            if rk64(rk64(proc + 8)) != proc {
                print("Reached end of list")
                print(String(format: "Allproc? 0x%llx", proc))
                all_proc = proc
                break
            }
            proc = rk64(proc + 8)
        }
    }
    
    public func findPort(port: mach_port_name_t) -> UInt64 {
        let ourTask = rk64(self.our_proc + offsets.proc.task)
        let itkSpace = rk64(ourTask + offsets.task.itk_space)
        let isTable = rk64(itkSpace + offsets.ipc_space.is_table)
        
        let portIndex = UInt32(port) >> 8
        let ipcEntrySz = UInt32(0x18)
        
        let portAddr = rk64(isTable + UInt64((portIndex * ipcEntrySz)))
        return portAddr
    }
    #endif
    
    public func populate_procs() {
        let our_pid = getpid()
        var proc = rk64(all_proc)
        while proc != 0 {
            let pid = rk32(proc + offsets.proc.pid)
            if pid == 0 {
                kernel_proc = proc
                print("found kernel proc")
            } else if pid == our_pid {
                print("found our pid")
                
                our_proc = proc
                our_task_addr = rk64(proc + offsets.proc.task)
                
                let our_flags = rk32(our_task_addr + offsets.task.flags)
                wk32(our_task_addr + offsets.task.flags, our_flags | consts.TF_PLATFORM)
                
                var our_csflags = rk32(our_proc + offsets.proc.csflags)
                our_csflags = our_csflags | consts.CS_PLATFORM_BINARY | consts.CS_INSTALLER | consts.CS_GET_TASK_ALLOW
                our_csflags &= ~(consts.CS_RESTRICT | consts.CS_HARD | consts.CS_KILL)
                wk32(our_proc + offsets.proc.csflags, our_csflags)
            } else if pid == 1 {
                print("found launchd")
                
                launchd_proc = proc
            } else {
                let nameptr = proc + offsets.proc.name
                var name = [UInt8](repeating: 0, count: 32)
                kread(nameptr, &name, 32)
                //print("found proc name: ", String(cString: &name))
                
                if String(cString: &name) == "amfid" {
                    print("found amfid")
                    amfid_pid = pid
                }
            }
            proc = rk64(proc)
        }
    }
    
    public func find_proc(pid: UInt32) -> UInt64 {
        var proc = rk64(all_proc)
        while proc != 0 {
            let proc_pid = rk32(proc + offsets.proc.pid)
            if proc_pid == pid {
                return proc
            }
            proc = rk64(proc)
        }
        return proc
    }
    
    #if MAINAPP
    private func getRoot() -> JAILBREAK_RETURN_STATUS {
        let kern_ucred = rk64(kernel_proc + offsets.proc.ucred)
        let self_ucred = rk64(our_proc + offsets.proc.ucred)
        
        let our_label = rk64(self_ucred + offsets.ucred.cr_label)
        wk64(self_ucred + offsets.ucred.cr_label, rk64(kern_ucred + offsets.ucred.cr_label))
        wk32(self_ucred + offsets.ucred.cr_svuid, UInt32(0))
        
        setuid(0)
        setuid(0)
        
        wk64(self_ucred + offsets.ucred.cr_label, our_label)
        
        guard getuid() == 0 else {
            return .ERR_JAILBREAK
        }
        return .ERR_NOERR
    }
    
    private func cleanupCreds() {
        setuid(501)
        print("Reset creds")
    }
    
    public func jailbreak() -> JAILBREAK_RETURN_STATUS {
        print("Starting Electra...")
        guard tfp0 != MACH_PORT_NULL else {
            return .ERR_TFP0
        }
        
        let mobile_realhost = mach_host_self()
        
        var err: JAILBREAK_RETURN_STATUS = .ERR_NOERR
        
        find_allproc()
        populate_procs()
        
        let slide = getKernSlide(our_proc: our_proc)
        print(String(format: "kernel slide is at 0x%016llx", slide))
        kernel_slide = slide
        
        print(String(format: "our proc is at 0x%016llx", our_proc))
        print(String(format: "kern proc is at 0x%016llx", kernel_proc))
        
        err = getRoot()
        if err != .ERR_NOERR {
            return err
        }
        defer { cleanupCreds() }
        
        print(String(format: "our uid is %d", getuid()))
        
        let nvram = NVRamUtil(electra: self)
        _ = nvram.setNonce(nonce: nonce) //Not fatal is nonce setting fails
        
        let remount = Remount(our_proc: our_proc, kernel_proc: kernel_proc)
        if !remount.remount(launchd_proc: launchd_proc) {
            return .ERR_REMOUNT
        }
        if restore_rootfs {
            if !remount.restore_rootfs() {
                return .ERR_ROOTFS_RESTORE
            }
        }
        
        _ = SetHSP4(electra: self, tfp0: tfp0, slide: slide, kernel_proc: kernel_proc, our_proc: our_proc, old_realhost: mobile_realhost)
        
        try? FileManager.default.removeItem(atPath: "/chimera")
        
        mkdir("/chimera", 0755)
        chown("/chimera", 0, 0)
        
        mkdir("/chimera/bin/", 0755)
        chown("/chimera/bin/", 0, 0)
        
        try? FileManager.default.copyItem(at: Bundle.main.url(forResource: "amfidebilitate", withExtension: "")!,
                                          to: URL(fileURLWithPath: "/chimera/amfidebilitate"))
        try? FileManager.default.copyItem(at: Bundle.main.url(forResource: "bash", withExtension: "")!,
                                          to: URL(fileURLWithPath: "/chimera/bin/bash"))
        try? FileManager.default.copyItem(at: Bundle.main.url(forResource: "binbag", withExtension: "")!,
                                          to: URL(fileURLWithPath: "/chimera/bin/binbag"))
        try? FileManager.default.copyItem(at: Bundle.main.url(forResource: "dropbear", withExtension: "")!,
                                          to: URL(fileURLWithPath: "/chimera/dropbear"))
        try? FileManager.default.copyItem(at: Bundle.main.url(forResource: "kmap", withExtension: "")!,
                                          to: URL(fileURLWithPath: "/chimera/kmap"))
        try? FileManager.default.copyItem(at: Bundle.main.url(forResource: "kmem", withExtension: "")!,
                                          to: URL(fileURLWithPath: "/chimera/kmem"))
        chmod("/chimera/amfidebilitate", 0755)
        chmod("/chimera/bin/bash", 0755)
        chmod("/chimera/bin/binbag", 0755)
        chmod("/chimera/dropbear", 0755)
        chmod("/chimera/kmap", 0755)
        chmod("/chimera/kmem", 0755)
        chown("/chimera/amfidebilitate", 0, 0)
        chown("/chimera/bin/bash", 0, 0)
        chown("/chimera/bin/binbag", 0, 0)
        chown("/chimera/dropbear", 0, 0)
        chown("/chimera/kmap", 0, 0)
        chown("/chimera/kmem", 0, 0)
        
        symlink("/chimera/bin/bash", "/chimera/bin/sh")
        symlink("/chimera/bin/binbag", "/chimera/bin/mount")
        symlink("/chimera/bin/binbag", "/chimera/bin/umount")
        symlink("/chimera/bin/binbag", "/chimera/bin/chflags")
        symlink("/chimera/bin/binbag", "/chimera/bin/chmod")
        symlink("/chimera/bin/binbag", "/chimera/bin/chown")
        symlink("/chimera/bin/binbag", "/chimera/bin/chgrp")
        symlink("/chimera/bin/binbag", "/chimera/bin/cp")
        symlink("/chimera/bin/binbag", "/chimera/bin/dd")
        symlink("/chimera/bin/binbag", "/chimera/bin/df")
        symlink("/chimera/bin/binbag", "/chimera/bin/du")
        symlink("/chimera/bin/binbag", "/chimera/bin/ln")
        symlink("/chimera/bin/binbag", "/chimera/bin/ls")
        symlink("/chimera/bin/binbag", "/chimera/bin/mkdir")
        symlink("/chimera/bin/binbag", "/chimera/bin/mkfifo")
        symlink("/chimera/bin/binbag", "/chimera/bin/mknod")
        symlink("/chimera/bin/binbag", "/chimera/bin/mv")
        symlink("/chimera/bin/binbag", "/chimera/bin/readlink")
        symlink("/chimera/bin/binbag", "/chimera/bin/rm")
        symlink("/chimera/bin/binbag", "/chimera/bin/rmdir")
        symlink("/chimera/bin/binbag", "/chimera/bin/stat")
        symlink("/chimera/bin/binbag", "/chimera/bin/touch")
        symlink("/chimera/bin/binbag", "/chimera/bin/netcat")
        symlink("/chimera/bin/binbag", "/chimera/bin/arp")
        symlink("/chimera/bin/binbag", "/chimera/bin/ifconfig")
        symlink("/chimera/bin/binbag", "/chimera/bin/netstat")
        symlink("/chimera/bin/binbag", "/chimera/bin/ping")
        symlink("/chimera/bin/binbag", "/chimera/bin/traceroute")
        symlink("/chimera/bin/binbag", "/chimera/bin/basename")
        symlink("/chimera/bin/binbag", "/chimera/bin/date")
        symlink("/chimera/bin/binbag", "/chimera/bin/dirname")
        symlink("/chimera/bin/binbag", "/chimera/bin/expr")
        symlink("/chimera/bin/binbag", "/chimera/bin/find")
        symlink("/chimera/bin/binbag", "/chimera/bin/id")
        symlink("/chimera/bin/binbag", "/chimera/bin/hexdump")
        symlink("/chimera/bin/binbag", "/chimera/bin/hostname")
        symlink("/chimera/bin/binbag", "/chimera/bin/jot")
        symlink("/chimera/bin/binbag", "/chimera/bin/killall")
        symlink("/chimera/bin/binbag", "/chimera/bin/mktemp")
        symlink("/chimera/bin/binbag", "/chimera/bin/nice")
        symlink("/chimera/bin/binbag", "/chimera/bin/nohup")
        symlink("/chimera/bin/binbag", "/chimera/bin/printf")
        symlink("/chimera/bin/binbag", "/chimera/bin/pwd")
        symlink("/chimera/bin/binbag", "/chimera/bin/renice")
        symlink("/chimera/bin/binbag", "/chimera/bin/script")
        symlink("/chimera/bin/binbag", "/chimera/bin/seq")
        symlink("/chimera/bin/binbag", "/chimera/bin/sleep")
        symlink("/chimera/bin/binbag", "/chimera/bin/systime")
        symlink("/chimera/bin/binbag", "/chimera/bin/tee")
        symlink("/chimera/bin/binbag", "/chimera/bin/test")
        symlink("/chimera/bin/binbag", "/chimera/bin/time")
        symlink("/chimera/bin/binbag", "/chimera/bin/uname")
        symlink("/chimera/bin/binbag", "/chimera/bin/w")
        symlink("/chimera/bin/binbag", "/chimera/bin/what")
        symlink("/chimera/bin/binbag", "/chimera/bin/which")
        symlink("/chimera/bin/binbag", "/chimera/bin/xargs")
        symlink("/chimera/bin/binbag", "/chimera/bin/hostinfo")
        symlink("/chimera/bin/binbag", "/chimera/bin/lsmp")
        symlink("/chimera/bin/binbag", "/chimera/bin/mean")
        symlink("/chimera/bin/binbag", "/chimera/bin/nvram")
        symlink("/chimera/bin/binbag", "/chimera/bin/reboot")
        symlink("/chimera/bin/binbag", "/chimera/bin/sc_usage")
        symlink("/chimera/bin/binbag", "/chimera/bin/stackshot")
        symlink("/chimera/bin/binbag", "/chimera/bin/sysctl")
        symlink("/chimera/bin/binbag", "/chimera/bin/trace")
        symlink("/chimera/bin/binbag", "/chimera/bin/cat")
        symlink("/chimera/bin/binbag", "/chimera/bin/cut")
        symlink("/chimera/bin/binbag", "/chimera/bin/comm")
        symlink("/chimera/bin/binbag", "/chimera/bin/grep")
        symlink("/chimera/bin/binbag", "/chimera/bin/head")
        symlink("/chimera/bin/binbag", "/chimera/bin/split")
        symlink("/chimera/bin/binbag", "/chimera/bin/sort")
        symlink("/chimera/bin/binbag", "/chimera/bin/tail")
        symlink("/chimera/bin/binbag", "/chimera/bin/tr")
        symlink("/chimera/bin/binbag", "/chimera/bin/uniq")
        symlink("/chimera/bin/binbag", "/chimera/bin/wc")
        symlink("/chimera/bin/binbag", "/chimera/bin/tar")
        
        try? FileManager.default.createDirectory(atPath: "/chimera/etc/dropbear/", withIntermediateDirectories: true, attributes: nil)
        
        if let dirList = try? FileManager.default.contentsOfDirectory(atPath: "/chimera/") {
            print(dirList)
        }
        
        let amfidtakeover = AmfidTakeover(electra: self)
        guard amfidtakeover.grabEntitlements(our_proc: our_proc) else {
            return .ERR_JAILBREAK
        }
        amfidtakeover.takeoverAmfid(amfid_pid: amfid_pid)
        guard amfidtakeover.spawnAmfiDebilitate(allProc: all_proc) else {
            return .ERR_JAILBREAK
        }
        
        print("waiting for amfidebilitate...")
        while !amfidtakeover.amfidebilitate_spawned {
            usleep(1000)
        }
        amfidtakeover.resetEntitlements(our_proc: our_proc)
        
        guard spawnDropbear() else {
            return .ERR_JAILBREAK
        }
        
        return err
    }
    
    public func spawnDropbear() -> Bool {
        // swiftlint:disable line_length
        let dict = xpc_dictionary_create(nil, nil, 0)
        let request = xpc_dictionary_create(nil, nil, 0)
        let submitJob = xpc_dictionary_create(nil, nil, 0)
        let environmentVariables = xpc_dictionary_create(nil, nil, 0)
        let programArguments = xpc_array_create(nil, 0)
        
        xpc_dictionary_set_string(environmentVariables, "PS1", "\\h:\\w \\u\\$ ")
        xpc_dictionary_set_string(environmentVariables, "PATH",
                                  "/chimera:/chimera/usr/local/sbin:/chimera/usr/local/bin:/chimera/usr/sbin:/chimera/usr/bin:/chimera/sbin:/chimera/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/bin/X11:/usr/games")
        
        xpc_array_append_value(programArguments, xpc_string_create("dropbear"))
        xpc_array_append_value(programArguments, xpc_string_create("-S"))
        xpc_array_append_value(programArguments, xpc_string_create("/chimera"))
        xpc_array_append_value(programArguments, xpc_string_create("-p"))
        xpc_array_append_value(programArguments, xpc_string_create("22"))
        xpc_array_append_value(programArguments, xpc_string_create("-p"))
        xpc_array_append_value(programArguments, xpc_string_create("2222"))
        xpc_array_append_value(programArguments, xpc_string_create("-F"))
        
        xpc_dictionary_set_bool(submitJob, "KeepAlive", true)
        xpc_dictionary_set_bool(submitJob, "RunAtLoad", true)
        xpc_dictionary_set_string(submitJob, "UserName", "root")
        xpc_dictionary_set_string(submitJob, "Program", "/chimera/dropbear")
        xpc_dictionary_set_string(submitJob, "Label", "dropbear")
        xpc_dictionary_set_value(submitJob, "EnvironmentVariables", environmentVariables)
        xpc_dictionary_set_value(submitJob, "ProgramArguments", programArguments)
        
        xpc_dictionary_set_value(request, "SubmitJob", submitJob)
        xpc_dictionary_set_value(dict, "request", request)
        
        xpc_dictionary_set_uint64(dict, "subsystem", 7)
        xpc_dictionary_set_uint64(dict, "type", 7)
        xpc_dictionary_set_uint64(dict, "handle", 0)
        xpc_dictionary_set_uint64(dict, "routine", UInt64(ROUTINE_SUBMIT))
        
        var outDict: xpc_object_t?
        let rc = xpc_pipe_routine(xpc_bootstrap_pipe(), dict, &outDict)
        if rc == 0,
            let outDict = outDict {
            let rc2 = Int32(xpc_dictionary_get_int64(outDict, "error"))
            if rc2 != 0 {
                print(String(format: "Error submitting service: %s", xpc_strerror(rc2)))
                return false
            }
        } else if rc != 0 {
            print(String(format: "Error submitting service (no outdict): %s", xpc_strerror(rc)))
            return false
        }
        return true
    }
    #endif
}
