#!/bin/bash
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero 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 Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public
# License along with this program.  If not, see
# <http://www.gnu.org/licenses/>.

# Load.
source $APP_BASE/lib/hydra/functions || exit 1
hydra_config_load

# Set bootless folder
function hydra_bootless_folder {
  # Check for a bootless repository
  if [ -e "$HYDRA_FOLDER/bootless" ]; then
    BOOTLESS_DIR="$HYDRA_FOLDER/bootless"
  elif [ -e "$HYDRA_FOLDER/conf/bootless" ]; then
    BOOTLESS_DIR="$HYDRA_FOLDER/conf/bootless"
  else
    echo "Please make a symlink $HYDRA_FOLDER/bootless pointing to your devel puppet modules."
    exit 1
  fi
}

# Make a boot device
function hydra_bootless_make {
  # Set folder
  hydra_bootless_folder

  # Set sudo config
  local sudo device rsync
  if [ "`whoami`" != 'root' ]; then
    sudo="sudo"
  fi
  
  # Script description
  cat <<EOF
Bootless Install Script
=======================

This script installs a bootable copy of the bootless repository in a device
partition. It will ask for a device partition, clone the repository in the
partition and make the device bootable.

Press any key to continue, or ^C to abort.
EOF

  read tmp
  
  # Git repo consistency check
  ( cd ${BOOTLESS_DIR} && git status > /dev/null )
  if [ $? -ne 0 ]; then
    echo "Error: '${BOOTLESS_DIR}' is not a git repository."
    exit 1
  fi
  
  if [ -z "$1" ]; then
    echo -n "Target device: "
    read device
  else
    device=$1
  fi

  if [ "$2" == "--rsync" ]; then
    rsync="true"
  fi

  usbdevice=`echo ${device} | sed -e s/[0-9]\$//g`
  
  # Issue a warning
  cat <<EOF

******************
*  ATTENTION!!!  *
******************

If you continue, all data in device "${device}" will be destroyed!

EOF

  echo -n "Are you sure you want to continue? Type uppercase \"YES\": "
  read go
  
  if [ "${go}" != "YES" ]; then
    echo "Aborting..."
    exit 1
  fi

  hydra_bootless_device ${usbdevice} ${device}

  # Mount
  tmpdir=`mktemp -d`
  $sudo mount ${device} ${tmpdir}
  if [ $? != 0 ]; then
    echo "Error: failed to mount \"${device}\" filesystem in \"${tmpdir}\" (errno: $?)."
    exit 1
  fi
  
  # Copy data
  if [ "$rsync" == "true" ]; then
    $sudo rsync -Cav ${BOOTLESS_DIR}/ ${tmpdir}/boot/
  else
    $sudo git clone --depth=1 file:///${BOOTLESS_DIR} ${tmpdir}/boot
    if [ $? != 0 ]; then
      echo "Error: failed to clone repository \"${BOOTLESS_DIR}\" in \"${tmpdir}\" (errno: $?)."
      exit 1
    fi
  fi
  
  # Grub legacy
  #devicemap=`mktemp`
  #grub-mkdevicemap --no-floppy --device-map=${devicemap}
  #usbdevice=`echo ${device} | sed -e s/[0-9]\$//`
  #grubroot=`grep "${usbdevice}" ${devicemap} | cut -f1`
  #grubdevice=`echo ${grubroot} | sed -e s/\)/,0\)/`
  #echo "root ${grubdevice}
  #setup ${grubroot}
  #quit" | grub --device-map=${devicemap} --batch
  
  # Grub 2
  #$sudo grub-install --root-directory=${tmpdir} ${usbdevice} --force
  $sudo grub-install --boot-directory=${tmpdir} ${usbdevice} --force
  if [ $? != 0 ]; then
    echo "Error: grub-install failed (errno: $?)."
    exit 1
  fi
  
  # Finalize
  #rm -f ${devicemap}
  $sudo umount ${tmpdir}
  $sudo rm -rf ${tmpdir}
}

# Update a boot device
function hydra_bootless_update {
  # Set folder
  hydra_bootless_folder

  # Set sudo config
  local sudo
  if [ "`whoami`" != 'root' ]; then
    sudo="sudo"
  fi

  if [ -z "$1" ]; then
    echo -n "Target device: "
    read device
  else
    device=$1
  fi

  # Target device consistency check
  if [ ! -b ${device} ]; then
    echo "Error: device \"${device}\" not found."
    exit 1
  elif [ ! -d "${tmpdir}/boot" ]; then
    echo "Error: bootless repository not found at ${device}."
    exit 1
  fi

  tmpdir=`mktemp -d`
  $sudo mount ${device} ${tmpdir}
  if [ $? != 0 ]; then
    echo "Error: failed to mount \"${device}\" filesystem in \"${tmpdir}\" (errno: $?)."
    exit 1
  fi

  # Copy data
  if [ -d "${tmpdir}/boot/.git" ]; then
    ( cd ${tmpdir}/boot && $sudo git pull origin master )
  else
    $sudo rsync -Cav ${BOOTLESS_DIR}/ ${tmpdir}/boot/
  fi

  # Finalize
  $sudo umount ${tmpdir}
  $sudo rm -rf ${tmpdir}
}

