#!/bin/sh -efu

. shell-error

show_help(){
cat <<EOF

Helper for dealing with images for network installations.

    Usage:

$0 <options>

    Options:

 -h        -- show this help
 -l        -- list numbers and descriptions for all avalable images
 -L        -- list current pxelinux.cfg/default boot labels
 -c        -- show number of image which is currently used for installation
 -C        -- show number of default boot label
 -r <num>  -- show image info
 -s <num>  -- select image <num> for network installations
 -b <num>  -- select number of pxelinux.cfg/default label to boot from
 -A <url>  -- start downloading of new image
 -D        -- stop downloading
 -S        -- show downloading status

    Files:

Images are stored in  /srv/public/netinst/<num>.img
Some information about images is stored in  /srv/public/netinst/list
Information about last download is stored in  /srv/public/netinst/download
Current image is selected by /srv/public/netinst/current  symlink
Files for boot over network are installed in /var/lib/tftpboot

    Requirements for network installations

* /var/lib/tftpboot must be accessable from client via tftp;
* directory /srv/netinst must be accessable from client via nfs;
* dhcp options next-server filename root-path must be setup
  for client in this way:
    next-server <ip>;
    filename "pxelinux.0";
    option root-path "/srv/netinst/current"

  If alterator-dhcp is installed these settings are made
  automatically by this script.

EOF
}

DATADIR="/srv/public/netinst"

# we want to search directory in /proc/mounts,
# so we need to dereference possible symlinks:
[ ! -h "$DATADIR" -a ! -d "$DATADIR" ] || DATADIR="$(readlink -f "$DATADIR" ||:)"

LIST="$DATADIR/list"
CURR="$DATADIR/current"

DPID="$DATADIR/download/pid"
DURL="$DATADIR/download/url"
DNUM="$DATADIR/download/num"
DSIZE="$DATADIR/download/size"
DSTAT="$DATADIR/download/status"
DERR="$DATADIR/download/error"

MNT="$DATADIR/mnt"
MNTTMP="$DATADIR/mnt_tmp"
TFTPDIR="/var/lib/tftpboot"

show_list(){
  sed -n -e '/^[0-9]\+[[:space:]]/p' "$LIST" 2>/dev/null | cut -d'	' -f1,3 ||:
}

show_curr(){
  if cat /proc/mounts | cut -f2 -d' ' | grep -q "^$MNT\$" && [ -h "$CURR" ]; then
    num="$(readlink "$CURR")"
    basename "$num" ".img"
  else
    echo "0"
  fi
}

show_info(){
  local num url name
  sed -n \
       -e "/^$1[[:space:]]/ {p;q}" \
       -e "\$ a $1" \
       "$LIST" 2>/dev/null |
  while IFS='	' read num url name; do
    echo "num:   $num"
    echo "name:  $name"
    echo "url:   $url"
  done
}

show_labels(){
  local cfgfile="${1:-$TFTPDIR/pxelinux.cfg/default}"
  sed -n 's/^label[[:space:]]\+//p' < "$cfgfile" | cat -n | sed 's/^[[:space:]]*//'
}

show_default_label(){
  local cfgfile="${1:-$TFTPDIR/pxelinux.cfg/default}"
  sed -n 's/^default[[:space:]]\+//p' < "$cfgfile"
}

show_default_nlabel() {
  local deflabel="$(show_default_label "$@")"
  show_labels | sed -n "s/^\([0-9][0-9]*\)[[:space:]]\+$deflabel$/\1/p"
}

guess_default_label() {
  local cfgfile="${1:-$TFTPDIR/pxelinux.cfg/default}"
  local default
  for default in linux live rescue; do
    show_labels "$cfgfile" | egrep -q "^[0-9]+[[:space:]]+$default$" || continue
    break
  done
  test -n "$default" || default="$(show_labels "$cfgfile" | head -1 | sed 's/^[0-9]+[[:space:]]+//')"
  echo "$default"
}

set_default_label(){ # new_label_number
  local lname="$(show_labels | sed -n "s/^$1[[:space:]]\+\(.*\)/\1/p")"
  test -n "$lname" || lname="$(show_labels | sed -n "s/^1[[:space:]]\+\(.*\)/\1/p")"
  test -n "$lname" && sed -i "s/^\(default\)[[:space:]]\+.*/\1 $lname/" "$TFTPDIR/pxelinux.cfg/default"
}

##### downloading images

