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

import Foundation

public class Remount {
    let our_proc: UInt64
    let kernel_proc: UInt64
    
    init(our_proc: UInt64, kernel_proc: UInt64) {
        self.our_proc = our_proc
        self.kernel_proc = kernel_proc
    }
    
    private let offsets = Offsets.shared
    
    let mntpathSW = "/var/rootfsmnt"
    let mntpath = strdup("/var/rootfsmnt")
    
    private func findRootVnode(launchd_proc: UInt64) -> UInt64 {
        let textvp = rk64(launchd_proc + offsets.proc.textvp)
        var nameptr = rk64(textvp + offsets.vnode.name)
        var name = [UInt8](repeating: 0, count: 20)
        kread(nameptr, &name, 20)
        
        print("found vnode: ", String(cString: &name))
        
        let sbin = rk64(textvp + offsets.vnode.parent)
        nameptr = rk64(sbin + offsets.vnode.name)
        kread(nameptr, &name, 20)
        
        print("found vnode (should be sbin): ", String(cString: &name))
        
        let rootvnode = rk64(sbin + offsets.vnode.parent)
        nameptr = rk64(rootvnode + offsets.vnode.name)
        kread(nameptr, &name, 20)
        
        print("found vnode (should be root): ", String(cString: &name))
        
        let flags = rk32(rootvnode + offsets.vnode.flag)
        print(String(format: "vnode flags: 0x%x", flags))
        
        return rootvnode
    }
    
    private func isOTAMounted() -> Bool {
        let path = strdup("/var/MobileSoftwareUpdate/mnt1")
        defer {
            free(path)
        }
        
        var buffer = stat()
        if lstat(path, &buffer) != 0 {
            return false
        }
        
        let S_IFMT = 0o170000
        let S_IFDIR = 0o040000
        
        guard Int(buffer.st_mode) & S_IFMT == S_IFDIR else {
            return false
        }
        
        let cwd = getcwd(nil, 0)
        chdir(path)
        
        var p_buf = stat()
        lstat("..", &p_buf)
        
        if let cwd = cwd {
            chdir(cwd)
            free(cwd)
        }
        
        return buffer.st_dev != p_buf.st_dev || buffer.st_ino == p_buf.st_ino
    }
    
    private func isRenameRequired() -> Bool {
        let fd = open("/", O_RDONLY, 0)
        print("fd:", fd)
        
        var alist = attrlist()
        alist.commonattr = ATTR_BULK_REQUIRED
        
        var abuf = [UInt8](repeating: 0, count: 2048)
        let count = fs_snapshot_list(fd, &alist, &abuf[0], 2048, 0)
        print("snapshot count: ", count)
        close(fd)
        
        return count == -1
    }
    
    private func find_boot_snapshot() -> String? {
        let chosen = IORegistryEntryFromPath(0, "IODeviceTree:/chosen")
        guard let data = IORegistryEntryCreateCFProperty(chosen,
                                                         "boot-manifest-hash" as CFString,
                                                         kCFAllocatorDefault, 0).takeUnretainedValue() as? Data else {
            return nil
        }
        IOObjectRelease(chosen)
        
        var manifestHash = ""
        let buf = [UInt8](data)
        for byte in buf {
            manifestHash = manifestHash.appendingFormat("%02X", byte)
        }
        
        let systemSnapshot = "com.apple.os.update-" + manifestHash
        
        print("System Snapshot: ", systemSnapshot)
        return systemSnapshot
    }
    
    private func mountRealRootfs(rootvnode: UInt64) -> Int32 {
        let vmount = rk64(rootvnode + offsets.vnode.mount)
        let dev = rk64(vmount + offsets.mount.devvp)
        
        let nameptr = rk64(dev + offsets.vnode.name)
        var name = [UInt8](repeating: 0, count: 20)
        kread(nameptr, &name, 20)
        print("found dev vnode name: ", String(cString: &name))
        
        let specinfo = rk64(dev + offsets.vnode.specinfo)
        let flags = rk32(specinfo + offsets.specinfo.flags)
        print("found dev flags: ", flags)
        
        wk32(specinfo + offsets.specinfo.flags, 0)
        
        let fspec = strdup("/dev/disk0s1s1")
        
        var mntargs = hfs_mount_args()
        mntargs.fspec = fspec
        mntargs.hfs_mask = 1
        gettimeofday(nil, &mntargs.hfs_timezone)
        
        let retval = mount("apfs", mntpath, 0, &mntargs)
        
        free(fspec)
        
        print("mount completed with status ", retval)
        
        return retval
    }
    
    private func findNewMount(rootvnode: UInt64) -> UInt64? {
        var vmount = rk64(rootvnode + offsets.vnode.mount)
        
        vmount = rk64(vmount + offsets.mount.mnt_next)
        while vmount != 0 {
            let dev = rk64(vmount + offsets.mount.devvp)
            if dev != 0 {
                let nameptr = rk64(dev + offsets.vnode.name)
                var name = [UInt8](repeating: 0, count: 20)
                kread(nameptr, &name, 20)
                let devName = String(cString: &name)
                print("found dev vnode name: ", devName)
                
                if devName == "disk0s1s1" {
                    return vmount
                }
            }
            
            vmount = rk64(vmount + offsets.mount.mnt_next)
        }
        return nil
    }
    
