#!/bin/bash

PROGNAME=$(basename "$0")
FORCE=0
AUTO_MODE=
IMG_SIZE=
BOOTIMAGE_ISO=
GRAPHICAL_SUFFIX=
IMG_FILE=
AUTO_ISO=
DHCPV6C=yes
OAM_DEV=ens3
IPV4_GW_ADDR=
IPV6_GW_ADDR=
AWS_COMPATIBLE=0
declare -A PASSWORDS
: KVM=
KVM_OPTS=()
TEMPFILES_DIR=
SUDO=0
GRAPHICAL=0
TTY_SETTINGS=
RPM_ADDON_LIST=()

# Print out the help message
usage() {
    echo "\
Usage: $0 OPTIONS...
Create a QCOW2/QEMU image with StarlingX pre-installed

 -f,--force        overwrite output file if it exists

 -m,--mode={controller|aio|aio_lowlatency}
                   create a controller or an all-in-one/low latency system
                   (default: aio)

    --sudo         Use sudo to mount the ISO, rather than udisks

 -s,--size=nnnG    image file size, must end with "G" (default: 500G)

 -g,--graphical    create a graphical installation, rather than console

 -e,--oam-dev=OAM_DEV
                   OAM network device (default: ens3)

 -4,--ipv4         don't configure IPv6 in the generated image

 -w,--ipv4-default-gateway=GW_IPV4_ADDR
                   Add a default IPv4 route via this gateway address

 -W,--ipv6-default-gateway=GW_IPV6_ADDR
                   Add a default IPv6 route via this gateway address

 -p,--password=USER:PASSWORD
                   Unlock USER account and set its password in the generated
                   image.
                   USER must exist -- e.g., root, sysadmin.
                   This option may be repeated.

                   WARNING: this option is not recommended because the
                            password will be visible to anyone listing the
                            processes. Use \`--passwords-from' instead.

 -P,--passwords-from=PASSWORD_FILE
                   Unlock and set passwords of each user account from
                   PASSWORD_FILE, which must contain one or more lines
                   of the form

                      USER:PASSWORD

                   USERs must exist -- e.g., root, sysadmin.

 -S,--passwords-from-stdin
                   Same as \`--passwords-from=/dev/stdin'

 -i,--iso=BOOTIMAGE_ISO
                   use this iso file as input, it must have been generated
                   by build-iso with default options
                   (default: \$MY_WORKSPACE/export/bootimage.iso)

 -o,--output=IMG_FILE
                   output image file name
                   Default:
                      \$MY_WORKSPACE/export/stx_\${MODE}.qcow2)
                   Default with --graphical:
                      \$MY_WORKSPACE/export/stx_\${MODE}_graphical.qcow2)

 --aws
                   Prepare an image that can be loaded onto an AWS EC2
                   instance
 --addon
                   Specify additional rpms to add to the qcow2 image

ENVIRONMENT

 MY_REPO           source repo directory
 MY_WORKSPACE      build workspace directory
 KVM               path to kvm executable (default: auto)
"
}

# Delete temporary files
cleanup() {
   # QEMU changes terminal settings, restore them before exiting
   [[ -z $TTY_SETTINGS ]] || stty "$TTY_SETTINGS" <&1
   # remove temporary files
   rm -rf "$TEMPFILES_DIR"
   rm -f "$IMG_FILE.tmp"
}

# Clean up before exiting due to a signal
handle_sig() {
   trap - EXIT
   cleanup
   exit 1
}

# Clean up before normal exit
handle_exit() {
  local rv="$?"
  trap - EXIT
  cleanup
  exit $rv
}

# Print out an error message
error() {
    echo "$PROGNAME: error: $*" >&2
}

# Print out an error message and exit
die() {
    error "$*"
    exit 1
}

# Print out a command-line error message and exit
cmdline_error() {
    if [ "$#" -gt 0 ] ; then
        error "$*"
    fi
    echo "Type \`$0 --help' for more info." >&2
    exit 2
}

# Encrypt a password for /etc/passwd
encrypt_password() {
    export ARG="$1"
    python -c '
import crypt, os, binascii, sys
salt = binascii.b2a_hex(os.urandom (8)).decode("ascii")
encrypted = crypt.crypt (os.environ["ARG"], "$5$" + salt + "$")
print (encrypted)
' "$1"
    local status="$?"
    unset ARG
    [[ $status -eq 0 ]] || exit 1
}

# Save username/password to $PASSWORDS
save_password() {
    local passwd_str="$1"
    local error_prefix="$2"
    if [[ ! $passwd_str =~ : ]] ; then
        error "${error_prefix}expecting USER:PASSWORD"
        return 1
    fi
    local user="${passwd_str%%:*}"
    local passwd="${passwd_str#*:}"
    if [[ -z $user || -z $passwd ]] ; then
        error "${error_prefix}expecting USER:PASSWORD"
        return 1
    fi
    if [[ $user =~ [^a-zA-Z0-9._-] ]] ; then
        error "${error_prefix}username must only contain characters [a-zA-Z0-9._-]"
        return 1
    fi
    PASSWORDS[$user]="$passwd"
    return 0
}

# Read passwords from file or STDIN
read_passwords() {
    local filename="$1"
    local -i lineno=0
    local numchar="#"
    # Open password file or STDIN as file descriptor 3
    if [[ -z $filename || $filename == - ]] ; then
        filename=STDIN
        exec 3<&0 || exit 1
    else
        exec 3<"$filename" || exit 1
    fi
    while read line <&3 ; do
        let lineno++
        # skip empty lines and comments
        # ${numchar} is "#" to avoid tripping up VI's syntax highlighting
        if [[ ! $line =~  ^[[:space:]]*(${numchar}.*)?*$ ]] ; then
            save_password "$line" "$filename:$lineno: " || exit 1
        fi
    done
    # close file descriptor 3
    exec 3<&-
}

# Check if an IPv4 address is valid
is_ipv4_addr() {
    # make sure we have python
    python -c 'import socket' || exit 1
    # parse the address via python
    python -c 'import socket,sys;socket.inet_aton(sys.argv[1])' "$1" >/dev/null 2>&1
}

# Check if an IPv6 address is valid
is_ipv6_addr() {
    # make sure we have python
    python -c 'import socket' || exit 1
    # parse the address via python
    python -c 'import socket,sys;socket.inet_pton(socket.AF_INET6,sys.argv[1])' "$1" >/dev/null 2>&1
}

# find QEMU/KVM
find_kvm() {
    local kvm
    if [[ -n "$KVM" ]] ; then
        kvm=$(which "$KVM")
        [[ -n $kvm ]] || exit 1
    else
        for kvm_basename in qemu-kvm kvm ; do
            kvm=$(export PATH=$PATH:/usr/bin:/usr/libexec ; which $kvm_basename 2>/dev/null || :)
            [[ -n $kvm ]] && break || :
        done
        [[ -n $kvm ]] || die "unable to find kvm executable"
    fi
    KVM="$kvm"
    if [[ -c /dev/kvm ]] ; then
        KVM_OPTS+=("-enable-kvm")
    fi
}

# Perform setup work for an image to run on AWS
# Create config files for adding ENA driver module, network scripts, and for
# regenerating a generic initramfs image
add_aws_setup(){
    local ks_addon=$1
    AWS_OAM_IF=ens5
    AWS_MGMT_IF=ens6
    cat >>"$ks_addon" <<_END

# Comment out deprecated virtio by-path rules to avoid duplicate symlinks
sed -i 's/^\(KERNEL.*disk\/by-path\/virtio\)/#\1/' /usr/lib/udev/rules.d/60-persistent-storage.rules

cat >/etc/modules-load.d/ena.conf <<END
ena
END

cat >/etc/dracut.conf.d/add-ena.conf <<END
add_drivers+=" ena "
END

cat >/etc/dracut.conf.d/no-hostonly.conf <<END
hostonly="no"
END

cat >/etc/sysconfig/network-scripts/ifcfg-${AWS_OAM_IF} <<END
DEVICE=${AWS_OAM_IF}
BOOTPROTO=dhcp
ONBOOT=yes
TYPE=Ethernet
USERCTL=yes
PEERDNS=yes
DHCPV6C=yes
DHCPV6C_OPTIONS=-nw
PERSISTENT_DHCLIENT=yes
RES_OPTIONS="timeout:2 attempts:5"
DHCP_ARP_CHECK=no
END

cat >/etc/sysconfig/network-scripts/ifcfg-${AWS_MGMT_IF} <<END
DEVICE=${AWS_MGMT_IF}
BOOTPROTO=dhcp
ONBOOT=yes
TYPE=Ethernet
USERCTL=yes
PEERDNS=yes
DHCPV6C=yes
DHCPV6C_OPTIONS=-nw
PERSISTENT_DHCLIENT=yes
RES_OPTIONS="timeout:2 attempts:5"
DHCP_ARP_CHECK=no
END

if [ ! -d /var/tmp ]; then
    mkdir -m 1777 /var/tmp
fi

KERNEL_VERSION=\$(rpm -q kernel --qf '%{version}-%{release}.%{arch}')
/sbin/dracut -f /boot/initramfs-\$KERNEL_VERSION.img \$KERNEL_VERSION
_END
}

# Process command line
init() {
    local temp
    temp=$(getopt -o hf4w:W:e:p:P:Sm:gs:i:o: --long help,force,ipv4,ipv4-default-gateway:,ipv6-default-gateway:,oam-dev:,password:,passwords-from:,passwords-from-stdin,mode:,graphical,sudo,size:,iso:,output:,aws,addon: -n "$PROGNAME" -- "$@") || cmdline_error
    eval set -- "$temp"
    while true ; do
        case "$1" in
            -h|--help)
                usage
                exit 0
                ;;
            -f|--force)
                FORCE=1
                shift
                ;;
            -4|--ipv4)
                DHCPV6C=no
                shift
                ;;
            -w|--ipv4-default-gateway)
                is_ipv4_addr "$2" || cmdline_error "invalid IP address \`$2'"
                IPV4_GW_ADDR="$2"
                shift 2
                ;;
            -W|--ipv6-default-gateway)
                is_ipv6_addr "$2" || cmdline_error "invalid IP address \`$2'"
                IPV6_GW_ADDR="$2"
                shift 2
                ;;
            -e|--oam-dev)
                OAM_DEV="$2"
                shift 2
                ;;
            -P|--passwords-from)
                read_passwords "$2"
                shift 2
                ;;
            -S|--passwords-from-stdin)
                read_passwords -
                shift
                ;;
            -p|--password)
                save_password "$2" "invalid $1: " || cmdline_error
                shift 2
                ;;
            -m|--mode)
                [[ "$2" =~ ^(controller|aio|aio_lowlatency)$ ]] || cmdline_error "invalid --mode"
                AUTO_MODE="$2"
                shift 2
                ;;
            -g|--graphical)
                GRAPHICAL=1
                GRAPHICAL_SUFFIX=_graphical
                shift
                ;;
            --sudo)
                SUDO=1
                shift
                ;;
            -s|--size)
                [[ $2 =~ ^[0-9]{1,5}G$ ]] || cmdline_error "invalid --size"
                IMG_SIZE="$2"
                shift 2
                ;;
            -i|--iso)
                BOOTIMAGE_ISO="$2"
                shift 2
                ;;
            -o|--output)
                IMG_FILE="$2"
                shift 2
                ;;
            --aws)
                AWS_COMPATIBLE=1
                shift
                ;;
            --addon)
                RPM_ADDON_LIST+=("$2")
                shift 2
                ;;
            --)
                shift
                break
                ;;
            -?*)
                cmdline_error
                ;;
            *)
                break
                ;;
        esac
    done
    [[ $# -le 0 ]] || cmdline_error "too many arguments"

    # These are required
    [[ -n $MY_WORKSPACE ]] || die "MY_WORKSPACE is not set"
    [[ -n $MY_REPO ]] || die "MY_REPO is not set"

    # Defaults
    : ${AUTO_MODE:=aio}
    : ${IMG_SIZE:=500G}
    : ${BOOTIMAGE_ISO:=$MY_WORKSPACE/export/bootimage.iso}
    : ${IMG_FILE:=$MY_WORKSPACE/export/stx_${AUTO_MODE}${GRAPHICAL_SUFFIX}.qcow2}
}

# main
init "$@"

# make sure we clean up before exiting
trap handle_sig  INT TERM PIPE HUP
trap handle_exit EXIT

# make sure update-iso-centos.sh exists
UPDATE_ISO=$MY_REPO/stx/utilities/utilities/platform-util/scripts/update-iso-centos.sh
: <"$UPDATE_ISO" || exit 1

# make sure input ISO file exists
: <"$BOOTIMAGE_ISO" || exit 1

# make sure patch_build.sh exists
PATCH_BUILD=$MY_REPO/stx/update/extras/scripts/patch_build.sh
: <"$PATCH_BUILD" || exit 1

# find patch-iso
which patch-iso >/dev/null || exit 1

# find QEMU/KVM
find_kvm

# find qemu-img
which qemu-img >/dev/null || exit 1

# refuse to overwrite existing output file
if [[ -e "$IMG_FILE" ]] && [[ $FORCE -ne 1 ]] ; then
    die "output file $IMG_FILE already exist, delete it first or use --force"
fi

# which menu item to use?
menu_item=
case "$AUTO_MODE" in
    controller)     menu_item=0 ;;
    aio)            menu_item=2 ;;
    aio_lowlatency) menu_item=4 ;;
    *)              die "internal error" ;;
esac

# create a directory for temporary files
TEMPFILES_DIR=$(mktemp -d -t build_img.XXXXXXXX) || exit 1

# create an updated iso with the menu item pre-selected
auto_iso="$TEMPFILES_DIR/bootimage_${AUTO_MODE}${GRAPHICAL_SUFFIX}.iso"
rm -f "$auto_iso"
cmd=()
if [[ $SUDO == 1 ]] ; then
    cmd+=(sudo)
fi
cmd+=("$UPDATE_ISO" -i "$BOOTIMAGE_ISO" -o "$auto_iso" -d "$menu_item" -t 3)

if [[ $AWS_COMPATIBLE == 1 ]] ; then
    cmd+=(-p rdloaddriver=ena)
fi

# generate a kickstart add-on
ks_addon="$TEMPFILES_DIR/ks_addon.sh"
echo "#### start ks-addon.cfg" >"$ks_addon"
# configure $OAM_DEV
cat >>"$ks_addon" <<_END
# configure $OAM_DEV
uuid=\$(uuidgen)
cat >/etc/sysconfig/network-scripts/ifcfg-$OAM_DEV <<END
UUID=\$uuid
DEVICE=$OAM_DEV
NAME=$OAM_DEV
TYPE=Ethernet
PROXY_METHOD=none
BROWSER_ONLY=no
BOOTPROTO=dhcp
DEFROUTE=yes
IPV4_FAILURE_FATAL=no
IPV6INIT=yes
IPV6_AUTOCONF=no
IPV6_DEFROUTE=yes
IPV6_FAILURE_FATAL=no
IPV6_ADDR_GEN_MODE=stable-privacy
ONBOOT=yes
DHCPV6C=$DHCPV6C
END
_END

# Add default routes
if [[ -n "$IPV4_GW_ADDR" ]] ; then
    cat >>"$ks_addon" <<_END
# Add a default IPv4 route
echo "default via $IPV4_GW_ADDR dev $OAM_DEV metric 1" >/etc/sysconfig/network-scripts/route-$OAM_DEV
_END
fi
if [[ -n "$IPV6_GW_ADDR" ]] ; then
    cat >>"$ks_addon" <<_END
# Add a default IPv6 route
echo "default via $IPV6_GW_ADDR dev $OAM_DEV metric 1" >/etc/sysconfig/network-scripts/route6-$OAM_DEV
_END
fi

# Disable cloud-init networking if cloud-init is installed
cat >>"$ks_addon" <<_END
if [ -d /etc/cloud/cloud.cfg.d/ ]; then
    echo "network: {config: disabled}" > /etc/cloud/cloud.cfg.d/99-disable-networking.cfg
fi
_END

# Set passwords
for user in "${!PASSWORDS[@]}" ; do
    encrypted=$(encrypt_password "${PASSWORDS[$user]}")
    [[ $? -eq 0 ]] || exit 1
    cat >>"$ks_addon" <<_END
# set ${user}'s password
usermod -e '' -p '$encrypted' '$user' || exit 1
chage --inactive -1 --maxdays -1 --lastday \$(date '+%Y-%m-%d') '$user' || exit 1
_END
done

# Comment-out global_filter in lvm.conf
# The installer normally sets it to the installer hard drive's bus address,
# and LVM doesn't come up when booted in different emulation environment.
cat >>"$ks_addon" <<'_END'
# Comment-out global_filter in lvm.conf
sed -r -i 's!^(\s*)global_filter\s*=.*!\1# global_filter = [ "a|.*/|" ]!' /etc/lvm/lvm.conf
_END

# Change grub parameters to boot to graphical console.
# The installer sets these to use the serial port when we install
# in text mode.
if [[ $GRAPHICAL -eq 1 ]] ; then
    cat >>"$ks_addon" <<'_END'
# Boot in graphical mode
sed -r -i \
    -e '/^\s*GRUB_SERIAL_COMMAND=/       d' \
    -e '/^\s*GRUB_TERMINAL(_OUTPUT)?=/   s/=.*/="console"/' \
    -e '/^\s*GRUB_CMDLINE_LINUX=/        s/\bconsole=ttyS0,\S+/console=tty0/' \
    /etc/default/grub
if [ -d /sys/firmware/efi ] ; then
  grub2-mkconfig -o /boot/efi/EFI/centos/grub.cfg
else
  grub2-mkconfig -o /boot/grub2/grub.cfg
fi
_END
fi

# Add necessary setup work for an aws image to the ks_addon script
if [[ $AWS_COMPATIBLE == 1 ]] ; then
    add_aws_setup $ks_addon
fi

echo "#### end ks-addon.cfg" >>"$ks_addon"
cmd+=(-a "$ks_addon")

# execute update_iso.sh
echo "${cmd[@]}"
"${cmd[@]}" || exit 1

# patch the iso if additional rpms are specified
if [ ${#RPM_ADDON_LIST[@]} -gt 0 ] ; then
    # Patch build will drop the generated patch file into the current directory.
    # We want that to be $MY_WORKSPACE.
    pushd $MY_WORKSPACE
    patch_file="PATCH.img-addon"
    patched_iso="$TEMPFILES_DIR/bootimage_${AUTO_MODE}${GRAPHICAL_SUFFIX}_patched.iso"
    cmd=("$PATCH_BUILD" --id "${patch_file}" --summary "additional packages for qcow2 image" --desc "Adds customizations to qcow2 image" --status "REL" --reboot-required "N")
    for rpm_addon in "${RPM_ADDON_LIST[@]}"; do
        cmd+=(--all-nodes "${rpm_addon}")
    done
    # create the patch file
    echo "${cmd[@]}"
    "${cmd[@]}" || exit 1
    cmd=(patch-iso -i "$auto_iso" -o "$patched_iso" "${MY_WORKSPACE}/${patch_file}.patch")
    # execute patch-iso
    echo "${cmd[@]}"
    "${cmd[@]}" || exit 1
    mv ${patched_iso} ${auto_iso}
    popd
fi

# create a blank image file
rm -f "$IMG_FILE.tmp"
cmd=(qemu-img create "$IMG_FILE.tmp" -f qcow2 "$IMG_SIZE")
echo "${cmd[@]}"
"${cmd[@]}" || exit 1

# run the installer in QEMU
cmd=(
    "$KVM"
    "${KVM_OPTS[@]}"
    -m 8192
    -drive file="$IMG_FILE.tmp",if=ide
    -cdrom "$auto_iso"
    -boot d
    -no-reboot
    -nographic
    -smp 4
)
# if STDOUT is a terminal, save current terminal settings
# so that we can restore them later
if [[ -t 1 ]] ; then
    TTY_SETTINGS=$(stty -g <&1)
# otherwise, disable QEMU's terminal features
else
    cmd+=(-serial file:/dev/stdout)
fi
# execute qemu
echo "${cmd[@]}"
"${cmd[@]}" 2>&1 | tee $TEMPFILES_DIR/kvm.log
if [[ ${PIPESTATUS[0]} -ne 0 || ${PIPESTATUS[1]} -ne 0 ]] ; then
    die "qemu: installation failed"
fi

# QEMU exits with status=0 even when killed by a signal. Check its output
# for a known message to detect this case
if tail "$TEMPFILES_DIR/kvm.log" | grep -q -E "(qemu|kvm).*: terminating on signal" ; then
    die "qemu terminated by a signal"
fi

# rename tmp image file to the final name
mv -f "$IMG_FILE.tmp" "$IMG_FILE" || exit 1

# done
echo "
Created $IMG_FILE

To use this image, type:
"
if [[ $GRAPHICAL -eq 1 ]] ; then
    echo "    $KVM ${KVM_OPTS[@]} -m 16384 -drive file=$IMG_FILE,if=ide -boot c -smp 4"
    echo
    echo "(requires a graphical console)"
else
    echo "    $KVM ${KVM_OPTS[@]} -m 16384 -drive file=$IMG_FILE,if=ide -boot c -nographic -smp 4"
fi