downloader(){
  local url="$1"

  echo ""          > "$DERR"
  echo ""          > "$DSIZE"
  echo "progress"  > "$DSTAT"
  echo "$url"      > "$DURL"

  if [ -z "$url" ]; then
    echo "fail" > "$DSTAT"
    echo "Error: empty url" &> /dev/stderr
    return
  fi

  local max="0"
  for n in $(ls -1 "$DATADIR" | grep "^[0-9]\+.img$"); do 
    [ "$((${n%.img} > $max))" = 0 ] || max=${n%.img};
  done
  local num="$(($max+1))"
  echo "$num"      > "$DNUM"

  [ -z "${url##*:*}" ] || url="file:$url"

  local size
  if [ "$url" = "cdrom:" ]; then
     dev="/dev/sr0"
     # wait for a CDROM
     eject -X "$dev" &> /dev/null
     size="$(cat /sys/block/${dev##*/}/size)"
     size="$(($size*512))"
  else
     size="$(curl -I -s -S "$url" | \
             sed -n 's/Content-Length:[[:space:]]\+\([0-9]\+\)[^0-9]\+/\1/p')"
  fi

  if [ -z "$size" ]; then
    echo "fail" > "$DSTAT"
    return
  fi

  echo "$size"     > "$DSIZE"
  mkdir -p "$MNTTMP" ||:

  if [ "$url" = "cdrom:" ]; then
    dd if=$dev of="$DATADIR/$num.img" &
    echo $! > "$DPID"          ## pid of dd
  else
    curl -s -S -o "$DATADIR/$num.img" "$url" &
    echo $! > "$DPID"          ## pid of curl
  fi

  wait "$(cat "$DPID")" 2>/dev/null
  local retcode="$?"

  local newsize="$(stat -c %s $DATADIR/$num.img 2>/dev/null || echo "0")"
  local status="done"

  if [ "$retcode" -gt 126 ]; then
    status="stopped"
  elif [ "$newsize" -le 0 -o "$(( $size - $newsize ))" -gt 2048 ]; then
    status="fail"
  elif ! mount "$DATADIR/$num.img" "$MNTTMP" -o loop; then
    echo "unable to mount image" >&2
    status="fail"
  elif [ ! -f "$MNTTMP/.disk/info" ]; then
    echo "unable to find disk info" >&2
    status="fail"
  else
    echo -e "$num\t$url\t$(cat "$MNTTMP/.disk/info")" >> "$LIST"
    ls --block-size=1 -s "$DATADIR/$num.img" | cut -f1 -d " " > "$DSIZE"
    status="done"
  fi

  echo "$status" >"$DSTAT"

  grep -q "$MNTTMP" /proc/mounts && umount -f "$MNTTMP" ||:

  [ "$status" = "done" ] || rm -f -- $DATADIR/$num.img
  rm -rf -- "$MNTTMP"
  rm -f -- "$DPID"
}

downloading_start(){

  mkdir -p "$DATADIR/download" || fatal "Can't create $DATADIR/download"

  [ ! -f "$DPID" ] ||
    fatal "Error: can't start downloading while previous one is active"

  downloader "$1" &> "$DERR" &
}

downloading_status(){
 [ -d "$DATADIR" ] || fatal "Can't find $DATADIR"

 local num="$(cat $DNUM 2>/dev/null)"
 echo "status:   $(cat $DSTAT 2>/dev/null)"
 local url="$(cat $DURL 2>/dev/null)"
 [ "$url" != "cdrom:" ] || url="CD/DVD"
 echo "url:      $url"
 echo "target:   $DATADIR/$num.img"
 echo "num:      $num"
 echo "pid:      $(cat $DPID 2>/dev/null)"
 local ts="$(cat $DSIZE 2>/dev/null | tr -d -c '[[:digit:]]')"
 local cs="$(ls --block-size=1 -s "$DATADIR/$num.img" 2>/dev/null |\
            cut -f1 -d " " | tr -d -c '[[:digit:]]')"
 echo "currsize: $cs"
 echo "size:     $ts"
 if [ -n "$cs" -a -n "$ts" ]; then echo "percent:  $((100 * $cs / $ts ))";
 else echo "percent:  0"; fi
 echo "error:    $(cat $DERR 2>/dev/null|sed -n '1h;1!H;${x;s/\n/; /Mg;p}')"
}

downloading_stop(){
  [ -d "$DATADIR" ] || fatal "Can't find $DATADIR"

  [ -x "$DATADIR" -a -w "$DATADIR" ] ||
    fatal "Error: not enough permissins for accessing $DATADIR"

  if [ -f "$DPID"  ]; then
    kill "$(cat "$DPID" 2>/dev/null)" ||:
    rm -f "$DPID" ||:
  fi
}

##### 

setup_vnc(){
  local params="$1"

  [ -d "$DATADIR" ] || fatal "Can't find $DATADIR"

  echo $params > $DATADIR/vncparams
  select_image $(show_curr)
}

del_image(){
  local num="$1"

  [ -d "$DATADIR" ] || fatal "Can't find $DATADIR"

  [ -x "$DATADIR" -a -w "$DATADIR" ] ||
    fatal "Error: not enough permissins for accessing $DATADIR"

  if [ "$(show_curr)" = "$num" ]; then
      grep -q "$MNT" /proc/mounts && umount -f "$MNT" ||:
  fi

  sed -i -e "/^$num[[:space:]]/d" "$LIST"
  rm -f "$DATADIR/$num.img" ||:
}

