#!/bin/sh
# Copyright (C) 2016 Arista Networks, Inc.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# Aboot stage 0 boot script
#
# This boot0 script can be used in different scenario
#  - Installation and boot from Aboot (manual or reboot on a new image)
#  - EOS to SONiC fast-reboot (installation and kexec in EOS)
#  - SONiC to SONiC fast-reboot (installation and kexec in SONiC)
#
# Use it this way:
#   - swipath=path/to/swi install=true boot0
#   - swipath=path/to/swi install=true force=true boot0
#   - swipath=path/to/swi kexec=true boot0
#
# The logic defaults to the first scenario but providing extra environment variable
# will affect the behavior of the script.
# The list of variables is maintained here
#   - KERNEL : relative path to the kernel to execute
#   - INITRD : relative path to the initrd to load
#   - CMDLINE : place to find the default kernel cmdline to use for the platform
#   - VERBOSE : setting it to 1 will enable debug traces
#
# By default the boot0 script will behave for an Aboot based behavior.
# Options can be provided to only run some features of this script.
#
# Extra kernel parameters can be provided at runtime by the user by adding them
# into a kernel-params file.

image_name="image-%%IMAGE_VERSION%%"
dockerfs="{{ FILESYSTEM_DOCKERFS }}"

do_not_clean="do-not-clean"
kernel_params="kernel-params"

aboot_machine="arista_unknown"

info() { printf "%04.2f: $@\n" "$(cut -f1 -d' ' /proc/uptime)"; }
err() { info "Error: $@"; }
warn() { info "Warning: $@"; }

# extract mount point from the swi path, e.g., /mnt/flash/sonic.swi --> /mnt/flash
if [ -z "$target_path" ]; then
    if [ -z "$swipath" ]; then
       err "target_path= is required when swipath= is not provided"
       exit 1
    fi
    target_path=$(df "$swipath" | tail -1 | tr -s " " | cut -d ' ' -f6)
fi
image_path="$target_path/$image_name"
hook_path="$image_path/platform/hooks"
data_path="$image_path/platform/data"

cmdline_base="$target_path/kernel-params-base"
cmdline_image="$image_path/kernel-cmdline"
boot_config="$target_path/boot-config"

bootconfigvars="KERNEL INITRD CONSOLESPEED PASSWORD NETDEV NETAUTO NETIP NETMASK NETGW NETDOMAIN NETDNS NETHW memtest"
flash_re=" /mnt/flash| /host"

# for backward compatibility with the sonic_upgrade= behavior
install="${install:-${sonic_upgrade:-}}"

parse_environment_config() {
    for n in ${bootconfigvars}; do
        eval v="\$$n"
        if [ "$v" ]; then
            echo "$n=$v"
        fi
    done
}

clean_flash() {
    ## Remove all the other unnecssary files except swi file, boot-config
    for f in $(ls -A $target_path); do
        if [ $f != "${swipath##*/}" ] &&
           [ $f != "boot-config" ] &&
           [ $f != "$kernel_params" ] &&
           [ $f != "$cmdline_base" ] &&
           [ $f != "aquota.user" ] &&
           [ $f != "old_config" ] &&
           [ $f != "minigraph.xml" ]
        then
            rm -rf "$target_path/$f"
        fi
    done
}

update_boot_config() {
    local key="$1"
    local value="$2"

    if grep "$key" "$boot_config" 2>&1 > /dev/null; then
        sed -i "s#^$key=.*\$#$key=$value#" "$boot_config"
    else
        echo "$key=$value" >> "$boot_config"
    fi
}

get_boot_config() {
    local key="$1"

    sed -n "s#^$key=##p" "$boot_config" || :
}

update_next_boot() {
    local default="$(get_boot_config SWI_DEFAULT)"

    if [ -z "$default" ]; then
        warn "boot-config has no variable SWI_DEFAULT"
    else
        update_boot_config SWI "$default"
    fi
}

extract_image() {
    mkdir -p "$image_path"

    info "Extracting swi content"
    ## Unzip the image except boot0 and dockerfs archive
    unzip -oq "$swipath" -x boot0 "$dockerfs" -d "$image_path"

    ## detect rootfs type
    rootfs_type=`grep " $target_path " /proc/mounts | cut -d' ' -f3`

    info "Extracting $dockerfs from swi"
    ## vfat does not support symbol link
    if [ "$rootfs_type" != "vfat" ]; then
        mkdir -p "$image_path/{{ DOCKERFS_DIR }}"

        if [ -n "$install" ]; then
            TAR_EXTRA_OPTION="--numeric-owner --warning=no-timestamp"
        fi

        ## extract docker archive
        info "Unpacking $dockerfs"
        unzip -oqp "$swipath" "$dockerfs" | tar xzf - -C "$image_path/{{ DOCKERFS_DIR }}" $TAR_EXTRA_OPTION
    else
        ## save dockerfs archive in the image directory
        unzip -oq "$swipath" "$dockerfs" -d "$image_path"
        info "Unpacking $dockerfs delayed to initrd because $target_path is $rootfs_type"
    fi

    ## use new reduced-size boot swi
    local swi_boot_path="flash:$image_name/{{ ABOOT_BOOT_IMAGE }}"
    update_boot_config SWI "$swi_boot_path"
    update_boot_config SWI_DEFAULT "$swi_boot_path"

    ## Remove installer swi as it has lots of redundunt contents
    rm -f "$swipath"

    ## sync disk operations
    sync
}

