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

import Foundation

class NVRamUtil {
    private var electra: Electra
    init(electra: Electra) {
        self.electra = electra
    }
    
    private func getObject(serv: mach_port_t) -> UInt64 {
        let port = electra.findPort(port: serv)
        guard port != 0 else {
            return 0
        }
        return rk64(port + Offsets.shared.ipc_port.ip_kobject)
    }
    
    private func getOfDict(nvram_serv: mach_port_t) -> UInt64 {
        let nvramObj = getObject(serv: nvram_serv)
        guard nvramObj != 0 else {
            return 0
        }
        return rk64(nvramObj + Offsets.shared.nvram.io_dt_nvram_of_dict)
    }
    
    private func lookupKeyInOsDict(dict: UInt64, key: UInt64) -> UInt64 {
        var value = UInt64(0)
        
        let count = rk32(dict + Offsets.shared.osobject.os_dict_count)
        
        let entriessize = Int(count) * MemoryLayout<dict_entry_t>.size
        let os_dict_entries = malloc(entriessize)
        
        let entry_ptr = rk64(dict + Offsets.shared.osobject.os_dict_dict_entry)
        kread(entry_ptr, os_dict_entries, entriessize)
        
        value = lookup_key_in_dicts(os_dict_entries?.assumingMemoryBound(to: dict_entry_t.self), count, key)
        free(os_dict_entries)
        return value
    }
    
    private func generateNonce(nonce_serv: mach_port_t) -> kern_return_t {
        var nonce_conn: io_connect_t = 0
        var ret = KERN_FAILURE
        
        var nonce_d = [UInt8](repeating: 0, count: Offsets.shared.nvram.sha384_digest_length)
        var nonce_d_siz = Offsets.shared.nvram.sha384_digest_length
        if IOServiceOpen(nonce_serv, mach_task_self_, 0, &nonce_conn) == KERN_SUCCESS,
            nonce_conn != MACH_PORT_NULL {
            ret = IOConnectCallStructMethod(nonce_conn, Offsets.shared.nvram.ap_nonce_generate_nonce_sel, nil, 0, &nonce_d, &nonce_d_siz)
        }
        return ret
    }
    
    private func getBootNonceOSSymbol(nonce_serv: mach_port_t) -> UInt64 {
        let nonceObj = getObject(serv: nonce_serv)
        return rk64(nonceObj + Offsets.shared.nvram.ap_nonce_boot_nonce_os_symbol)
    }
    
    private func syncNVRam(nvram_serv: mach_port_t) -> Bool {
        guard IORegistryEntrySetCFProperty(nvram_serv,
                                           "chimera.tempkey" as CFString,
                                           "Freeze! Don't move!" as CFString) == KERN_SUCCESS else {
            return false
        }
        guard IORegistryEntrySetCFProperty(nvram_serv,
                                           "IONVRAM-DELETE-PROPERTY" as CFString,
                                           "chimera.tempkey" as CFString) == KERN_SUCCESS else {
            return false
        }
        guard IORegistryEntrySetCFProperty(nvram_serv,
                                           "IONVRAM-FORCESYNCNOW-PROPERTY" as CFString,
                                           "com.apple.System.boot-nonce" as CFString) == KERN_SUCCESS else {
            return false
        }
        return true
    }
    
    private func isGeneratorValid(generator: String) -> Bool {
        let generatorInput = generator.cString(using: .utf8)
        var rawGeneratorValue = UInt64(0)
        var generatorString = [Int8](repeating: 0, count: 22)
        
        var retval = false
        withUnsafeMutablePointer(to: &rawGeneratorValue) { rawGeneratorValuePtr in
            let args: [CVarArg] = [rawGeneratorValuePtr]
            vsscanf(generatorInput, "0x%16llx", getVaList(args))
            _ = snprintf(ptr: &generatorString, 22, "0x%016llx", rawGeneratorValuePtr.pointee)
            retval = (strcmp(generatorString, generatorInput) == 0)
        }
        return retval
    }
    
    public func setNonce(nonce: String) -> Bool {
        guard isGeneratorValid(generator: nonce) else {
            print("Invalid nonce")
            return false
        }
        
        let nvram = IOServiceGetMatchingService(kIOMasterPortDefault, IOServiceMatching("IODTNVRAM"))
        guard nvram != 0 else {
            return false
        }
        
        let apnonce = IOServiceGetMatchingService(kIOMasterPortDefault, IOServiceMatching("AppleMobileApNonce"))
        guard apnonce != 0 else {
            return false
        }
        
        let bootNonceSymbol = getBootNonceOSSymbol(nonce_serv: apnonce)
        guard bootNonceSymbol != 0 else {
            return false
        }
        
        /*let symStringPtr = rk64(bootNonceSymbol + Offsets.shared.osobject.os_string_string)
        var name = [UInt8](repeating: 0, count:27)
        kread(symStringPtr, &name, 27)
        
        print("found symbol: ", String(cString: &name))*/
        
        let ofDict = getOfDict(nvram_serv: nvram)
        guard ofDict != 0 else {
            return false
        }
        
        var osString = lookupKeyInOsDict(dict: ofDict, key: bootNonceSymbol)
        if osString == 0 {
            //nonceSet = false
            
            #if DEBUG
            print("Failed to get os_string. Trying to generate nonce")
            #endif
            if generateNonce(nonce_serv: apnonce) != KERN_SUCCESS {
                return false
            }
            osString = lookupKeyInOsDict(dict: ofDict, key: bootNonceSymbol)
            guard osString != 0 else {
                return false
            }
        }
        
        let stringPtr = rk64(osString + Offsets.shared.osobject.os_string_string)
        if stringPtr == 0 {
            return false
        }
        
        let nonceSize = 19 //2 * 4 + 3
        var nonceHex = [UInt8](repeating: 0, count: nonceSize)
        kread(stringPtr, &nonceHex, nonceSize - 1)
        
        print("Current nonce: ", String(cString: &nonceHex))
        
        kwrite(stringPtr, nonce.cString(using: .utf8), nonceSize - 1)
        
        bzero(&nonceHex, nonceSize)
        kread(stringPtr, &nonceHex, nonceSize - 1)
        return syncNVRam(nvram_serv: nvram)
    }
}
