#!/bin/sh
CMD="help create extension grub2"
BIN="ukit"
VERBOSE=0
create_tools_needed() {
    printf ""
}
create_usage() {
    printf "Helper of the create command\n"
    exit 2
}
create_helper() {
    create_usage
}
create_tools_needed() {
    print ""
}
create_exec() {
    printf "Execute command create\n"
}

EXTENSION_TYPE_DEFAULT="raw"
EXTENSION_FORMAT_DEFAULT="squashfs"
EXTENSION_PART_UUID="0fc63daf-8483-4772-8e79-3d69d8477de4"
EXTENSION_PART_LABEL="Linux filesystem"
EXTENSION_LIST_DEPS=""
EXTENSION_INITRD_RELEASE=""
_extension_deps_packages() {
    pkg="$1"
    found=0
    rpm_cmd=$(rpm -qR "$pkg")
    while read -r require; do
        dep_pkg=$(rpm -q --whatprovides "$require")
        if [ $? -eq 1 ]; then
            continue
        fi
        for elem in $EXTENSION_LIST_DEPS; do
            if [ "$elem" = "$dep_pkg" ]; then
                found=1
                break
            fi
        done
        if [ $found -eq 0 ]; then
            EXTENSION_LIST_DEPS="$EXTENSION_LIST_DEPS $dep_pkg"
            _extension_deps_packages "$dep_pkg"
        else
            found=0
        fi
    done <<rpm_cmd_input
$rpm_cmd
rpm_cmd_input
    # while IFS= read -r require; do
    #     dep_pkg=$(rpm -q --whatprovides "$require")
    #     if [ $? -eq 1 ]; then
    #         continue
    #     fi
    #     for elem in "${EXTENSION_LIST_DEPS[@]}"; do
    #         if [ "$elem" == "$dep_pkg" ]; then
    #             found=1
    #             break
    #         fi
    #     done
    #     if [[ $found -eq 0 ]]; then
    #         EXTENSION_LIST_DEPS+=("$dep_pkg")
    #         _extension_deps_packages "$dep_pkg"
    #     else
    #         found=0
    #     fi
    # done < <(rpm -qR "$pkg")
}
_extension_usage() {
    usage_str="$BIN extension [-n | --name] [-p | --package] [-f | --format ] \
[ -t | --type]
    -n|--name: Extension's name
    -p|--packages: List of packages to install into the extension
    -f|--format: Extension format (squashfs by default)
    -t|--type: Type of the extension (dir, raw)
    -u|--uki: Path to the referenced UKI (installed one by default)
    -a|--arch: Specify an architecture
        See https://uapi-group.org/specifications/specs/extension_image
        For the list of potential value.
    help: Print this helper
Info:
    Generate an extension for an UKI 'name-ext.format'
Example:
    $BIN extension -n \"debug\" -p \"strace,gdb\" -t \"raw\""
    printf "%s\n" "$usage_str"
}
_extension_size_partition() {
    fs_size=$1
    inode_size=$(head -n 7 /etc/mke2fs.conf\
        | grep "inode_size" | awk -F' = ' '{print $2}')
    inode_ratio=$(head -n 7 /etc/mke2fs.conf\
        | grep "inode_ratio" | awk -F' = ' '{print $2}')
    echo "$((fs_size*1000/(1000-(inode_size*1000)/(inode_ratio*1000)-50)))"
}
_extension_create() {
    if [ $# -lt 6 ]; then
        echo_debug "Missing arguments"
        return 1
    fi
    name="${1}"
    img_name="${name}-ext.raw"
    pkgs="$2"
    format="$3"
    type="$4"
    uki="$5"
    arch="$6"
    echo_info "Create the extension '$img_name' with '$pkgs' in format $format\
 at type $type for the uki $uki"
    for pkg in $(printf "%s" "$pkgs" | sed 's/,/ /g'); do
        if [ "$pkg_list" = "" ]; then
            pkg_list="$pkg";
        else
            pkg_list="$pkg_list $pkg"
        fi
    done
    # Check if all packages requires are installed
    for pkg in $pkg_list; do
        if ! rpm -q "$pkg" > /dev/null 2>&1; then
            echo_error "'$pkg' is not installed"
            exit 1
        fi
    done
    # Get dependencies of packages
    echo_info "Get all dependencies to install..."
    if [ "$EXTENSION_LIST_DEPS" = "" ]; then
        EXTENSION_LIST_DEPS="$pkg_list";
    else
        EXTENSION_LIST_DEPS="$EXTENSION_LIST_DEPS $pkg_list"
    fi
    for pkg in $pkg_list; do
        _extension_deps_packages "$pkg"
    done
    
    # Get list of files to install
    tmp_dir=$(mktemp -d)
    for pkg in $EXTENSION_LIST_DEPS; do
        for file in $(rpm -ql "$pkg" | sed 's/\n/ /g'); do
            if [ -f "$file" ]; then
                cp --parents "$file" "$tmp_dir"
            fi
        done
    done
    ext_name="${img_name%.*}"
    ext_dir=$tmp_dir/usr/lib/extension-release.d
    ext_file=$ext_dir/extension-release.$ext_name
    mkdir -p "$ext_dir"
    touch "$ext_file"
    id_arg=$(echo "$EXTENSION_INITRD_RELEASE" | grep "^ID=\"*\"")
    ver_id_arg=$(echo "$EXTENSION_INITRD_RELEASE" | grep "^VERSION_ID=\"*\"")
    {
        echo "SYSEXT_LEVEL=2"
        echo "$id_arg"
        echo "$ver_id_arg"
        echo "SYSEXT_ID=$name"
        # scope=[initrd,system,portable]
        echo "SYSEXT_SCOPE=initrd"
        echo "ARCHITECTURE=$arch"
    } > "$ext_file"
    # Create an empty disk raw image extensions
    sized=$(du -s --block-size=1M "$tmp_dir" | awk '{print $1}')
    sized=$((sized+1)) # Add at minimum 1M
    part_sized=$(_extension_size_partition ${sized})
    echo_info "Create an image of sized ${part_sized}M..."
    if [ "$format" = "$EXTENSION_FORMAT_DEFAULT" ]; then
        mksquashfs \
            "$tmp_dir" \
            "./$img_name" \
            -quiet
    else
        dd if=/dev/zero of="./$img_name" bs=1M count="$part_sized" \
            > /dev/null 2>&1
        mkfs."$format" \
            -U "$EXTENSION_PART_UUID" \
            -L "$EXTENSION_PART_LABEL" \
            -d "$tmp_dir" \
            -q \
            "./$img_name"
    fi
    # Clean
    [ "$tmp_dir" ] && rm -r "$tmp_dir"
    echo_info "extension image created at ./$img_name"
    return 0
}
extension_helper() {
    _extension_usage
}
extension_tools_needed() {
    printf "objcopy lsinitrd mksquashfs mktemp"
}
extension_exec() {
    [ $# -lt 2 ] \
        && echo_error "Missing arguments"\
        && _extension_usage && exit 2
    args=$(getopt -a -n extension -o n:p:f:t:u:a:\
        --long name:,packages:,format:,type:,uki:,arch: -- "$@")
    eval set --"$args"
    while :
    do
        case "$1" in
            -n | --name)        name="$2"     ; shift 2 ;;
            -p | --packages)    packages="$2" ; shift 2 ;;
            -t | --type)        type="$2"     ; shift 2 ;;
            -f | --format)      format="$2"   ; shift 2 ;;
            -u | --uki)         uki="$2"      ; shift 2 ;;
            -a | --arch)        arch="$2"     ; shift 2 ;;
            --)                 shift               ; break   ;;
            *) echo_warning "Unexpected option: $1"; _extension_usage   ;;
        esac
    done
    if [ ! ${packages+x} ]; then
        echo_error "Missing packages to install in the extension"
        _extension_usage
        exit 2
    fi
    if [ ! ${uki+x} ]; then
        count=0
        for uki_f in /usr/share/unified/efi/uki*.efi; do
            if [ -e "$uki_f" ]; then
                count=$((count +1))
                uki="$uki_f"
            fi
        done
        if [ "$count" -eq 0 ]; then
            echo_error "No UKI installed, please provides one"
            exit 2
        elif [ "$count" -ne 1 ]; then
            echo_error "More tahn one UKI installed, please select one"
            exit 2
        fi
    else
        if [ ! -f "$uki" ]; then
            echo_error "Cannot find the UKI at $uki"
            exit 2
        fi
    fi
    echo_info "Check the uki $uki and extract the initrd..."
    objcopy --dump-section .initrd=initrd-tmp "$uki"
    #EXTENSION_LSINITRD=$(lsinitrd ./initrd-tmp | grep " usr/")
    EXTENSION_INITRD_RELEASE=$(lsinitrd -f usr/lib/initrd-release ./initrd-tmp)
    rm initrd-tmp
    if [ ! ${arch+x} ]; then
        arch=$(printf "%s" "$(uname -m)" | sed 's/_/-/g')
    fi
    [ ! ${type+x} ] && type="$EXTENSION_TYPE_DEFAULT"
    if [ ! ${format+x} ]; then 
        format="$EXTENSION_FORMAT_DEFAULT"
    elif [ ! -f "/usr/sbin/mkfs.$format" ]; then
        echo_error "No mkfs.$format found, use another format"
        exit 1
    fi
    _extension_create "$name" "$packages" "$format" "$type" "$uki" "$arch"
    exit $?
}

