#!/bin/bash

kmutilErrorCheck () {
    if [ $? -ne 0 ]
    then
        echo 'kmutil failed. See above output for more information.'
        echo 'patch-kexts.sh cannot continue.'
        exit 1
    fi
}

IMGVOL="/Volumes/Image Volume"
# Make sure we're inside the recovery environment. This may not be the best
# way to check, but it's simple and should work in the real world.
if [ ! -d "$IMGVOL" ]
then
    echo 'You must use this script from inside the Recovery environment.'
    echo 'Please restart your Mac from the patched Big Sur installer'
    echo 'USB drive, then open Terminal and try again.'
    echo
    echo '(The ability to use this script without rebooting into the'
    echo 'Recovery environment is planned for a future patcher release.)'
    exit 1
fi

VOLUME="$1"
echo "$VOLUME"
echo

# Make sure a volume has been specified. (Without this, other error checks
# eventually kick in, but the error messages get confusing.)
if [ -z "$VOLUME" ]
then
    echo 'You must specify a target volume (such as /Volumes/Macintosh\ HD)'
    echo 'on the command line.'
    exit 1
fi

# Sanity check to make sure that the specified $VOLUME isn't an obvious mistake
#
# DO NOT check for /Syste/Library/CoreServices here, or Big Sur data drives
# as well as system drives will pass the check!
if [ ! -d "$VOLUME/System/Library/Extensions" ]
then
    echo "Unable to find /System/Library/Extensions on the volume."
    echo "Cannot proceed. Make sure you specified the correct volume."
    echo "(Make sure to specify the system volume, not the data volume.)"
    exit 1
fi

# Check that the $VOLUME has macOS build 20*. This version check will
# hopefully keep working even after Apple bumps the version number to 11.
SVPL="$VOLUME"/System/Library/CoreServices/SystemVersion.plist
SVPL_VER=`fgrep '<string>10' "$SVPL" | sed -e 's@^.*<string>10@10@' -e 's@</string>@@' | uniq -d`
SVPL_BUILD=`grep '<string>[0-9][0-9][A-Z]' "$SVPL" | sed -e 's@^.*<string>@@' -e 's@</string>@@'`

if echo $SVPL_BUILD | grep -q '^20'
then
    echo -n "Volume appears to have a Big Sur installation (build" $SVPL_BUILD
    echo "). Continuing."
else
    if [ -z "$SVPL_VER" ]
    then
        echo 'Unable to detect macOS version on volume. Make sure you chose'
        echo 'the correct volume. Or, perhaps a newer patcher is required.'
    else
        echo 'Volume appears to have an older version of macOS. Probably'
        echo 'version' "$SVPL_VER" "build" "$SVPL_BUILD"
        echo 'Please make sure you specified the correct volume.'
    fi

    exit 1
fi

# Also check to make sure $VOLUME is an actual volume and not a snapshot.
# Maybe I'll add code later to handle the snapshot case, but in the recovery
# environment for Developer Preview 1, I've always seen it mount the actual
# volume and not a snapshot.
DEVICE=`df "$VOLUME" | tail -1 | sed -e 's@ .*@@'`
echo 'Volume is mounted from device: ' $DEVICE
# The following code is somewhat convoluted for just checking if there's
# a slice within a slice, but it should make things easier for future
# code that will actually handle this case.
POPSLICE=`echo $DEVICE | sed -E 's@s[0-9]+$@@'`
POPSLICE2=`echo $POPSLICE | sed -E 's@s[0-9]+$@@'`

if [ $POPSLICE = $POPSLICE2 ]
then
    echo 'Mounted volume is an actual volume, not a snapshot. Proceeding.'
else
    echo
    echo 'ERROR:'
    echo 'Mounted volume appears to be an APFS snapshot, not the underlying'
    echo 'volume. The patcher was not expecting to encounter this situation'
    echo 'within the Recovery environment, and an update to the patcher will'
    echo 'be required. Kext installation will not proceed.'
    exit 1
