#!/usr/bin/bash


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

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


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

SHIM_LOCATION="EFI/boot/bootx64.efi"
MMX_LOCATION="EFI/boot/mmx.efi"
GRUB_LOCATION="EFI/boot/grubx64.efi"
GURB_CFG_LOCATION="boot/grub/grub.cfg"

RET_SKIP=2
RET_ERR=1

download() {
    echo "Downloading new files"
    if ! wget "${NEW_SHIM_URL}" -P "${TMP_DIR}" -q
    then
        echo "failed to download new shim"
        exit ${RET_ERR}
    fi
    if ! wget "${NEW_MMX_URL}" -P "${TMP_DIR}" -q
    then
        echo "failed to download new mok manager"
        exit ${RET_ERR}
    fi
    if ! wget "${NEW_GRUB_URL}" -P "${TMP_DIR}" -q
    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 wget  &> /dev/null
    then
        echo "You need to install wget 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

    if [ ! -f "${update_dir}/${GRUB_LOCATION}" ]
    then
        echo "No grub.cfg found"
        return ${RET_ERR}
    fi

    if grep -q "dejavu_sans_mono_12" "${update_dir}/${GURB_CFG_LOCATION}"
    then
        if ! sed -i 's#if loadfont /boot/grub/dejavu_sans_mono_12.pf2#if loadfont (memdisk)/fonts/unicode.pf2#g' "${update_dir}/${GURB_CFG_LOCATION}"
        then
            echo "Couldn't patch grub.cfg"
            return ${RET_ERR}
        fi    
    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