#!/bin/bash

offset=$((2*1024*1024))
filesize=0
prog_req=(mkimage file)
tmp_files=()

function remove_tmp_files()
{
    for i in "${tmp_files[@]}"
    do
        rm -f "$i"
    done
}

function cleanup_and_return_err()
{
    rm -f $UBOOT_SOURCE $UBOOT_SCRIPT
    remove_tmp_files
    exit 1
}

function add_device_tree_kernel()
{
    local path=$1
    local addr=$2
    local size=$3

    echo "fdt mknod $path module$addr" >> $UBOOT_SOURCE
    echo "fdt set $path/module$addr compatible \"multiboot,kernel\" \"multiboot,module\"" >> $UBOOT_SOURCE
    echo "fdt set $path/module$addr reg <0x0 "$addr" 0x0 "$size">" >> $UBOOT_SOURCE
    echo "fdt set $path/module$addr bootargs \"console=ttyAMA0\"" >> $UBOOT_SOURCE
}

function add_device_tree_ramdisk()
{
    local path=$1
    local addr=$2
    local size=$3

    echo "fdt mknod $path module$addr" >> $UBOOT_SOURCE
    echo "fdt set $path/module$addr compatible \"multiboot,ramdisk\" \"multiboot,module\"" >> $UBOOT_SOURCE
    echo "fdt set $path/module$addr reg <0x0 "$addr" 0x0 "$size">" >> $UBOOT_SOURCE
}

function add_device_tree_passthrough()
{
    local path=$1
    local addr=$2
    local size=$3

    echo "fdt mknod $path module$addr" >> $UBOOT_SOURCE
    echo "fdt set $path/module$addr compatible \"multiboot,device-tree\" \"multiboot,module\"" >> $UBOOT_SOURCE
    echo "fdt set $path/module$addr reg <0x0 "$addr" 0x0 "$size">" >> $UBOOT_SOURCE
}

function device_tree_editing()
{
    local device_tree_addr=$1

    echo "fdt addr $device_tree_addr" >> $UBOOT_SOURCE
    echo "fdt resize 1024" >> $UBOOT_SOURCE
    echo "fdt set /chosen \#address-cells <0x2>" >> $UBOOT_SOURCE
    echo "fdt set /chosen \#size-cells <0x2>" >> $UBOOT_SOURCE
    echo "fdt set /chosen xen,xen-bootargs \"$XEN_CMD\"" >> $UBOOT_SOURCE
    echo "fdt mknod /chosen dom0" >> $UBOOT_SOURCE
    echo "fdt set /chosen/dom0 compatible \"xen,linux-zimage\" \"xen,multiboot-module\"" >> $UBOOT_SOURCE
    echo "fdt set /chosen/dom0 reg <0x0 "$dom0_kernel_addr" 0x0 "$dom0_kernel_size">" >> $UBOOT_SOURCE
    if test "$DOM0_RAMDISK"
    then
        echo "fdt set /chosen xen,dom0-bootargs \"$DOM0_CMD\"" >> $UBOOT_SOURCE
        if test $dom0_ramdisk_addr != "-"
        then
            echo "fdt mknod /chosen dom0-ramdisk" >> $UBOOT_SOURCE
            echo "fdt set /chosen/dom0-ramdisk compatible \"xen,linux-initrd\" \"xen,multiboot-module\"" >> $UBOOT_SOURCE
            echo "fdt set /chosen/dom0-ramdisk reg <0x0 "$dom0_ramdisk_addr" 0x0 "$dom0_ramdisk_size">" >> $UBOOT_SOURCE
        fi
    else
        echo "fdt set /chosen xen,dom0-bootargs \"$DOM0_CMD\"" >> $UBOOT_SOURCE
    fi

    i=0
    while test $i -lt $NUM_DOMUS
    do
        echo "fdt mknod /chosen domU$i" >> $UBOOT_SOURCE
        echo "fdt set /chosen/domU$i compatible \"xen,domain\"" >> $UBOOT_SOURCE
        echo "fdt set /chosen/domU$i \#address-cells <0x2>" >> $UBOOT_SOURCE
        echo "fdt set /chosen/domU$i \#size-cells <0x2>" >> $UBOOT_SOURCE
        echo "fdt set /chosen/domU$i memory <0x0 "${DOMU_MEM[$i]}">" >> $UBOOT_SOURCE
        echo "fdt set /chosen/domU$i cpus <"${DOMU_VCPUS[$i]}">" >> $UBOOT_SOURCE
        echo "fdt set /chosen/domU$i vpl011 <0x1>" >> $UBOOT_SOURCE
        add_device_tree_kernel "/chosen/domU$i" ${domU_kernel_addr[$i]} ${domU_kernel_size[$i]}
        if test "${domU_ramdisk_addr[$i]}"
        then
            add_device_tree_ramdisk "/chosen/domU$i" ${domU_ramdisk_addr[$i]} ${domU_ramdisk_size[$i]}
        fi
        if test "${domU_passthrough_dtb_addr[$i]}"
        then
            add_device_tree_passthrough "/chosen/domU$i" ${domU_passthrough_dtb_addr[$i]} ${domU_passthrough_dtb_size[$i]}
        fi
        i=$(( $i + 1 ))
    done
}

