#!/bin/bash
Usage() {
   cat <<EOF
${0##*/} out-directory release

By default, the cloud image for 'release' is downloaded from the
daily cloud image stream.

Additional arguments can change behavior of what is downloaded.
Some supported arguments:
   * cloud-daily: use daily stream [default]
   * cloud-release: use release stream rather than daily stream.
   * maas-release, maas-daily: download maas images instead of cloud-image
   * hwe-N: download the hwe kernel from release 'N' instead of default
     this is only respected if maas stream is used.
   * YYYYMMDD[.X] : download the specific build version

example:
 ${0##*/} xenial_dir xenial cloud-daily
EOF
}
TEMP_D=""

NO_IMAGE_AVAILABLE_RC=3

cleanup() {
   [ ! -d "${TEMP_D}" ] || rm -Rf "${TEMP_D}"
}
VERBOSITY=1
inargs() {
  local n=$1 x="";
  shift;
  for x in "$@"; do [ "$n" = "$x" ] && return 0; done;
  return 1
}
debug() { [ $1 -ge $VERBOSITY ] || return 0; shift; echo "$@" 1>&2; }
error() { echo "$@" 1>&2; }
fail() { [ $# -eq 0 ] || error "$@"; exit 1; }

VERBOSITY=0
TEMP_D=$(mktemp -d ${TMPDIR:-/tmp}/${0##*/}.XXXXXX) || exit 1
trap cleanup EXIT

arches=( ppc64el i386 amd64 )
releases=( $(ubuntu-distro-info --all) )
mbase="http://maas.ubuntu.com/images/ephemeral-v2"
cbase="http://cloud-images.ubuntu.com"
def_stream_type="cloud"
stream_src=""
curbase="$mbase"
def_rel="xenial"
vername=""

[ "$1" = "-h" -o "$1" = "--help" ] && { Usage; exit 0; }
[ $# -lt 2 ] && { Usage 1>&2; exit 1; }
out_d="$1"
shift

myarch=$(uname -m)
case "$myarch" in
   ppc64*) def_arch="ppc64el";;
   i?86) def_arch="i386";;
   x86_64) def_arch="amd64";;
   arm*) def_arch="armel";;
esac

pt=( )
for x in "$@"; do
   inargs "$x" "${arches[@]}" && pt[${#pt[@]}]=arch=$x && arch="$x" && continue
   inargs "$x" "${releases[@]}" && pt[${#pt[@]}]="release=$x" && release=$x && continue
   [ "$x" = "released" -o "$x" = "release" ] && x="releases"
   case "$x" in
      cloud-released|cloud-release) x="cloud-releases";;
      maas-released|maas-release) x="maas-releases";;
   esac

   case "$x" in
      maas|cloud) stream_type="$x";;
      cloud-daily|maas-daily|cloud-releases|maas-releases)
          stream_type="${x%-*}"
          stream_sub=${x#*-};;
      http://*) stream=$x;;
      hwe-*) subarch="$x";;
      subarch=*) subarch=${x#*=};;
      *=*) pt[${#pt[@]}]="$x";;
      *~*) pt[${#pt[@]}]="$x";;
      [0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9]|[0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9].[0-9])
         pt[${#pt[@]}]="version_name=$x"
         vername="$x";;
      *) fail "Confused by input token '$x'";;
   esac
done

[ -n "$stream_type" ] || stream_type="${def_stream_type}"
if [ -z "$stream" ]; then
    stream_sub=${stream_sub:-daily}
    case "$stream_type" in
        maas) stream="$mbase/$stream_sub";;
        cloud) stream="$cbase/$stream_sub";;
        *) fail "unknown stream_type '$stream_type'";;
    esac
fi

[ -n "$arch" ] || pt[${#pt[@]}]="arch=$def_arch"

if [ "$stream_type" = "maas" ]; then
    if [ -z "$subarch" ]; then
        t=${release#?}
        first_letter=${release%${t}}
        subarch="hwe-${first_letter}"
        error "selected subarch=${subarch} for $release"
    fi
    pt[${#pt[@]}]="subarch=$subarch"
    ofmt="%(sha256)s %(ftype)s %(subarch)s %(item_url)s"
else
    ofmt="%(sha256)s %(ftype)s %(item_url)s"
fi

OIFS="$IFS"
mkdir -p "$out_d"
case "$stream" in
   *.json) :;;
   *) keyring="/usr/share/keyrings/ubuntu-cloudimage-keyring.gpg";;
esac

needs=""
[ -n "$keyring" -a ! -f "$keyring" ] &&
   needs="${needs} ubuntu-cloudimage-keyring"

for pair in sstream-query:simplestreams tgt-admin:tgt; do
   cmd=${pair%:*}
   pkg=${pair#*:}
   command -v "$cmd" >/dev/null 2>&1 || needs="$needs $pkg"
done

if [ -n "${needs# }" ]; then
   error "missing dependencies"
   fail "sudo apt-get install -qy ${needs# }"
fi

qcmd=( sstream-query ${keyring:+"--keyring=$keyring"} ${vername:---max=1}
       --output-format="$ofmt" "${stream}" datatype="image-downloads" 
       "${pt[@]}" )
echo "${qcmd[@]}"
"${qcmd[@]}" > "${TEMP_D}/qout"
roots=$(awk '$2 == "root-image.gz" || $2 == "disk1.img" { print $2 }' \
        "${TEMP_D}/qout" | sort -u | wc -l)
if [ "$roots" = "0" ]; then
   error "No image available for $release in $stream${pt:+ (${pt[*]})}"
   exit ${NO_IMAGE_AVAILABLE_RC}
elif [ "$roots" != "1" ]; then
   error "query resulted in '$roots' root images. expect 1 and only 1"
   error "cmd was ${qcmd[*]}"
   cat "${TEMP_D}/qout"
   fail
fi

set -f
while read line; do
    set -- $line
    if [ "$stream_type" = "maas" ]; then
        sha256="$1"; ftype="$2"; subarch="$3"; url="$4"
    else
        sha256="$1"; ftype="$2"; url="$3"
    fi
    case "$ftype" in
        boot-kernel|boot-initrd|root-image.gz|disk1.img) :;;
        *) continue
    esac
    if [ "$ftype" = "root-image.gz" ]; then
       fname="$ftype"
       image_src="$out_d/$fname"
       image="$out_d/root-image"
    elif [ "$ftype" = "disk1.img" ]; then
       fname="disk.img.dist"
       image_src="$out_d/$fname"
       image="$out_d/disk.img"
    else
       fname="$subarch/$ftype"
    fi
    out_f="${out_d}/$fname"
    if [ -f "$out_f" ]; then
        debug 1 "reusing existing '$out_f' rather than downloading $url."
        continue
    fi
    [ -d "${out_f%/*}" ] || mkdir -p "${out_f%/*}"
    tmp="${out_f}.tmp"
    wget --progress=dot:mega -O "$tmp" "$url" || {
        error "Failed to download $url to $tmp."
        rm -f "$tmp"
        exit 3
    }
    debug 1 "checksumming $tmp from $url"
    out=$(sha256sum "$tmp") || {
        error "Failed to checksum $tmp"
        rm -f "$tmp"
        exit 3;
    }
    found=${out%  *}
    if [ "$found" != "$sha256" ]; then
        error "$tmp from $url found sha256 of $found. expected $sha256."
        ls -l "$tmp" 1>&2
        rm -f "$tmp"
        exit 3
    fi
    mv "$tmp" "$out_f" || fail "failed renaming $tmp to $out_f"
    echo "$fname" "$url" >> "$out_d/contents"
done < "${TEMP_D}/qout"

if [ ! -f "$image" ]; then
    case "$image_src" in
        */root-image.gz)
           debug 1 "decompressing $image_src"
           zcat "$image_src" > "$image.tmp" &&
             mv "$image.tmp" "$image" ||
             { rm -f "$image.tmp"; exit 1; }
           ;;
        */disk.img.dist)
           debug 1 "converting qcow2 in $image_src -> raw in $image"
           qemu-img convert -O raw "$image_src" "$image.tmp" &&
               mv "$image.tmp" "$image" ||
               { rm -f "$image.tmp"; exit 1; }
           ;;
    esac
fi

# vi: ts=4 expandtab
