#!/bin/bash

VERBOSITY=0
TEMP_D=""
DEF_WPORT=32600
IPV4_NET="10.0.12"
DEF_NETDEV="user,net=${IPV4_NET}.0/24,host=${IPV4_NET}.2"

ISCSI_PORT=3260
TGT_CONF=""
HTTP_PID=""
TGT_NAME=""

_DEF_CMDLINE_TMPL=(
    nomodeset {icmdline_parms} {cmdline_ip} ro
    net.ifnames=0 BOOTIF_DEFAULT=eth0
    root={root} {overlay_drive}
    console=ttyS0 {cmdline_ds}
    # LP: #1796137
    systemd.mask=snapd.seeded.service systemd.mask=snapd.service
)
DEF_CMDLINE_TMPL="${_DEF_CMDLINE_TMPL[*]}"

error() { echo "$@" 1>&2; }
fail() { [ $# -eq 0 ] || error "$@"; exit 1; }

cleanup() {
    [ ! -d "$TEMP_D" ] || rm -Rf "$TEMP_D"
    if [ -f "$TGT_CONF" ]; then
       error "cleaning up tgt mount ${TGT_NAME}"
       sudo tgt-admin --force --delete="$TGT_NAME"
       sudo rm -f "$TGT_CONF"
    fi
    if [ -n "$HTTP_PID" ]; then
       kill $HTTP_PID
    fi
}

Usage() {
    cat <<EOF
Usage: ${0##*/} [ options ] image kernel initrd

   put 'image' into tgt, boot 'kernel' and 'initrd' in a kvm
   guest pointing to that iscsi target.

   options:
   -p | --webserv-port PORT  serve cloud-init seed on port PORT
                             default: ${DEF_WPORT}
   -O | --overlay-disk DISK  attach DISK to guest and use it for the overlayroot
                             target.  Use 'temp' to use a temp file.
                             use an existing file to use that.  Use a
                             non-existing file to create.
   -n | --netdev DEV         passed through to xkvm.
                             if none passed, default is:
                               $DEF_NETDEV
   -a | --host-addr ADDR     will be used for iscsi host address and
                             cloud-init seed address.  Default is read
                             from 'host=' in --netdev arg.
   -l | --serial-log FILE    write serial log to FILE default is to let output
                             go to console.
        --cmdline    TMPL    template for kernel command line.
                             default: ${DEF_CMDLINE_TMPL}
        --cmdline-add ARGS   add the args to the kernel command line. default: ""
        --user-data  FILE    user-data to provide to instance.
                             default: builtin
        --user-data-add FILE append FILE to default user-data.
EOF
}

bad_Usage() { Usage 1>&2; [ $# -eq 0 ] || error "$@"; return 1; }

debug() {
    local level=${1}; shift;
    [ "${level}" -gt "${VERBOSITY}" ] && return
    error "${@}"
}

write_userdata() {
    ## set up a user-data and meta-data for cloud-init
    ## and we'll seed that through the kernel command line
    ##
    ## The user data we provide will allow us to log in
    ## as the 'ubuntu' user with 'passw0rd'.
    local out="$1" keys="${2:-[]}"
    cat > "$out" <<"EOF"
#cloud-config
chpasswd:
  expire: False
  list: |
    ubuntu:passw0rd
    root:root
ssh_pwauth: True
EOF
}

get_ssh_pubkeys() {
    # get users ssh pubkeys in a yaml list format
    local keys="" d=""
    keys=$( (ssh-add -L 2>/dev/null ;
            [ -f ~/.ssh/id_rsa.pub ] && cat ~/.ssh/id_rsa.pub ) | sort -u )
    if [ -n "$keys" ]; then
        d=$(IFS=$'\n'; for i in ${keys}; do echo "'$i',"; done)
        keys="[${d%,}]"
    else
        keys="[]"
    fi
    _RET="$keys"
}

write_metadata() {
    local iid="" out="$1" pubkeys="$2"
    iid=$(uuidgen 2>/dev/null || echo i-abcdefg)
    {
    echo "instance-id: $iid"
    if [ -n "$pubkeys" ]; then
        echo "public-keys: $pubkeys"
    fi
    } > "$out"
}

register_image() {
    local conf="$1" name="$2" image="$3"
    sudo tee "$conf" >/dev/null <<EOF
<target ${name}>
    readonly 1
    backing-store "$image"
    allow-in-use yes
</target>
EOF
    [ $? -eq 0 ] || { error "failed to write '$conf'"; return 1; }

    sudo tgt-admin "--update=$name" ||
        { error "failed tgt-admin update=$name"; return 1; }
}

address_on_dev() {
    local dev="$1" atype="${2:-inet}" out="" addr=""
    out=$(ip address show dev "$dev" 2>/dev/null) ||
        { error "could not get address on $1"; return 1; }
    addr=$(echo "$out" |
        awk '$1 == atype && $4 == "global" {
                sub(/\/.*/, "", $2); print $2; exit(0); }' "atype=$atype" )
    [ -n "$addr" ] || return
    _RET="$addr"
}

find_host_addr() {
    # run through --netdev args and fine 'host=' or make one up base
    # on there being a net= and assume .2 for host.
    # example: type=user,id=net00,net=10.0.12.0/24,host=10.0.12.2
    #   ipv6-net=fec0::/64,ipv6-host=fec0::2
    #   lxcbr0
    local nd="" atype="" bridge=""
    local oifs="$IFS"
    for nd in "$@"; do
        local host="" netaddr="" cur="" ipv4="x" ipv6="x" ntype=""
        bridge=${nd%%,*}
        if [ "$bridge" != "user" -a -e "/sys/class/net/$bridge" ]; then
            for atype in inet inet6; do
                address_on_dev "$bridge" "$atype" || continue
                debug 2 "took $atype address on bridge '$bridge': $_RET";
                return 0;
            done
        fi
        cur=${nd}
        while :; do
            first=${cur%%,*}
            rest=${cur#*,}
            case "$first" in
                user) ntype=user;;
                ipv4=on|ipv4) ipv4=on;;
                ipv4=off) ipv4=off;;
                ipv6=on|ipv6) ipv6=on;;
                ipv6=off) ipv6=off;;
                ipv6-net=*) tmp=${first#*=};
                    tmp=${tmp%/*}
                    netaddr="${tmp%.*}:2"
                    ;;
                net=*) tmp=${first#*=};
                    tmp=${tmp%/*}
                    netaddr="${tmp%.*}.2"
                    ;;
                host=*|ipv6-host=) host=${first#*=};;
            esac
            if [ "$first" = "$rest" ]; then
                # there was no ','
                break
            fi
            cur=$rest
        done
        local found="" msg=""
        if [ -n "$host" ]; then
            found="$host"
            msg="host= in '$nd'"
        elif [ -n "$netaddr" ]; then
            found="$netaddr"
            msg="second address on net= in '$nd'";
        elif [ "$ntype" = "user" ]; then
            case "$ipv4/$ipv6" in
                */on)
                    msg="known user net default host in ipv6 in '$nd'"
                    found="fec0::2";;
                x/x|on/*)
                    msg="known user net default host in ipv4 in '$nd'"
                    found="10.0.2.2";;
            esac
        fi
        if [ -n "$found" ]; then
            debug 2 "found host address '$found' via $msg"
            _RET="$found"
            return 0
        fi
    done
    return 1
}


get_rfc4173() {
    # protocol would normally be '6' (tcp but empty is fine)
    # https://tools.ietf.org/html/rfc4173
    # https://bugs.debian.org/cgi-bin/bugreport.cgi?bug=804162
    local server="$1" name="$2" port="$3" lun="$4" user="$5" pass="$6" proto="$7"
    local rfc="" auth="${user:+${user}${pass:+:${pass}}}"
    rfc="iscsi:"
    if [ -n "${user}" ]; then
        rfc="$rfc:${user}${pass:+:${pass}}"
    fi
    case "$server" in
        \[*\]) :;;
        *:*) server="[$server]";;
    esac
    rfc="${server}:$proto:${port}:${lun}:${name}"
    _RET="$rfc"
}

get_root_parm() {
    local server="$1" name="$2" port="${3:-$ISCSI_PORT}"
    local lun="$4" user="$5" pass="$6" proto="$7" part="$8"
    local ret=""
    # for ipv6 addresses, path just contains : for the server
    _RET="/dev/disk/by-path/ip-$server:$port-iscsi-${name}-lun-${lun}${part:+-part$part}"
}

get_cmdline_params() {
    local server="$1" name="$2" port="$3" lun="$4" user="$5" pass="$6"
    local proto="$7" initiator="$8" ret=""
    ret="${initiator:+iscsi_initiator=${initiator} }"
    ret="${ret}iscsi_target_name=$name "
    ret="${ret}iscsi_target_ip=$server "
    ret="${ret}${port:+iscsi_target_port=$port }"
    ret="${ret}${initiator:+iscsi_initiator=$initiator }"
    _RET=${ret% }
}

get_iscsi_cmd() {
    local server="$1" name="$2" port="$3" lun="$4" user="$5" pass="$6"
    local proto="$7" initiator="$8" ret="" group=${9:-1}
    _RET=( iscsistart -i "$initiator"
            -t "$name" -g "$group" -a "$server" -p "$port"
            ${user:+-u "$user"} ${pass:+-w "${pass}"} )
}

serv_http() {
    local dir="$1" port="$2" check="$3"
    local python="" kid=""
    command -v python3 >/dev/null 2>&1 && python=python3 || python=python2
    cat > "$TEMP_D/webserv" <<EOF
import os, socket, sys
try:
    from BaseHTTPServer import HTTPServer
    from SimpleHTTPServer import SimpleHTTPRequestHandler
except ImportError:
    from http.server import HTTPServer, SimpleHTTPRequestHandler

class HTTPServerV6(HTTPServer):
  address_family = socket.AF_INET6

if __name__ == "__main__":
    dir = sys.argv[1]
    port = int(sys.argv[2])
    os.chdir(dir)
    server = HTTPServerV6(("::", port), SimpleHTTPRequestHandler)
    server.serve_forever()
EOF
    $python "$TEMP_D/webserv" "$dir" "$port" &
    HTTP_PID=$!
    local i=0 site="127.0.0.1"
    local url="http://$site:$port/$check"
    while [ -d /proc/$HTTP_PID -a "$i" -lt 10 ] && i=$(($i+1)); do
        no_proxy=$site wget -q -O /dev/null "$url" && return 0
        sleep .5
    done
    return 1
}

render_templ() {
    local tmpl="$1"
    shift
    local key val i="" cur="$tmpl"
    for i in "$@"; do
        key=${i%%=*}
        val=${i#*=}
        cur=${cur//\{${key}\}/${val}}
    done
    _RET="$cur"
}

get_partition() {
    # return in _RET the 'auto' partition for a image.
    # return partition number for a partitioned image
    # return 0 for unpartitioned
    # return 0 if image is partitioned, 1 if not
    local img="$1"
    out=$(LANG=C sfdisk --list -uS "$img" 2>&1) || {
        error "failed determining if partitioned: $out";
        return 1;
    }
    if echo "$out" | grep -q 'Device.*Start.*End'; then
        _RET=1
    else
        _RET=0
    fi
}

main() {
    local short_opts="a:hn:O:p:l:v"
    local long_opts="help,cmdline:,cmdline-add:,cmdline-ip:,disk:,host-addr:,overlay-disk:,netdev:,serial-log:,user-data:,user-data-add:,verbose,webserv-port:"
    local getopt_out=""

    getopt_out=$(getopt --name "${0##*/}" \
        --options "${short_opts}" --long "${long_opts}" -- "$@") &&
        eval set -- "${getopt_out}" ||
        { bad_Usage; return; }

    local cur="" next=""
    local overlay_disk="" webserv_port="${DEF_WPORT}" netdevs="" host_addr=""
    local serial_log="" cmdline_tmpl="${DEF_CMDLINE_TMPL}" user_data_file=""
    local cmdline_add="" user_data_add="" pt=""
    # The string 'BOOTIF' is replaced at runtime with the interface
    # name by cloud-initramfs-dyn-netconf.
    local cmdline_ip="ip=::::{iinitiator}:BOOTIF"
    netdevs=( )
    pt=( )

    while [ $# -ne 0 ]; do
        cur="$1"; next="$2";
        case "$cur" in
            -h|--help) Usage ; exit 0;;
               --cmdline) cmdline_tmpl=$next; shift;;
               --cmdline-add) cmdline_add=$next; shift;;
               --cmdline-ip) cmdline_ip=$next; shift;;
               --disk) pt[${#pt[@]}]="$1=$2"; shift;;
            -a|--host-addr) host_addr=$next; shift;;
            -O|--overlay-disk) overlay_disk=$next; shift;;
            -n|--netdev) netdevs[${#netdevs[@]}]="$next"; shift;;
            -v|--verbose) VERBOSITY=$((${VERBOSITY}+1));;
            -l|--serial-log) serial_log="$next"; shift;;
               --user-data-file) user_data_file="$next"; shift;;
               --user-data-add) user_data_add="$next"; shift;;
            -p|--webserv-port) webserv_port=$port; shift;;
            --) shift; break;;
        esac
        shift;
    done

    [ $# -eq 3 ] || { bad_Usage "expected 3 args, got $# ($*)"; return 1; }
    local image_in=$1 kernel_in=$2 initrd_in=$3
    for f in "$image_in" "$kernel_in" "$initrd_in"; do
        [ -f "$f" ] || fail "$f: not a file"
    done

    local image="" kernel="" initrd=""
    image=$(readlink -f "$image_in") ||
        fail "cannot get full path to $image_in"
    kernel=$(readlink -f "$kernel_in") ||
        fail "cannot get full path to $kernel_in"
    initrd=$(readlink -f "$initrd_in") ||
        fail "cannot get full path to $initrd_in"

    TEMP_D=$(mktemp -d "${TMPDIR:-/tmp}/${0##*/}.XXXXXX") ||
        fail "failed to make tempdir"
    TEMP_D=$(cd "$TEMP_D" && pwd)
    trap cleanup EXIT

    local t_image t_kernel t_initrd
    t_image="${TEMP_D}/image"
    t_kernel="${TEMP_D}/kernel"
    t_initrd="${TEMP_D}/initrd"

    if [ "${#netdevs[@]}" -eq 0 ]; then
        netdevs=( "$DEF_NETDEV" )
    fi

    local keys=""
    mkdir -p "$TEMP_D/ci-seed"
    if [ -n "${user_data_file}" ]; then
        cat "${user_data_file}" > "$TEMP_D/ci-seed/user-data" ||
            fail "failed to copy '$user_data_file' to ci-seed/user-data"
    else
        write_userdata "$TEMP_D/ci-seed/user-data" ||
            fail "failed to write user-data"
    fi
    if [ -n "$user_data_add" ]; then
        cat "$user_data_add" >> "$TEMP_D/ci-seed/user-data" ||
            fail "failed to write additional user-data"
    fi

    get_ssh_pubkeys || fail "failed getting pubkeys"
    keys="$_RET"
    write_metadata "$TEMP_D/ci-seed/meta-data" "$keys" ||
        fail "failed writing metadata"

    if [ -z "$host_addr" ]; then
        find_host_addr "${netdevs[@]}" ||
            fail "failed to find host address. try --host-addr."
        host_addr="$_RET"
    fi
    debug 1 "using host address: '${host_addr}'"

    ## run a python web server to seed this
    serv_http "$TEMP_D/ci-seed" "$webserv_port" "user-data" ||
        fail "failed to web serv $TEMP_D/ci-seed on $webserv_port"
    debug 1 "web server in pid $HTTP_PID serving ${TEMP_D}/ci-seed/ on port $webserv_port"
    local seed_url="http://$host_addr:$webserv_port/"
    if [ "${host_addr#*:}" != "$host_addr" ]; then
        seed_url="http://[$host_addr]:$webserv_port/"
    fi
    debug 1 " ${seed_url}meta-data"

    if [ "${image%.gz}" != "$image" ]; then
        error "uncompressing $image_in -> temp-dir"
        gzcat "$image" > "${TEMP_D}/image" ||
        fail "failed uncompress of $image_in"
    else
        ln -s "$image" "$t_image" || fail "symlink to $image failed"
    fi

    ln -s "$kernel" "$t_kernel" &&
        ln -s "$initrd" "$t_initrd" ||
        fail "failed link to kernel or initrd"

    local partition=""
    get_partition "$t_image" ||
        fail "failed to read partition from $image"
    partition=${_RET}
    [ "$part" = "0" ] && partition=""

    TGT_NAME=${TEMP_D##*/}
    TGT_NAME=${TGT_NAME//./-}
    TGT_CONF="/etc/tgt/conf.d/$TGT_NAME.conf"
    register_image "$TGT_CONF" "$TGT_NAME" "$t_image" ||
        fail "failed to register image $t_image in tgt"

    overlay_drive_kernel="overlayroot=tmpfs"
    overlay_drive_xkvm=""
    if [ -n "${overlay_disk}" ]; then
        if [ "${overlay_disk}" = "temp" ]; then
            overlay_disk="${TEMP_D}/overlay.img"
        fi
        if [ ! -f "$overlay_disk" ]; then
            qemu-img create -f raw "$overlay_disk" 1G
            mkfs.ext2 -L "delta" -F "$overlay_disk" >/dev/null 2>&1 ||
                fail "failed mkfs on $overlay_disk"
            overlay_drive_kernel="overlayroot=device:dev=LABEL=delta"
            overlay_drive_xkvm="--disk=${overlay_disk},format=raw"
        fi
    fi

    local iserver="${host_addr}" iname="$TGT_NAME" iport="$ISCSI_PORT"
    local ilun="1" iuser="" ipass="" iproto="" iinitiator="maas-enlist"
    local rfc="" root="" iscsi_cmdline_parms="" iscsi_start_cmd=""
    local cmdline=""
    get_rfc4173 "${iserver}" "$iname" "$iport" "$ilun" \
        "$iuser" "$ipass" "$iproto"
    rfc="$_RET"

    get_cmdline_params "${iserver}" "$iname" "$iport" "$ilun" \
        "$iuser" "$ipass" "$iproto" "$iinitiator"
    iscsi_cmdline_parms="$_RET"

    get_root_parm "$iserver" "$iname" "$iport" "$ilun" \
        "$iuser" "$ipass" "$iproto" "$partition"
    root="$_RET"

    get_iscsi_cmd "$iserver" "$iname" "$iport" "$ilun" \
        "$iuser" "$ipass" "$iproto" "$iinitiator"
    iscsi_start_cmd=( "${_RET[@]}" )

    render_templ "${cmdline_ip}" iinitiator="$iinitiator" ||
        { error "failed rendering cmdline ip: ${cmdline_ip}"; return 1; }
    cmdline_ip="$_RET"

    render_templ "$cmdline_tmpl${cmdline_add:+ ${cmdline_add}}" \
        iserver="$host_addr" iname="$TGT_NAME" iport="$ISCSI_PORT" \
        ilun="$ilun" iuser="$iuser" ipass="$ipass" iproto="$iproto" \
        iinitiator="$iinitiator" "rfc4173=$rfc" "root=$root" \
        "seed_url=$seed_url" icmdline_parms="$iscsi_cmdline_parms" \
        "overlay_drive=${overlay_drive_kernel}" \
        "cmdline_ds=ds=nocloud-net;seedfrom=${seed_url}" \
        "cmdline_ip=${cmdline_ip}" ||
        { error "failed rendering cmdline"; return 1; }
    cmdline="$_RET"

    debug 1 "rfc4173: $rfc"
    debug 1 "iscsi_cmdline_parms: $iscsi_cmdline_parms"
    debug 1 "root: $root"
    debug 1 "iscsi_start_cmd: ${iscsi_start_cmd[*]}"
    debug 1 "cmdline=$cmdline"

    local netdev_args=""
    netdev_args=( )
    for i in "${netdevs[@]}"; do
        netdev_args[${#netdev_args[@]}]="--netdev=$i"
    done

    local mem="512" start="$SECONDS" ret=""
    local cmd=""
    cmd=(
        ${BOOT_TIMEOUT:+timeout --kill-after=1m --signal=TERM $BOOT_TIMEOUT}
        xkvm -v -v "${netdev_args[@]}" ${overlay_drive_xkvm} "${pt[@]}" --
        -echr 0x5
        -m $mem ${serial_log:+-serial "file:$serial_log"} -nographic
        -kernel "${t_kernel}" -initrd "${t_initrd}"
        -append "${cmdline}"
    )
    debug 1 "executing: ${cmd[*]}"
    "${cmd[@]}"
    ret=$?
    debug 1 "xkvm returned $ret in $((SECONDS-start))s"
    return $ret
}

main "$@"
# vi: ts=4 expandtab