select_image(){
  local num="$1"

  [ -d "$DATADIR" ] || fatal "Can't find $DATADIR"

  [ -x "$DATADIR" -a -w "$DATADIR" ] ||
    fatal "Error: not enough permissins for accessing $DATADIR"

  [ -f "$DATADIR/$num.img" -o "$num" = "0" ] ||
    fatal "can't find $DATADIR/$num.img"

  # removing old image

  grep -q "$MNT" /proc/mounts && umount -f "$MNT" ||:
  rm -f "$CURR" ||:
  rm -f "$TFTPDIR/pxelinux.0"    ||:
  rm -rf "$TFTPDIR/syslinux"     ||:
  rm -rf "$TFTPDIR/pxelinux.cfg" ||:
  sed -i "\%^$CURR%d" /etc/fstab

  [  "$num" != "0" ] || return 0

  # installing new image
  echo -e "$CURR\t$MNT\tiso9660\tloop,ro\t0 0" >> /etc/fstab
  mkdir -p "$MNT"
  ln -s  "$num.img" "$CURR"
  mount "$MNT" ||:

  [ -d "$MNT/syslinux" ] ||
    fatal "Error: bad image: no syslinux directory"
  [ -d "$TFTPDIR" ] ||
    fatal "Error: no $TFTPDIR"
  cp -f "/usr/lib/syslinux/pxelinux.0" "$TFTPDIR/pxelinux.0" ||
    fatal "Error: can't copy pxelinux.0 to $TFTPDIR"
  cp -fr "$MNT/syslinux" "$TFTPDIR/syslinux" ||
    fatal "Error: can't copy syslinux directory to $TFTPDIR"
  mkdir -p "$TFTPDIR/pxelinux.cfg" ||
    fatal "Error: can't create $TFTPDIR/pxelinux.cfg"

  local tz="$(sed -n -e 's/^ZONE=\([^[:space:]]\+\)/\1/p' /etc/sysconfig/clock 2>/dev/null)"
  local lang="$(sed -n -e 's/^LANG=\([^[:space:]\.]\+\).*/\1/p' /etc/sysconfig/i18n 2>/dev/null)"
  local vnc="$(cat $DATADIR/vncparams 2>/dev/null)"
  local default="$(guess_default_label "$TFTPDIR/syslinux/isolinux.cfg")"

  # add 'krb5' option if domain master
  role_file="/etc/sysconfig/system"
  . $role_file ||:
  if [ "$SERVER_ROLE" = "master" ]; then
    krb5=yes
  fi

  # Setting up pxelinux.cfg
  sed -e "
      1i\
  default $default
      s%\(kernel\)[[:space:]]\+\([^[:space:]]\+\)%\1 syslinux/\2%
      s%timeout[[:space:]]\+\([^[:space:]]\+\)%timeout 100%
      /^default[[:space:]]/d
      s/.*gfxboot bootlogo.*//
      /^[[:space:]]*append/{
        s%initrd=\([^[:space:]]\+\)%initrd=syslinux/\1%
        s%automatic=[^[:space:]]\+%%
        s%\$% automatic=method:nfs,network:dhcp${tz:+ tz=$tz} ${lang:+ lang=$lang} ${krb5:+ krb5} $vnc%
      }"\
    "$TFTPDIR/syslinux/isolinux.cfg" \
    > "$TFTPDIR/pxelinux.cfg/default"

  # Setting up DHCP using alterator-dhcp
  if [ -f /usr/bin/alterator-dhcp-functions ];then
    . /usr/bin/alterator-dhcp-functions

    dhcp_config_set client_pxe_server '*'
    dhcp_config_set client_pxe_filename 'pxelinux.0'
    dhcp_config_set client_pxe_root "$DATADIR/current"
    dhcp_update_config
  fi

}

case "${1-}" in
  -l) show_list ;;
  -L) show_labels ;;
  -c) show_curr ;;
  -C) show_default_nlabel ;;
  -g) guess_default_label "$TFTPDIR/syslinux/isolinux.cfg";;
  -r) [ -n "${2-}" ] && show_info "$2" || show_help ;;
  -s) [ -n "${2-}" ] && select_image "$2" || show_help ;;
  -b) [ -n "${2-}" ] && set_default_label "$2" || show_help ;;
  -d) [ -n "${2-}" ] && del_image "$2" || show_help ;;
  -v) [ -n "${2-}" ] && setup_vnc "$2" || show_help ;;
  -A) [ -n "${2-}" ] && downloading_start "$2" || show_help ;;
  -D) downloading_stop   ;;
  -S) downloading_status ;;
  *) show_help ;;
esac