function add_size()
{
    local filename=$1
    local size=`stat -L --printf="%s" $filename`
    memaddr=$(( $memaddr + $size + $offset - 1))
    memaddr=$(( $memaddr & ~($offset - 1) ))
    memaddr=`printf "0x%X\n" $memaddr`
    filesize=$size
}

function load_file()
{
    local filename=$1

    local relative_path="$(realpath --relative-to=$PWD $filename)"

    echo "$LOAD_CMD $memaddr $relative_path" >> $UBOOT_SOURCE
    add_size $filename
}

function check_file_type()
{
    local filename=$1
    local type="$2"

    if [ ! -f $filename ]
    then
        echo "File $filename doesn't exist, exiting";
        cleanup_and_return_err
    fi

    # if file doesn't know what it is, it outputs data, so include that
    # since some executables aren't recongnized
    if [ "$type" = "executable" ]
    then
        type="executable\|data"
    # file in older distros (ex: RHEL 7.4) just output data for device
    # tree blobs
    elif [ "$type" = "Device Tree Blob" ]
    then
        type="Device Tree Blob\|data"
    fi

    file -L $filename | grep "$type" &> /dev/null
    if test $? != 0
    then
        echo Wrong file type "$filename". It should be "$type", exiting.
        cleanup_and_return_err
    fi
}

function check_compressed_file_type()
{
    local filename=$1
    local type="$2"

    if [ ! -f $filename ]
    then
        echo "File $filename doesn't exist, exiting";
        cleanup_and_return_err
    fi

    file -L $filename | grep "gzip compressed data" &> /dev/null
    if test $? == 0
    then
        local tmp=`mktemp`
        tmp_files+=($tmp)
        cat $filename | gunzip > $tmp
        filename=$tmp
    fi
    check_file_type $filename "$type"
}