fi

# It's likely that at least one of these was reenabled during installation.
# But as we're in the recovery environment, there's no need to check --
# we'll just redisable these. If they're already disabled, then there's
# no harm done.
#csrutil disable
#csrutil authenticated-root disable

# Remount the volume read-write
echo "Remounting volume as read-write..."
if mount -uw "$VOLUME"
then
    # Remount succeeded. Do nothing in this block, and keep going.
    true
else
    echo "Remount failed. Kext installation cannot proceed."
    exit 1
fi

echo "Checking and removing apfs snapshots"
IFS=';'; snaplist=($(diskutil apfs listSnapshots "$VOLUME" | egrep -o " [0-9A-F-]{36}" | tr '\n' ';')); unset IFS
if [[ ! ${#snaplist[@]} = 0 ]]; then 
    "$VOLUME"/System/Library/F*/apfs.fs/C*/R*/apfs_systemsnapshot -v "$VOLUME" -r ""
    for i in ${snaplist[@]}; do diskutil apfs deleteSnapshot "$VOLUME" -uuid $i; done
    echo
else
    echo "no snapshots found"
fi

# Move the old kext out of the way, or delete if needed. Then unzip the
# replacement.
pushd "$VOLUME/System/Library/Extensions" > /dev/null

if [[ -f "$IMGVOL"/kexts/IO80211Family.kext.zip ]]; then
echo 'Installing highvoltage12v patched IO80211Family.kext'
    if [ -d IO80211Family.kext.original ]
    then
        rm -rf IO80211Family.kext
    else
        mv IO80211Family.kext IO80211Family.kext.original
    fi

    unzip -q "$IMGVOL/kexts/IO80211Family.kext.zip"
    rm -rf __MACOSX
    chown -R 0:0 IO80211Family.kext
    chmod -R 755 IO80211Family.kext
fi

if [[ -f "$IMGVOL"/kexts/AppleGraphicsPowerManagement.kext.zip ]]; then
echo 'Installing patched AppleGraphicsPowerManagement.kext'
    if [ -d AppleGraphicsPowerManagement.kext.original ]
    then
        rm -rf AppleGraphicsPowerManagement.kext
    else
        mv AppleGraphicsPowerManagement.kext AppleGraphicsPowerManagement.kext.original
    fi

    unzip -q "$IMGVOL/kexts/AppleGraphicsPowerManagement.kext.zip"
    rm -rf __MACOSX
    chown -R 0:0 AppleGraphicsPowerManagement.kext
    chmod -R 755 AppleGraphicsPowerManagement.kext
fi

if [[ -f "$IMGVOL"/kexts/IOPlatformPluginFamily.kext.zip ]]; then
echo 'Installing patched IOPlatformPluginFamily.kext'
    if [ -d IOPlatformPluginFamily.kext.original ]
    then
        rm -rf IOPlatformPluginFamily.kext
    else
        mv IOPlatformPluginFamily.kext IOPlatformPluginFamily.kext.original
    fi

    unzip -q "$IMGVOL/kexts/IOPlatformPluginFamily.kext.zip"
    rm -rf __MACOSX
    chown -R 0:0 IOPlatformPluginFamily.kext
    chmod -R 755 IOPlatformPluginFamily.kext
fi

if [[ -f "$IMGVOL"/kexts/AppleHDA.kext.zip ]]; then
echo 'Installing High Sierra AppleHDA.kext'
    if [ -d AppleHDA.kext.original ]
    then
        rm -rf AppleHDA.kext
    else
        mv AppleHDA.kext AppleHDA.kext.original
    fi

    unzip -q "$IMGVOL/kexts/AppleHDA.kext.zip"
    rm -rf __MACOSX
    chown -R 0:0 AppleHDA.kext
    chmod -R 755 AppleHDA.kext
fi

if [[ -f "$IMGVOL"/kexts/LegacyUSBInjector.kext.zip ]]; then
echo 'Installing LegacyUSBInjector.kext'
    rm -rf LegacyUSBInjector.kext

    unzip -q "$IMGVOL/kexts/LegacyUSBInjector.kext.zip"
    rm -rf __MACOSX
    chown -R 0:0 LegacyUSBInjector.kext
    chmod -R 755 LegacyUSBInjector.kext

    # parameter for kmutil later on
    BUNDLE_PATH="--bundle-path /System/Library/Extensions/LegacyUSBInjector.kext"
fi

if [[ -f "$IMGVOL"/kexts/GeForceTesla.zip ]] && [[ -f "$IMGVOL"/kexts/NVDANV50HalTesla.kext.zip ]] && [[ -f "$IMGVOL"/kexts/NVDAResmanTesla.kext.zip ]]; then
echo 'Installing GeForce Tesla (9400M/320M) kexts'
    rm -rf *Tesla*

    unzip -q "$IMGVOL/kexts/GeForceTesla.zip"
    rm -rf __MACOSX
    unzip -q "$IMGVOL/kexts/NVDANV50HalTesla.kext.zip"
    rm -rf __MACOSX
    unzip -q "$IMGVOL/kexts/NVDAResmanTesla.kext.zip"
    rm -rf __MACOSX

    chown -R 0:0 *Tesla*
    chmod -R 755 *Tesla*
fi

if [[ -f "$IMGVOL"/kexts/nvenet.kext.zip ]]; then
echo 'Installing High Sierra nvenet.kext'
    pushd IONetworkingFamily.kext/Contents/Plugins > /dev/null
    rm -rf __MACOSX
    rm -rf nvenet.kext
    unzip -q "$IMGVOL/kexts/nvenet.kext.zip"
    rm -rf __MACOSX
    chown -R 0:0 nvenet.kext
    chmod -R 755 nvenet.kext
    popd > /dev/null
fi

popd > /dev/null

if [[ -d "$VOLUME"/System/Library/UserEventPlugins/com.apple.telemetry.plugin ]]; then
echo 'Deactivating com.apple.telemetry.plugin'
    pushd "$VOLUME/System/Library/UserEventPlugins" > /dev/null
    mv -f com.apple.telemetry.plugin com.apple.telemetry.plugin.disabled
    popd > /dev/null
fi

echo

# Update the kernel/kext collections.
# kmutil *must* be invoked separately for boot and system KCs when
# LegacyUSBInjector is being used, or the injector gets left out, at least
# as of Big Sur beta 2. So, we'll always do it that way (even without
# LegacyUSBInjector, it shouldn't do any harm).
#
# I suspect it's not supposed to require the chroot, but I was getting weird
# "invalid argument" errors, and chrooting it eliminated those errors.
# BTW, kmutil defaults to "--volume-root /" according to the manpage, so
# it's probably redundant, but whatever.
echo 'Using kmutil to rebuild boot collection...'
chroot "$VOLUME" kmutil create -n boot \
    --kernel /System/Library/Kernels/kernel \
    --volume-root / $BUNDLE_PATH \
    --boot-path /System/Library/KernelCollections/BootKernelExtensions.kc
kmutilErrorCheck

# When creating SystemKernelExtensions.kc, kmutil requires *both* --boot-path
# and --system-path!
echo 'Using kmutil to rebuild system collection...'
chroot "$VOLUME" kmutil create -n sys \
    --kernel /System/Library/Kernels/kernel \
    --volume-root / \
    --system-path /System/Library/KernelCollections/SystemKernelExtensions.kc \
    --boot-path /System/Library/KernelCollections/BootKernelExtensions.kc
kmutilErrorCheck

# The way you control kcditto's *destination* is by choosing which volume
# you run it *from*. I'm serious. Read the kcditto manpage carefully if you
# don't believe me!
"$VOLUME/usr/sbin/kcditto"

bless --folder "$VOLUME"/System/Library/CoreServices --bootefi --create-snapshot

echo 'Installed patch kexts successfully.'