#!/usr/bin/bash

OLD_SHIM_HASHES=("8061e727e908ca58583c3a83b318ebe61b69fcd36c545cfd15ca79c89c8cc2e2")
NEW_SHIM_HASH="f3d6ff4c66821c631bf4bfd0c23d4e04c11142a38b1deebffd26bb454fb3f64e"

NEW_SHIM_URL="https://github.com/Lernstick/lernstickAdvanced/raw/3ba2ce2f13a11214ed770ee74a451d963e267322/config/includes.binary/EFI/boot/bootx64.efi.lernstick"
NEW_MMX_URL="https://github.com/Lernstick/lernstickAdvanced/raw/3ba2ce2f13a11214ed770ee74a451d963e267322/config/includes.binary/EFI/boot/mmx64.efi"
NEW_GRUB_URL="https://github.com/Lernstick/lernstickAdvanced/raw/3ba2ce2f13a11214ed770ee74a451d963e267322/config/includes.binary/EFI/boot/grubx64.efi.lernstick"

SHIM_LOCATION="EFI/boot/bootx64.efi"
MMX_LOCATION="EFI/boot/mmx64.efi"
GRUB_LOCATION="EFI/boot/grubx64.efi"

TMP_DIR=$(mktemp -d)
MOUNT_POINT="${TMP_DIR}/mount"

RET_SKIP=2
RET_ERR=1

download() {
    echo "Downloading new files"
    if ! curl --location "${NEW_SHIM_URL}" --remote-name --output-dir "${TMP_DIR}" --fail --silent
    then
        echo "failed to download new shim"
        exit ${RET_ERR}
    fi
    if ! curl --location "${NEW_MMX_URL}" --remote-name --output-dir "${TMP_DIR}" --fail --silent
    then
        echo "failed to download new mok manager"
        exit ${RET_ERR}
    fi
    if ! curl --location "${NEW_GRUB_URL}" --remote-name --output-dir "${TMP_DIR}" --fail --silent
    then
        echo "failed to download new GRUB"
        exit ${RET_ERR}
    fi
}

check_root() {
    if [ "$(id -u)" -ne 0 ]
        then echo This script needs to be run as root!
        exit ${RET_ERR}
    fi
}

check_programs() {
    if ! command -v curl  &> /dev/null
    then
        echo "You need to install curl to run this script"
        exit ${RET_ERR}
    fi

    if ! command -v awk  &> /dev/null
    then
        echo "You need to install awk to run this script"
        exit ${RET_ERR}
    fi
}

update_partition() {
    local partition=$1
   
    mkdir -p "${MOUNT_POINT}"
    if ! mount "${partition}" "${MOUNT_POINT}" > /dev/null 2>&1
    then
        echo "could mount partition from ${partition}. Skipping..."
        return ${RET_SKIP}
    fi

    if awk '$4~/(^|,)ro($|,)/' /proc/mounts | grep -q "${MOUNT_POINT}"
    then
        echo "Filesystem is readonly. Skipping..."
        if ! umount "${MOUNT_POINT}"
        then
            echo "unmounting partition failed"
            return ${RET_ERR}
        fi
        return ${RET_SKIP}
    fi

    local ret
    ret=0
    if has_lernstick_shim "${MOUNT_POINT}/${SHIM_LOCATION}"
    then
        update_files "${MOUNT_POINT}"
        ret=$?
    else 
        ret=$?
    fi

    if ! umount "${MOUNT_POINT}"
    then
        echo "unmounting partition failed"
        return ${RET_ERR}
    fi

    return $ret
}

has_lernstick_shim() {
    local shim_path=$1
    local shim_hash
    if [ ! -f "${shim_path}" ]
    then
        echo "No shim found. Skipping..."
        return ${RET_SKIP}
    fi
    if ! shim_hash=$(sha256sum "${shim_path}" | awk '{print $1}')
    then 
        echo "Couldn't caclulate shim hash something went wrong"
        return ${RET_ERR}
    fi

    if [ "${shim_hash}" = "${NEW_SHIM_HASH}" ]
    then
        echo "Shim already up-to-date. Skipping..."
        return ${RET_SKIP}
    fi

    for old_hash in "${OLD_SHIM_HASHES[@]}"
    do
        if [ "${shim_hash}" = "${old_hash}" ]
        then
            return 0
        fi
    done
    echo "Not our shim. Skipping..."
    return ${RET_SKIP}
}

update_files() {
    local update_dir=$1

    echo "Updating to latest files..."
    if ! cp "${TMP_DIR}/bootx64.efi.lernstick" "${update_dir}/${SHIM_LOCATION}"
    then
        echo "Copying new Shim to partition failed"
        return ${RET_ERR}
    fi

    if ! cp "${TMP_DIR}/mmx64.efi" "${update_dir}/${MMX_LOCATION}"
    then
        echo "Copying new Mok Mananger to partition failed"
        return ${RET_ERR}
    fi

    if ! cp "${TMP_DIR}/grubx64.efi.lernstick" "${update_dir}/${GRUB_LOCATION}"
    then
        echo "Copying new GRUB to partition failed"
        return ${RET_ERR}
    fi

    return 0
}


update_efis() {
    local vfat_partitions
    
    vfat_partitions=$(lsblk -o name,fstype -lpn | awk '{if ($2 == "vfat") {print $1} }')
    for partition in ${vfat_partitions}
    do
        echo "Checking for: ${partition}"
        update_partition "${partition}"
        if [ $? = 1 ]
        then
            echo "Something went wrong. Aborting..."
            exit ${RET_ERR}
        fi
    done
}


cleanup() {
    if findmnt "${MOUNT_POINT}"
    then
        umount "${MOUNT_POINT}"
    fi
}

trap cleanup EXIT


check_programs
check_root
download
update_efis