#!/bin/bash
# Steven Shiau <steven _at_ clonezilla org>
# License: GPL
# Description: prepare the memtest86+ by using the files on the system or download and extract the required programs for DRBL.
DRBL_SCRIPT_PATH="${DRBL_SCRIPT_PATH:-/usr/share/drbl}"

. $DRBL_SCRIPT_PATH/sbin/drbl-conf-functions
[ -e /etc/drbl/drbl-ocs.conf ] && . /etc/drbl/drbl-ocs.conf
[ -e $DRBL_SCRIPT_PATH/sbin/ocs-functions ] && . $DRBL_SCRIPT_PATH/sbin/ocs-functions

# Local settings
# By default we only copy the required files to $pxelinux_binsrc_dir. This option will allow us to copy files to dir $pxecfg_pd.
deploy_to_system_too="no"

# Functions
USAGE() {
  echo "Put the memtest86+ files in DRBL package repository"
  echo "Usage: $0 [OPTION]"
  echo "OPTION:"
  echo "-u, --from-upstream   Use the binary from upstream ($MEMTEST86_URL). If this option is not assigned, DRBL will try to use the memtest86+ from your GNU/Linux distribution."
  echo "-mv, --memtest86-version VERSION  The memtest86+ version to be used with -u|--from-upstream. If not assigned, the version number specified in drbl.conf will be used."
  echo "-d, --deploy-to-system-too   Deploy the file to DRBL system ($pxecfg_pd), too."
  echo "Ex: To use the memtest86+ 4.20 from $MEMTEST86_URL, run '$0 -u -mv 4.20'"
}

#
put_memtest_bin_2_drbl_repo() {
  local memtest_bin_abs_path="$1"
  [ -z "$memtest_bin_abs_path" ] && return 1
  echo "Putting memtest86+ in DRBL package repository $memtest86_dir/... "
  memtest_f="$(shorten_memtest_filename `basename $memtest_bin_abs_path`)"
  cp -af $memtest_bin_abs_path $memtest86_dir/$memtest_f
  memtest_v="$(LC_ALL=C strings $memtest_bin_abs_path  | grep -Ei 'Memtest86\+[[:space:]]+v' | sed -r -e "s/^[[:space:]]*//g")"
  echo "Memtest86+ version: $memtest_v"
  echo $memtest_v > $memtest86_dir/VERSION
  echo "done!"
  if [ "$deploy_to_system_too" = "yes" ]; then
    echo -n "Putting memtest86+ in DRBL system dir $pxecfg_pd/... "
    cp -af $memtest_f $pxecfg_pd/
    echo $memtest_v > $pxecfg_pd/MEMTEST86+_VERSION
    echo "done!"
  fi
} # end of put_memtest_bin_2_drbl_repo