function check_depends()
{
    for ((i=0; i<${#prog_req[@]}; i++))
    do
        if ! command -v ${prog_req[i]} > /dev/null
        then
            echo "Please install the needed program: ${prog_req[i]}."
            exit 1
        fi
    done
}

function print_help
{
    script=`basename "$0"`
    echo "usage:"
    echo "	$script -c CONFIG_FILE -t UBOOT_TYPE -d DIRECTORY [-o FILE]"
    echo "	$script -h"
    echo "where:"
    echo "	CONFIG_FILE - configuration file"
    echo "	UBOOT_TYPE can be:"
    echo "		sd - alias for \"mmc load 0:1\" for uboot load commands"
    echo "		scsi - alias for \"scsi load 0:1\" for uboot load commands"
    echo "		tftp - alias for \"tftpb\" for uboot load cammnds"
    echo "		< > - used for uboot load commands"
    echo "	DIRECTORY - root directory where the files of CONFIG_FILE are located"
    echo "	FILE - output filename for the uboot script and its source, overrides option in CONFIG_FILE"
    echo "	-h - prints out the help message and exits "
    echo "Defaults:"
    echo "	CONFIG_FILE=$cfg_file, UBOOT_TYPE=\"LOAD_CMD\" env var, DIRECTORY=$uboot_dir"
    echo "Example:"
    echo "	$script -c ../config -d ./build42 -t \"scsi load 1:1\""
}

while getopts ":c:t:d:ho:" opt; do
    case ${opt} in
    t )
        case $OPTARG in
        scsi )
            LOAD_CMD="load scsi 0:1"
            ;;
        sd )
            LOAD_CMD="load mmc 0:1"
            ;;
        tftp )
            LOAD_CMD="tftpb"
            ;;
        * )
            LOAD_CMD="$OPTARG"
            ;;
        esac
        ;;
    c )
        cfg_file=$OPTARG
        ;;
    d )
        uboot_dir=$OPTARG
        ;;
    o )
        UBOOT_SCRIPT_ARG=$OPTARG
        ;;
    h )
        print_help
        exit 0
        ;;
    * )
        echo "Unknown option, see \"$0 -h\""
        exit 1
        ;;
    esac
done
shift $((OPTIND -1))

if [ -z "$LOAD_CMD" ] || [ -z "$cfg_file" ] || [ -z "$uboot_dir" ]
then
    echo "Undefined arguments, see \"$0 -h\""
    exit 1
fi

check_depends

source "$cfg_file"

# CLI ARG overrides what's in the config file
if [ ! -z "$UBOOT_SCRIPT_ARG" ]
then
    UBOOT_SCRIPT="$UBOOT_SCRIPT_ARG".scr
    UBOOT_SOURCE="$UBOOT_SCRIPT_ARG".source
fi

if [ -z "$XEN_CMD" ]
then
    XEN_CMD="console=dtuart dtuart=serial0 dom0_mem=1G dom0_max_vcpus=1 bootscrub=0 vwfi=native sched=null"
fi

if [ -z "$DOM0_CMD" ]
then
    DOM0_CMD="console=hvc0 earlycon=xen earlyprintk=xen clk_ignore_unused"