    private func unsetSnapshotFlag(newmnt: UInt64) -> Bool {
        let dev = rk64(newmnt + offsets.mount.devvp)
        
        let nameptr = rk64(dev + offsets.vnode.name)
        var name = [UInt8](repeating: 0, count: 20)
        kread(nameptr, &name, 20)
        print("found dev vnode name: ", String(cString: &name))
        
        let specinfo = rk64(dev + offsets.vnode.specinfo)
        let flags = rk32(specinfo + offsets.specinfo.flags)
        print("found dev flags: ", flags)
        
        var vnodelist = rk64(newmnt + offsets.mount.vnodelist)
        while vnodelist != 0 {
            print("vnodelist: ", vnodelist)

            let nameptr = rk64(vnodelist + offsets.vnode.name)
            let len = Int(kstrlen(nameptr))
            var name = [UInt8](repeating: 0, count: len)
            kread(nameptr, &name, len)
            
            let vnodeName = String(cString: &name)
            print("found vnode name: ", vnodeName)
            
            if vnodeName.hasPrefix("com.apple.os.update-") {
                let vdata = rk64(vnodelist + offsets.vnode.data)
                
                let flag = rk32(vdata + offsets.apfsData.flag)
                print("found apfs flag: ", flag)
                
                if (flag & 0x40) != 0 {
                    print("would unset the flag here to", flag & ~0x40)
                    wk32(vdata + offsets.apfsData.flag, flag & ~0x40)
                    return true
                }
            }
            
            usleep(1000)
            vnodelist = rk64(vnodelist + UInt64(0x20))
        }
        return false
    }
    
    public func remount(launchd_proc: UInt64) -> Bool {
        let rootvnode = findRootVnode(launchd_proc: launchd_proc)
        if self.isRenameRequired() {
            if FileManager.default.fileExists(atPath: mntpathSW) {
                try? FileManager.default.removeItem(atPath: mntpathSW)
            }
            
            mkdir(mntpath, 0755)
            chown(mntpath, 0, 0)
            
            if isOTAMounted() {
                print("OTA update already mounted")
                return false
            }
    
            let kern_ucred = rk64(kernel_proc + offsets.proc.ucred)
            let our_ucred = rk64(our_proc + offsets.proc.ucred)
            
            wk64(our_proc + offsets.proc.ucred, kern_ucred)
            
            guard let bootSnapshot = find_boot_snapshot(),
                mountRealRootfs(rootvnode: rootvnode) == 0 else {
                wk64(our_proc + offsets.proc.ucred, our_ucred)
                return false
            }
            
            let fd = open("/var/rootfsmnt", O_RDONLY, 0)
            guard fd > 0,
                fs_snapshot_revert(fd, bootSnapshot.cString(using: .utf8), 0) == 0 else {
                print("fs_snapshot_revert failed")
                wk64(our_proc + offsets.proc.ucred, our_ucred)
                return false
            }
            close(fd)
            
            unmount(mntpath, MNT_FORCE)
            guard mountRealRootfs(rootvnode: rootvnode) == 0,
                let newmnt = findNewMount(rootvnode: rootvnode),
                unsetSnapshotFlag(newmnt: newmnt) else {
                wk64(our_proc + offsets.proc.ucred, our_ucred)
                return false
            }
            
            let fd2 = open("/var/rootfsmnt", O_RDONLY, 0)
            guard fd2 > 0,
                fs_snapshot_rename(fd2, bootSnapshot.cString(using: .utf8), "orig-fs", 0) == 0 else {
                print("fs_snapshot_rename failed")
                wk64(our_proc + offsets.proc.ucred, our_ucred)
                return false
            }
            close(fd2)
            
            unmount(mntpath, 0)
            rmdir(mntpath)
            
            wk64(our_proc + offsets.proc.ucred, our_ucred)
            
            print("rebooting...")
            sleep(2)
            reboot(0)
        } else {
            let vmount = rk64(rootvnode + offsets.vnode.mount)
            let vflag = rk32(vmount + offsets.mount.flag) & ~(UInt32(MNT_NOSUID) | UInt32(MNT_RDONLY))
            wk32(vmount + offsets.mount.flag, vflag & ~UInt32(MNT_ROOTFS))
            
            var dev_path = strdup("/dev/disk0s1s1")
            let retval = mount("apfs", "/", MNT_UPDATE, &dev_path)
            free(dev_path)
            
            wk32(vmount + offsets.mount.flag, vflag)
            return retval == 0
        }
        return true
    }
    
    public func restore_rootfs() -> Bool {
        if !self.isRenameRequired() {
            guard let bootSnapshot = find_boot_snapshot() else {
                return false
            }
            
            try? FileManager.default.removeItem(atPath: "/var/cache")
            try? FileManager.default.removeItem(atPath: "/var/lib")
            
            let kern_ucred = rk64(kernel_proc + offsets.proc.ucred)
            let our_ucred = rk64(our_proc + offsets.proc.ucred)
            
            wk64(our_proc + offsets.proc.ucred, kern_ucred)
            
            let fd = open("/", O_RDONLY, 0)
            guard fd > 0,
                fs_snapshot_rename(fd, "orig-fs", bootSnapshot.cString(using: .utf8), 0) == 0 else {
                print("fs_snapshot_rename failed")
                wk64(our_proc + offsets.proc.ucred, our_ucred)
                return false
            }
            guard fs_snapshot_revert(fd, bootSnapshot.cString(using: .utf8), 0) == 0 else {
                print("fs_snapshot_revert failed")
                wk64(our_proc + offsets.proc.ucred, our_ucred)
                return false
            }
            close(fd)
            
            wk64(our_proc + offsets.proc.ucred, our_ucred)
            
            print("rebooting...")
            sleep(2)
            reboot(0)
        } else {
            print("rootfs restore not required")
        }
        return true
    }
}