#
put_upstream_memtest(){
  # Function to use the upstream binary directly
  local ver="$1"
  if [ -z "$ver" ]; then
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
    echo "No upstream memtest version was assigned. Program terminated."
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    USAGE
    exit 1
  fi
  memfile="memtest86+-${ver}.bin"
  tmp_wd="$(LC_ALL=C mktemp -d /tmp/memtest_tmp.XXXXXXX)"
  echo -n "Downloading ${ver} memtest86+ files from $MEMTEST86_URL... "
  LC_ALL=C wget $wget_opt -P $tmp_wd $MEMTEST86_URL/$ver/$memfile
  echo "done!"
  put_memtest_bin_2_drbl_repo $tmp_wd/$memfile
  if [ -d "$tmp_wd" -a -n "$(echo "$tmp_wd" | grep -i memtest_tmp)" ]; then
    rm -rf $tmp_wd
  fi
} # end of put_upstream_memtest
#
shorten_memtest_filename() {
  local inp_fn="$1"
  # memtest86+.bin -> mt86+x32.mbr
  # memtest86+x32.bin -> mt86+x32.mbr
  # memtest86+x32.efi -> mt86+x32.efi
  # memtest86+x64.bin -> mt86+x64.mbr
  # memtest86+x64.efi -> mt86+x64.efi
  case "$inp_fn" in
    memtest86+.bin)    echo mt86+x32.mbr;;
    memtest86+x32.bin) echo mt86+x32.mbr;;
    memtest86+x32.efi) echo mt86+x32.efi;;
    memtest86+x64.bin) echo mt86+x64.mbr;;
    memtest86+x64.efi) echo mt86+x64.efi;;
  esac
} # end of shorten_memtest_filename
#
put_distribution_memtest() {
  # Function to use the package from distribution.
  # Memtest86+ file name on different distributions:
  # Debian/Ubuntu: "memtest86+.bin"
  # SuSE: "memtest.bin"
  # Fedora 14: "/boot/memtest86+"
  # Centos 5.6: "/boot/memtest86+-1.65"
  if $query_pkglist_exist_cmd memtest86+ &>/dev/null; then
     # Found on the system
     # For Redhat-like system, there might return more than one file on a system, i.e. one file belongs to 2 rpm packages.
     # 2022/11/23 On Debian, memtest86+ v6 is available:
     # /boot/memtest86+x32.bin
     # /boot/memtest86+x32.efi
     # /boot/memtest86+x64.bin
     # /boot/memtest86+x64.efi
     # While for memtest86+ <= v5, the file is:
     # /boot/memtest86+.bin

     # For syslinux in FAT file system, 
     # "memtest86+.bin" is 9 characters, will go wrong when it's FAT (usb flash drive). We use memtest to overwrite the one comes from Debian live.
     # In addition, do NOT use memtest.bin, use memtest instead of memtest.bin
     # since "*.bin" has specially purpose for syslinux.
     # So here they are:
     # memtest86+.bin -> mt86+x32.mbr
     # memtest86+x32.bin -> mt86+x32.mbr
     # memtest86+x32.efi -> mt86+x32.efi
     # memtest86+x64.bin -> mt86+x64.mbr
     # memtest86+x64.efi -> mt86+x64.efi

     # For memtest <= v5.0
     memtest_bin="$(LC_ALL=C $query_pkglist_cmd memtest86+ | grep -Ew "(memtest86\+.bin$|memtest\.bin$|^/boot/memtest86+)" | sort | head -n 1)"
     # For memtest >= v6.0
     memtest_x32_bin="$(LC_ALL=C $query_pkglist_cmd memtest86+ | grep -Ew "^/boot/memtest86\+x32\.bin" | sort | head -n 1)"
     memtest_x32_efi="$(LC_ALL=C $query_pkglist_cmd memtest86+ | grep -Ew "^/boot/memtest86\+x32\.efi" | sort | head -n 1)"
     memtest_x64_bin="$(LC_ALL=C $query_pkglist_cmd memtest86+ | grep -Ew "^/boot/memtest86\+x64\.bin" | sort | head -n 1)"
     memtest_x64_efi="$(LC_ALL=C $query_pkglist_cmd memtest86+ | grep -Ew "^/boot/memtest86\+x64\.efi" | sort | head -n 1)"
     for im in $memtest_bin $memtest_x32_bin $memtest_x32_efi $memtest_x64_bin $memtest_x64_efi; do
       eval imfile="\$im"
       if [ -n "$imfile" ]; then
         echo "Found $imfile in this system, copying the memtest file to DRBL local repository..."
         put_memtest_bin_2_drbl_repo $imfile
       else
         [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
         echo "Package of memtest was installed, but memtest boot file not found!"
         [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
       fi
     done
  else
     # Download and extract it. The reason we do not install it is we do not want to bother the server's boot menu
     tmp_wd="$(LC_ALL=C mktemp -d /tmp/memtest_tmp.XXXXXXX)"
     echo "Trying to download memtest86+ from distribution repository..."
     if [ -e /etc/debian_version ]; then
       # Debian
       # The binary for memtest86+ in Debian Squeeze: /boot/memtest86+.bin
       # Clean all stale deb files in the cache dir. Otherwise there might be some old version of debs exist in the cache dir.
       LC_ALL=C apt-get clean  
       LC_ALL=C apt-get -d --reinstall install memtest86+
       (
        cd $tmp_wd
        LC_ALL=C dpkg --extract /var/cache/apt/archives/memtest86+*.deb .
        put_memtest_bin_2_drbl_repo "$(pwd)"/boot/memtest86+.bin
       )
     elif [ -e /etc/SuSE-release ]; then
       # SuSE
       # The binary for memtest86+ in openSuse 11.3 is: /boot/memtest.bin
       LC_ALL=C zypper install -d -f -y memtest86+
       memtest86_rpm="$(LC_ALL=C find /var/cache/zypp/packages/ -name "memtest86+-*.rpm" -print)"
       memtest_bin="$(LC_ALL=C rpm -qpl "$memtest86_rpm" | grep -Ew "memtest.bin$")"
       (
        cd $tmp_wd
        LC_ALL=C rpm2cpio "$memtest86_rpm" | cpio -idm
        put_memtest_bin_2_drbl_repo "$(pwd)"/$memtest_bin
       )
     else
       # RH-like
       # The binary for memtest86+ in Fedora 14 is: /boot/memtest86+-4.10
       LC_ALL=C yumdownloader --destdir $tmp_wd memtest86+
       memtest_bin="$(LC_ALL=C rpm -qpl "$tmp_wd/memtest86*.rpm" | grep -Ew "^/boot/memtest86+")"
       (
        cd $tmp_wd
        LC_ALL=C rpm2cpio $tmp_wd/memtest86*.rpm | cpio -idm
        put_memtest_bin_2_drbl_repo "$(pwd)"/$memtest_bin
       )
     fi
     if [ -d "$tmp_wd" -a -n "$(LC_ALL=C echo "$tmp_wd" | grep -i memtest_tmp)" ]; then
       rm -rf $tmp_wd
     fi
  fi
} # end of put_distribution_memtest

############
### MAIN ###
############

check_if_root
ask_and_load_lang_set

#
while [ $# -gt 0 ]; do
  case "$1" in
    -u|--from-upstream) shift; mode="from-upstream" ;;
    -mv|--memtest86-version)
        shift;
        if [ -z "$(echo $1 |grep ^-.)" ]; then
          # skip the -xx option, in case 
          memtest_ver_requested="$1"
          [ -z "$memtest_ver_requested" ] && USAGE && exit 1
	  shift
        fi
	;;
    -d|--deploy-to-system-too) shift; deploy_to_system_too="yes" ;;
    *)  USAGE && exit 1 ;;
  esac
done

# Check arch
cpu_arch="$(LC_ALL=C uname -m)"
if [ "$cpu_arch" != "i686" -a "$cpu_arch" != "x86_64" ]; then
  [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
  echo "CPU arch is not i386/i686 or x86_64!"
  echo "Skipping preparing memtest+!"
  [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
  echo "$msg_program_stop"
  exit 3
fi

# mode is essential
[ -z "$mode" ] && mode="from-distribution"
[ -z "$memtest_ver_requested" ] && memtest_ver_requested="$MEMTEST86_VER_DEF"

case "$mode" in
  from-distribution) put_distribution_memtest ;;
  from-upstream)     put_upstream_memtest $memtest_ver_requested ;;
esac

exit 0