fi
if [[ ! $DOM0_CMD =~ root= ]]
then
    if test -z "$DOM0_ROOTFS"
    then
        DOM0_CMD="$DOM0_CMD root=/dev/ram0"
    else
        DEV=${LOAD_CMD%:*}
        DEV=${DEV##* }
        PAR=${LOAD_CMD#*:}

        if [ -z "$DEV" ] || [ -z "$PAR" ]
        then
            echo "Could not parse device and partition."
            echo "Please make sure the load command is correct or manually set DOM0_CMD in the config file."
            exit 1
        fi

        PAR=$((PAR + 1))

        if [[ $LOAD_CMD =~ mmc ]]
        then
            DOM0_CMD="$DOM0_CMD root=/dev/mmcblk${DEV}p${PAR}"
        elif [[ $LOAD_CMD =~ scsi ]]
        then
            # converts number to a scsi device character
            DEV=$((DEV + 97))
            DEV=$(printf %x $DEV)
            DEV=$(printf "\x$DEV")

            DOM0_CMD="$DOM0_CMD root=/dev/sd${DEV}${PAR}"
        else
            echo "Only mmc and scsi are supported for automatically setting the root partition."
            echo "Manually set DOM0_CMD with the root device in the config file to bypass this."
            exit 1
        fi
    fi
fi
i=0
while test $i -lt $NUM_DOMUS
do
    if test -z "${DOMU_MEM[$i]}"
    then
        DOMU_MEM[$i]=512
    fi
    DOMU_MEM[$i]=$((${DOMU_MEM[$i]} * 1024))
    if test -z "${DOMU_VCPUS[$i]}"
    then
        DOMU_VCPUS[$i]=1
    fi
    if test "${DOMU_ROOTFS[$i]}"
    then
        echo "Cannot handle non-ramdisk rootfses for dom0less VMs."
        echo "DomUs with rootfses on disk need to be created from dom0 using xl."
        exit 1
    fi
    i=$(( $i + 1 ))
done

# the cd is needed so that the relative paths will match once we use
# tftp or move the files to a partition
cd "$uboot_dir"

rm -f $UBOOT_SOURCE $UBOOT_SCRIPT
memaddr=$(( $MEMORY_START + $offset ))
# 12582912 is 0xc00000, 12MB
if test $memaddr -lt 12582912
then
    memaddr="12582912"
fi
memaddr=`printf "0x%X\n" $memaddr`
uboot_addr=$memaddr
# 2MB are enough for a uboot script
memaddr=$(( $memaddr + $offset ))
memaddr=`printf "0x%X\n" $memaddr`

check_compressed_file_type $XEN "executable"
xen_addr=$memaddr
load_file "$XEN"

check_compressed_file_type $DOM0_KERNEL "executable"
dom0_kernel_addr=$memaddr
load_file $DOM0_KERNEL
dom0_kernel_size=$filesize

if test "$DOM0_RAMDISK"
then
    check_compressed_file_type $DOM0_RAMDISK "cpio archive"
    dom0_ramdisk_addr=$memaddr
    load_file "$DOM0_RAMDISK"
    dom0_ramdisk_size=$filesize
else
    dom0_ramdisk_addr="-"
fi

i=0
while test $i -lt $NUM_DOMUS
do
    check_compressed_file_type ${DOMU_KERNEL[$i]} "executable"
    domU_kernel_addr[$i]=$memaddr
    load_file ${DOMU_KERNEL[$i]}
    domU_kernel_size[$i]=$filesize
    if test "${DOMU_RAMDISK[$i]}"
    then
        check_compressed_file_type ${DOMU_RAMDISK[$i]} "cpio archive"
        domU_ramdisk_addr[$i]=$memaddr
        load_file ${DOMU_RAMDISK[$i]}
        domU_ramdisk_size[$i]=$filesize
    fi
    if test "${DOMU_PASSTHROUGH_DTB[$i]}"
    then
        check_compressed_file_type ${DOMU_PASSTHROUGH_DTB[$i]} "Device Tree Blob"
        domU_passthrough_dtb_addr[$i]=$memaddr
        load_file ${DOMU_PASSTHROUGH_DTB[$i]}
        domU_passthrough_dtb_size[$i]=$filesize
    fi
    i=$(( $i + 1 ))
done

check_file_type $DEVICE_TREE "Device Tree Blob"
device_tree_addr=$memaddr
load_file $DEVICE_TREE
device_tree_editing $device_tree_addr

# disable device tree reloation
echo "setenv fdt_high 0xffffffffffffffff" >> $UBOOT_SOURCE
echo "booti $xen_addr - $device_tree_addr" >> $UBOOT_SOURCE
mkimage -A arm64 -T script -C none -a $uboot_addr -e $uboot_addr -d $UBOOT_SOURCE "$UBOOT_SCRIPT" &> /dev/null

remove_tmp_files

memaddr=$(( $MEMORY_END - $memaddr - $offset ))
if test $memaddr -lt 0
then
    echo Error, not enough memory to load all binaries
    cleanup_and_return_err
fi

echo "Generated uboot script $UBOOT_SCRIPT, to be loaded at address $uboot_addr:"
echo "$LOAD_CMD $uboot_addr $UBOOT_SCRIPT; source $uboot_addr"