write_machine_config() {
    ## Detect SKU and create a hardware description file
    aboot_version=$(grep ^Aboot "$cmdline_base" | sed 's/^.*norcal.-//' | tail -n 1)
    if [ -x /bin/sysinit ]; then
       aboot_build_date=$(stat -c %y /bin/sysinit | sed 's/ /T/')
    else
       aboot_build_date="unknown"
    fi
    cat <<EOF > ${target_path}/machine.conf
aboot_version=$aboot_version
aboot_vendor=arista
aboot_platform=x86_64-$aboot_machine
aboot_machine=$aboot_machine
aboot_arch=x86_64
aboot_build_date=$aboot_build_date
EOF
    chmod a+r "${target_path}/machine.conf"
}

platform_specific() {
    local platform="$(grep -Eo 'platform=[^ ]+' "$cmdline_base" | cut -f2 -d=)"
    local sid="$(grep -Eo 'sid=[^ ]+' "$cmdline_base" | cut -f2 -d=)"

    # set varlog size to 100MB
    local varlog_size=100

    # detect the size of the flash partition from name in Aboot/EOS/SONiC
    local flash_size=$(($(df | grep -E "$flash_re" | tr -s ' ' | cut -f2 -d' ') / 1000))

    if [ "$platform" = "raven" ]; then
        aboot_machine=arista_7050_qx32
        flash_size=2000
        echo "modprobe.blacklist=radeon,sp5100_tco acpi=off" >>/tmp/append
    fi
    if [ "$platform" = "crow" ]; then
        aboot_machine=arista_7050_qx32s
        flash_size=3700
        echo "modprobe.blacklist=radeon,sp5100_tco" >>/tmp/append
    fi
    if [ "$sid" = "Upperlake" ] || [ "$sid" = "UpperlakeES" ]; then
        aboot_machine=arista_7060_cx32s
        flash_size=3700
        echo "amd_iommu=off" >> /tmp/append
    fi
    if [ "$sid" = "Gardena" ] || [ "$sid" = "GardenaSsd" ]; then
        aboot_machine=arista_7260cx3_64
        flash_size=28000
    fi
    if [ "$sid" = "Alhambra" ] || [ "$sid" = "AlhambraSsd" ]; then
        aboot_machine=arista_7170_64c
        flash_size=28000
        echo "hugepages=128" >> /tmp/append
    fi
    if [ "$sid" = "BlackhawkO" ]; then
        aboot_machine=arista_7060px4_32
        flash_size=28000
    fi
    if [ "$sid" = "BlackhawkDD" ]; then
        aboot_machine=arista_7060dx4_32
        flash_size=28000
    fi
    if [ "$platform" = "rook" ]; then
        echo "iommu=on intel_iommu=on tsc=reliable pcie_ports=native" >>/tmp/append
        echo "rhash_entries=1 usb-storage.delay_use=0" >>/tmp/append
        if [ -x /bin/readprefdl ]; then
           readprefdl -f /tmp/.system-prefdl -d > /mnt/flash/.system-prefdl
        elif [ -f /etc/prefdl ]; then
           cp /etc/prefdl /mnt/flash/.system-prefdl
           chmod a+r /mnt/flash/.system-prefdl
        fi
        echo "reassign_prefmem" >> /tmp/append
    fi

    if [ $flash_size -ge 28000 ]; then
        varlog_size=4096
    elif [ $flash_size -ge 3700 ]; then
        varlog_size=400
    fi

    echo "varlog_size=$varlog_size" >>/tmp/append
    # disable deterministic interface naming
    echo "net.ifnames=0" >>/tmp/append
}

get_uuid_for() {
    local dev="$1"

    if type lsblk 2>&1 > /dev/null; then
       lsblk "$dev" -n --output UUID
    elif type blkid 2>&1 > /dev/null; then
       blkid | grep "^$dev" | sed -n "s/^.* UUID=\"//p" | grep -Eo '[^"]+'
    fi
}