GRUB2_CMD_ADD=1
GRUB2_CMD_REMOVE=2
GRUB2_CONFIG_INITRD="43_ukit_initrd"
GRUB2_CONFIG_UKI="44_ukit_uki"
GRUB2_CONFIG_FILE="/boot/grub2/grub.cfg"
GRUB2_EFI_DISTRO_DIR="/boot/efi/EFI/opensuse"
GRUB2_TRANSACTIONAL_UPDATE=0
_grub2_get_dev_name() {
    df -h "$1" | tail -1 | cut -d ' ' -f1
}
_grub2_get_dev_uuid() {
    blkid "$1" | sed -e 's|.* UUID="\(.*\)|\1|' | sed 's|" .*||'
}
_grub2_get_dev_avail() {
    df --block-size="1M" --output="avail" "$1" | tail -1 | tr -d ' '
}
_grub2_grub_cfg() {
    if [ "$GRUB2_TRANSACTIONAL_UPDATE" -eq 1 ]; then
        transactional-update grub.cfg
    else
        grub2-mkconfig -o /boot/grub2/grub.cfg
    fi
}
_grub2_remove_menuentry() {
    grub_config_path="$1"
    file_path="$2"
    file=$(basename "$file_path")
    if [ -f "$grub_config_path" ]; then
        if grep -q "$file_path" "$grub_config_path"; then
            echo_info "Removing menuentry for $file_path ..."
            # Get the block to remove:
            start_line=$(grep -n "menuentry.*$file' {" "$grub_config_path"\
                | cut -d':' -f1 | head -n 1)
            start_line=$((start_line-1))
            end_line=0
            while IFS= read -r line; do
                end_line=$((end_line+1))
                if expr "$line" : "^EOF" > /dev/null; then
                    [ "$end_line" -gt "$start_line" ] && break
                fi
            done < "$grub_config_path"
            sed -i "${start_line},${end_line}d" "$grub_config_path"
            _grub2_grub_cfg
        else
            echo_warning "There isn't a menu entry for $file_path"
            return
        fi
    else
        echo_warning "Grub config file not already created."
    fi
}
_grub2_usage() {
    usage_str="$BIN grub2 [--add-entry | --remove-entry] [-k | --kerver] \
[-i | --initrd ] [ -u | --uki]
    --add-entry|--remove-entry: Add/Remove grub2 entry (mandatory)
    -k|--kerver: Kernel Version (uname -r output by default)
    -i|--initrd: Path to the initrd
    -u|--uki: Path to the UKI
    help: Print this helper
Info:
    Create or remove an entry to the grub2 menu. If initrd argurment is \
provided, uki shouldn't, and vice versa.
    If the initrd provided isn't in the boot partition, it will copy it in \
/boot
    If the uki provided isn't in the the efi partition, it will copy it in \
$GRUB2_EFI_DISTRO_DIR
Example:
    $BIN grub2 --add-entry -k 6.3.4-1-default -u $GRUB2_EFI_DISTRO_DIR/uki.efi"
    printf "%s\n" "$usage_str"
}
_grub2_initrd() {
    cmd=$1
    kerver="$2"
    initrd_path="$3"
    root_dev="$(_grub2_get_dev_name /)"
    root_uuid="$(_grub2_get_dev_uuid "$root_dev")"
    grub_config_path="/etc/grub.d/$GRUB2_CONFIG_INITRD"
    eof="EOF"
    initrd_file=$(basename "$initrd_path")
    echo_debug "UUID root fs: $root_uuid"
    if [ "$cmd" -eq "$GRUB2_CMD_ADD" ]; then
        if [ ! -f "${initrd_path}" ]; then
            echo_error "Initrd not found at ${initrd_path}."
            exit 2
        fi
        if [ ! -f "/boot/$initrd_file" ]; then
            echo_info "$initrd_file isn't in boot partition, copy it to \
/boot/$initrd_file"
            if ! cp "$initrd_path" "/boot/$initrd_file"; then
                echo_error "Error when adding the initrd to the boot partition"
                exit 2
            fi
        fi
        initrd_path="/boot/$initrd_file"
        if [ -f "$grub_config_path" ]; then
            if grep -q "$initrd_path" "$grub_config_path"; then
                echo_warning "There is already a menu entry for $initrd_path"
                return
            fi
        else
            cat > $grub_config_path <<EOF
#!/bin/sh
set -e
EOF
            chmod +x $grub_config_path
        fi
        echo_info "Add initrd menuentry for $initrd_path ..."
        cat >> $grub_config_path <<EOF
cat << $eof
menuentry 'Linux ${kerver} and initrd ${initrd_file}' {
    load_video
    set gfxpayload=keep
    insmod gzio
    insmod part_gpt
    search --no-floppy --fs-uuid --set=root ${root_uuid}
    echo "Loading Linux ${kerver} ..."
    linux /boot/vmlinuz-${kerver} root=UUID=${root_uuid}
    echo "Loading ${initrd_path}..."
    initrd ${initrd_path}
}
$eof
EOF
        _grub2_grub_cfg
    elif [ "$cmd" -eq "$GRUB2_CMD_REMOVE" ]; then
        _grub2_remove_menuentry "$grub_config_path" "$initrd_path"
    fi
}
_grub2_uki() {
    cmd=$1
    uki_path="$2"
    efi_dev="$(_grub2_get_dev_name /boot/efi)"
    efi_uuid="$(_grub2_get_dev_uuid "$efi_dev")"
    grub_config_path="/etc/grub.d/$GRUB2_CONFIG_UKI"
    uki_file=$(basename "$uki_path")
    eof="EOF"
    echo_debug "UUID boot partition: $efi_uuid"
    if [ "$cmd" -eq "$GRUB2_CMD_ADD" ]; then
        if [ ! -f "${uki_path}" ]; then
            echo_error "Unified Kernel Image not found at ${uki_path}."
            exit 2
        fi
        if [ ! -f "$GRUB2_EFI_DISTRO_DIR/${uki_file}" ]; then
            echo_info "$uki_file isn't in efi partition, copy it to \
$GRUB2_EFI_DISTRO_DIR/$uki_file"
            efi_avail="$(_grub2_get_dev_avail "$efi_dev")"
            uki_size="$(du -m0 "$uki_path" | cut -f 1)"
            echo_info "${efi_avail}M available on efi partition"
            echo_debug "Size of uki file: ${uki_size}M"
            if [ "$uki_size" -gt "$efi_avail" ]; then
                echo_error "No space left on efi partition to install uki"
                echo_error "Need ${uki_size}M, Available: ${efi_avail}M"
                exit 2
            fi
            mkdir -p "$GRUB2_EFI_DISTRO_DIR"
            if ! cp "$uki_path" "$GRUB2_EFI_DISTRO_DIR/$uki_file"; then
                echo_error "Error when adding the uki to the EFI partition"
                exit 2
            fi
            
        fi
        uki_path="/EFI/opensuse/$uki_file"
        if [ -f "$grub_config_path" ]; then
            if grep -q "$uki_path" "$grub_config_path"; then
                echo_warning "There is already a menu entry for $uki_path"
                return
            fi
        else
            cat > $grub_config_path <<EOF
#!/bin/sh
set -e
EOF
            chmod +x $grub_config_path
        fi
        echo_info "Add UKI menuentry for $uki_path..."
        cat >> $grub_config_path <<EOF
cat << $eof
menuentry 'Unified Kernel Image ${uki_file}' {
    insmod part_gpt
    insmod btrfs
    insmod chain
    search --no-floppy --fs-uuid --set=root ${efi_uuid}
    echo "Loading unified kernel image ${uki_file} ..."
    chainloader /EFI/opensuse/${uki_file}
}
$eof
EOF
        _grub2_grub_cfg
    elif [ "$cmd" -eq "$GRUB2_CMD_REMOVE" ]; then
        _grub2_remove_menuentry "$grub_config_path" "/EFI/opensuse/$uki_file"
    fi
}
grub2_tools_needed() {
    printf ""
}
grub2_helper() {
    _grub2_usage
}
grub2_exec() {
    [ $# -lt 2 ] \
        && echo_error "Missing arguments"\
        && _extension_usage && exit 2
    [ ! -f "$GRUB2_CONFIG_FILE" ] \
        && echo_error "grub2 is not installed!" \
        && exit 2
    args=$(getopt -a -n extension -o k:i:u:\
        --long add-entry,remove-entry,kerver:,initrd:,uki: -- "$@")
    eval set --"$args"
    while :
    do
        case "$1" in
            --add-entry)        cmd_add=1         ; shift 1 ;;
            --remove-entry)     cmd_remove=1      ; shift 1 ;;
            -k | --kerver)      kerver="$2"       ; shift 2 ;;
            -i | --initrd)      initrd_path="$2"  ; shift 2 ;;
            -u | --uki)         uki_path="$2"     ; shift 2 ;;
            --)                 shift             ; break   ;;
            *) echo_warning "Unexpected option: $1"; _grub2_usage   ;;
        esac
    done
    # Check transactional update system
    if command -v transactional-update; then
        GRUB2_TRANSACTIONAL_UPDATE=1
    fi
    # Check the command
    if [ ! ${cmd_add+x} ] && [ ! ${cmd_remove+x} ]; then
        echo_error "Need \"add-entry\" or \"remove-entry\" command"
        _grub2_usage
        exit 2
    elif [ ${cmd_add+x} ] && [ ${cmd_remove+x} ]; then
        echo_error "Please choose between add or remove a menue entry. Not\
both!"
        _grub2_usage
        exit 2
    elif [ ${cmd_add+x} ]; then
        cmd=$GRUB2_CMD_ADD
    else
        cmd=$GRUB2_CMD_REMOVE
    fi
    # Check the mode
    if [ ${initrd_path+x} ] && [ ${uki_path+x} ]; then
        echo_error "Please choose between initrd or uki arguments. Not both!"
        _grub2_usage
        exit 2
    elif [ ! ${initrd_path+x} ] && [ ! ${uki_path+x} ]; then
        echo_error "Missing initrd path OR uki path to add to the menu entry"
        _grub2_usage
        exit 2
    elif [ ${uki_path+x} ]; then
        _grub2_uki $cmd "$uki_path"
    else
        # Check the kernel version
    	if [ ! ${kerver+x} ]; then
        	kerver=$(uname -r)
    	fi
    	if [ ! -f "/boot/vmlinuz-${kerver}" ]; then
        	echo_error "Unable to find the Kernel file: /boot/vmlinuz-${kerver}\
, wrong kernel version ?"
        	exit 2
   	    fi
        _grub2_initrd $cmd "$kerver" "$initrd_path"
    fi
}