# Repository initiator
function hydra_bootless_init {
  mkdir -p $HYDRA_FOLDER

  if [ ! -z "$1" ]; then
    # Clone from url
    git clone $1 $HYDRA_FOLDER/bootless
    exit $?
  fi

  # Create a fresh repository
  mkdir -p $HYDRA_FOLDER/bootless/{default,custom,grub}
  mkdir -p $HYDRA_FOLDER/bootless/default/{debian,memtest,ubuntu}
  ( cd $HYDRA_FOLDER/bootless && ln -s . boot)

  if [ -f "/boot/memtest86+.bin" ]; then
    cp /boot/memtest86+.bin $HYDRA_FOLDER/bootless/default/memtest
  else
    echo "No memtest image found. Please install memtest86+ package"
    echo "and manually copy /boot/memtest86+.bin if you want memtest support"
  fi

  # Grub configuration
  cat > $HYDRA_FOLDER/bootless/grub/grub.cfg <<-EOF    
# This is grub.cfg for use with Bootless Management System

### BEGIN header ###
if [ -s $prefix/grubenv ]; then
  load_env
fi
set default="0"
if [ "${prev_saved_entry}" ]; then
  set saved_entry="${prev_saved_entry}"
  save_env saved_entry
  set prev_saved_entry=
  save_env prev_saved_entry
  set boot_once=true
fi

function savedefault {
  if [ -z "${boot_once}" ]; then
    saved_entry="${chosen}"
    save_env saved_entry
  fi
}

function load_video {
}

set timeout=5
### END header ###

### BEGIN debian_theme ###
set menu_color_normal=white/blue
set menu_color_highlight=yellow/red
### END debian_theme ###
EOF

  # Initialize git repository
  (
  cd $HYDRA_FOLDER/bootless
  git init
  git add .
  git commit -a -m "Initial import"
  )

  echo "Now add your boot images and edit $HYDRA_FOLDER/bootless/grub/grub.cfg to suit your needs."
}

# Git wrapper
function hydra_bootless_git {
  # Set folder
  hydra_bootless_folder

  (
  cd $BOOTLESS_DIR && git $*
  )
}

# Create a target device
# TODO: properly support $subdevice in parted
function hydra_bootless_device {
  local device="$1"
  local subdevice="$2"

  # Check if device is mounted
  if [ "`mount | grep ${subdevice}`" != "" ]; then
    echo "Error: device \"${subdevice}\" is mounted."
    exit 1
  fi

  # Target device consistency check
  if [ ! -b ${subdevice} ]; then
    echo "Error: device \"${subdevice}\" not found."
    exit 1
  fi

  # Remove old partitions
  for partition in `$sudo parted -s -- ${device} print | awk '/^ / {print $1}'`; do
    $sudo parted -s -- ${device} rm $partition
  done

  # Create a single partition
  $sudo parted -s -- ${device} mkpart primary ext2 2 -1
  $sudo parted -s -- ${device} set 1 boot on

  # Create filesystem
  $sudo mke2fs ${subdevice}
  if [ $? != 0 ]; then
    echo "Error: mke2fs failed in \"${subdevice}\" (errno: $?)."
    exit 1
  fi

  # Tune
  $sudo tune2fs -c 0 -i 0 ${subdevice}
  if [ $? != 0 ]; then
    echo "Error: tune2fs failed in \"${subdevice}\" (errno: $?)."
    exit 1
  fi
}

# Generate a bootless disk image
function hydra_bootless_image {
  local output="$1"

  if [ -z "$output" ]; then
    output="bootless.img"
  fi

  # Set folder
  hydra_bootless_folder

  # Copy data
  tmpdir=`mktemp -d`
  $sudo rsync -Cav ${BOOTLESS_DIR}/ ${tmpdir}/boot/

  # Make rescue disk
  grub-mkrescue -o ${output} ${tmpdir}

  # Cleanup
  rm -rf ${tmpdir}
  echo "Image saved at ${output}"
}

# Usage
function hydra_bootless_usage {
  echo "Usage: `basename $0` <action> [arguments]"
  exit 1
}

# Check for requirements
for req in parted kpartx; do
  hydra_install_package $req
done

# Parameter verification
if [ -z "$1" ]; then
  hydra_bootless_usage
elif [ -z "$HYDRA_FOLDER" ]; then
  echo "Parameter HYDRA_FOLDER not configured."
  exit 1
elif [ "$1" == "make" ]; then
  shift
  hydra_bootless_make $*
elif [ "$1" == "init" ]; then
  shift
  hydra_bootless_init $*
elif [ "$1" == "update" ]; then
  shift
  hydra_bootless_update $*
elif [ "$1" == "git" ]; then
  shift
  hydra_bootless_git $*
elif [ "$1" == "image" ]; then
  shift
  hydra_bootless_image $*
else
  hydra_bootless_usage
fi