write_boot_configs() {
    if $in_aboot; then
        # generate the default kernel parameters for the platform
        echo "$append" > $cmdline_base
        cat /etc/cmdline | sed "/^\(${bootconfigvars// /\|}\|crashkernel\|loglevel\|ignore_loglevel\)\(\$\|=\)/d;/^\$/d" >> $cmdline_base
        parse_environment_config >> $cmdline_base
    elif [ ! -f "$cmdline_base" ]; then
        # some systems were started with other versions of this script and therefore
        # do not have the $cmdline_base file. we assume that we are on Sonic or EOS.
        cat /proc/cmdline | sed -E 's/^(.*) rw .*$/\1/' | tr ' ' '\n' > $cmdline_base
    fi

    cp $cmdline_base /tmp/append

    platform_specific
    echo "rw loop=$image_name/fs.squashfs loopfstype=squashfs apparmor=1 security=apparmor quiet" >> /tmp/append

    # Pass the MAC address to the new kernel as a command line parameter. This makes it
    # possible to restore the MAC address in the new kernel without requiring driver modifications.
    if [ -f /sys/class/net/ma1/address ]; then
        echo "hwaddr_ma1=$(cat /sys/class/net/ma1/address)" >> /tmp/append
    elif [ -f /sys/class/net/eth0/address ]; then
        echo "hwaddr_ma1=$(cat /sys/class/net/eth0/address)" >> /tmp/append
    else
        err "Management port not found."
    fi

    # use extra parameters from kernel-params hook if the file exists
    if [ -f "$target_path/$kernel_params" ]; then
        cat "$target_path/$kernel_params" >> /tmp/append
    fi

    # setting root partition if not overridden by kernel-params
    if ! grep -q "root=" /tmp/append; then
        rootdev="$(mount | grep -E "$flash_re" | cut -f1 -d' ')"
        rootfstype="$(mount | grep -E "$flash_re" | cut -f5 -d' ')"
        rootuuid="$(get_uuid_for $rootdev)"
        if [ -z "$rootuuid" ] || [ "$rootfstype" = "vfat" ] ; then
            echo "root=$rootdev" >> /tmp/append
        else
            echo "root=UUID=$rootuuid" >> /tmp/append
        fi
    fi

    mv /tmp/append $cmdline_image
    [ -e ${target_path}/machine.conf ] || write_machine_config
}

run_kexec() {
    local cmdline="$(cat $cmdline_image | tr '\n' ' ')"
    local kernel="${KERNEL:-$(find $image_path/boot -name 'vmlinuz-*' -type f | head -n 1)}"
    local initrd="${INITRD:-$(find $image_path/boot -name 'initrd.img-*' -type f | head -n 1)}"

    if ! $verbose; then
       # Start showing systemd information from the first failing unit if any.
       # systemd.show_status=false or quiet can be used to silence systemd entierly
       cmdline="$cmdline systemd.show_status=auto"
    fi

    kexec --load --initrd="$initrd" --append="$cmdline" "$kernel"

    [ -z "$testonly" ] || exit 0
    info "Kexecing..."
    kexec --exec
}

get_sorted_hooks() {
    echo $(find "$1" -name '[0-9][0-9]-*' -type f)
}

run_hooks() {
    if [ -d "$hook_path/$1" ]; then
        for hook in $(get_sorted_hooks "$hook_path/$1"); do
            if [ ! -z "$hook" ]; then
                info "Running hook $(basename $hook)"
                . "$hook"
            fi
        done
    fi
}

# In Aboot no option will be provided therefore these are the default values to use
in_aboot=true
do_clean=true
do_install=true
do_kexec=true

# prevent the flash from being cleaned if the do-not-clean file exists
if [ -f "$target_path/$do_not_clean" ]; then
    do_clean=false
fi

# Parse the cmdline options (used from EOS or from SONiC)
if [ ! -z "$install" ]; then
    # install from SONiC or EOS
    in_aboot=false
    do_clean=false
    do_kexec=false
elif [ ! -z "$kexec" ]; then
    # kexec from SONiC or EOS
    in_aboot=false
    do_install=false
    do_clean=false
elif [ $# -ne 0 ]; then
    echo "usage: $0 (see code)"
    exit 1
fi

# Verbosity can be defined by the caller, default to false otherwise
verbose=${verbose:-false}
if [ -f "$target_path/verbose-boot" ] ||
   [ "$(get_boot_config VERBOSE)" = "1" ] ||
   ! $in_aboot; then
   verbose=true
fi

# enable shell debug mode to get the most verbosity
if $verbose; then
   set -x
fi

# install the image if newer
if $do_install; then
    # we expect the swi to install to be a non empty file
    if [ ! -s "$swipath" ]; then
        err "The swipath= environment variable does not point to a valid SWI"
        exit 1
    fi

    # check the hash file in the image, and determine to install or just skip
    GIT_REVISION=$(unzip -p "$swipath" .imagehash)
    LOCAL_IMAGEHASH=$(cat $image_path/.imagehash 2>/dev/null || true)

    if [ "$GIT_REVISION" != "$LOCAL_IMAGEHASH" ] || [ ! -z "$force" ]; then
        if $do_clean; then
            info "Cleaning flash content $target_path"
            clean_flash
        fi

        info "Installing image under $image_path"
        extract_image

        info "Generating boot-config, machine.conf and cmdline"
        write_boot_configs

        run_hooks post-install
    else
        info "Using previously installed image"
    fi
fi

# chainloading using kexec
if $do_kexec; then
    run_hooks pre-kexec
    update_next_boot
    run_kexec
fi