TOOLS_NEEDED=""
echo_warning() {
    color="\033[0;33m"
    color_light="\033[1;33m"
    color_none="\033[0m"
    printf "%b[WARNING]%b %s%b\n" "${color}" "${color_light}" "$1" \
"${color_none}"
}
echo_error() {
    color="\033[0;31m"
    color_light="\033[1;31m"
    color_none="\033[0m"
    printf "%b[ERROR]%b %s%b\n" "${color}" "${color_light}" "$1" "${color_none}"
}
echo_info() {
    color="\033[0;32m"
    color_light="\033[1;32m"
    color_none="\033[0m"
    printf "%b[INFO]%b %s%b\n" "${color}" "${color_light}" "$1" "${color_none}"
}
echo_debug() {
    [ "$VERBOSE" -eq 0 ] && return
    color="\033[0;34m"
    color_light="\033[1;34m"
    color_none="\033[0m"
    printf "%b[DEBUG]%b %s%b\n" "${color}" "${color_light}" "$1" "${color_none}"
}
usage() {
    usage_str="$BIN [help] [verbose] COMMAND [help | COMMAND OPTION]
    - help: Print this helper
    - verbose: Print debug information to the output
    - COMMAND help: Print the helper of the command
    - COMMAND [OPTION]: Execute the command with additional options.
List of COMMAND:"
    for cmd in $CMD; do
        usage_str=$(printf "%s\n\t- %s" "$usage_str" "$cmd")
    done
    printf "%s\n" "$usage_str"
}
check_tools_needed() {
    for dep in $TOOLS_NEEDED; do
        if ! command -v "$dep" > /dev/null 2>&1; then
            if [ ${missing_deps+x} ]; then
                missing_deps="$missing_deps $dep"
            else
                missing_deps="$dep"
            fi
        fi
    done
    if [ ${missing_deps+x} ]; then
        echo_error "Some tools are missing on your system: $missing_deps"
        exit 1
    fi
}
if [ $# -lt 1 ]; then
    echo_error "Missing command"
    usage & exit 2
fi
cmd_in="$1"
if [ "$cmd_in" = "help" ]\
    || [ "$cmd_in" = "--help" ]\
    || [ "$cmd_in" = "-h" ]; then
        usage
        exit 0
elif [ "$cmd_in" = "verbose" ]\
    || [ "$cmd_in" = "-v" ]; then
    VERBOSE=1
    cmd_in="$2"
    shift 1
fi
found=0;
for cmd in $CMD; do
    if [ "$cmd" = "$cmd_in" ]; then
        found=1
        if [ "$2" = "help" ]\
            || [ "$2" = "--help" ]\
            || [ "$2" = "-h" ]; then
                "${cmd}_helper"
        else
            # Get dependencies of the command and check them
            TOOLS_NEEDED="$("${cmd}_tools_needed")"
            check_tools_needed
            # Exec the command
            "${cmd}_exec" "$@"
        fi
    fi
done
if [ $found -eq 0 ]; then
    echo_error "Unknown command \"$cmd_in\""
    usage & exit 1
fi
exit 0
