#!/bin/sh
#
# Copyright 2015 iXsystems (Kris Moore)
# All rights reserved
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted providing that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
# IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

# Script to detect TrueOS / PC-BSD disk installations and create grub
# entries for them

# Check if we have beadm installed, if not skip this file
if [ ! -e "/usr/local/sbin/beadm" ] ; then exit 0; fi

ROOTFS=`mount | awk '/ \/ / {print $1}'`
BEDS="$( echo ${ROOTFS} | awk -F '/' '{print $2}' )"
if [ "$BEDS" = "dev" ] ; then BEDS="ROOT"; fi

serial=0;
console=0;
for x in ${GRUB_TERMINAL_OUTPUT}; do
    if [ xserial = "x$x" ]; then
        serial=1;
    fi
    if [ xgfxterm = "x$x" ]; then
        console=1;
    fi
    if [ xconsole = "x$x" ]; then
        console=1;
    fi
done

if [ "x$serial" = "x1" ]; then
    if [ "x$console" = "x1" ]; then
        out="-Dh"
    else
        out="-h"
    fi
fi

display_loaderopts()
{
  # Optional ARG1, set to a ZFS dataset to mount and read values from
  if [ -n "$1" ] ; then
     fMnt="/mnt.$$"
     mkdir $fMnt
     mount -t zfs  ${1} $fMnt >/dev/null 
     if [ $? -ne 0 ] ; then
        echo "Failed to mount ${1}" >&2
        return
     fi
  else
     fMnt=""
  fi
  
  # Create our placeholder to save _load entries to parse
  touch /tmp/.lRObjs.$$
  touch /tmp/.lSysCtls.$$

  # Loader files, in order of which to read them
  lFiles="${fMnt}/boot/loader.conf.local ${fMnt}/boot/loader.conf ${fMnt}/boot/loader.conf.trueos ${fMnt}/boot/loader.conf.pcbsd"
  for f in $lFiles
  do
     if [ ! -e "$f" ] ; then continue ; fi

     # Lets parse any of the _load= lines
     grep "_load=" ${f} | grep -v "^#" >/tmp/.lObjs.$$
     while read line
     do
       loadVal="`echo $line | cut -d '=' -f 1`"
       # Is this value already set in a higher priority file?
       grep -q "^${loadVal}" /tmp/.lRObjs.$$
       if [ $? -eq 0 ];then continue; fi

       # Save this value for later
       echo "$line" >> /tmp/.lRObjs.$$
       haveObjs=1

     done < /tmp/.lObjs.$$
     rm /tmp/.lObjs.$$

     # Lets look for any sysctls to set
     grep "." ${f} | grep "=" | grep -v "^#" | grep -v "_load" >/tmp/.sObjs.$$
     while read line
     do
       loadVal="`echo $line | cut -d '=' -f 1`"
       # Is this value already set in a higher priority file?
       grep -q "^${loadVal}" /tmp/.lSysCtls.$$
       if [ $? -eq 0 ];then continue; fi

       # Save this value for later
       echo "$line" >> /tmp/.lSysCtls.$$
       haveSysCtls=1

     done < /tmp/.sObjs.$$
     rm /tmp/.sObjs.$$
  done

  # Using GELI encryption?
  haveGELI="false"

  # Now lets echo out the modules to load
  if [ "$haveObjs" = "1" ] ; then
    while read line
    do
      echo "$line" | grep -q '"YES"'
      if [ $? -ne 0 ] ; then continue ; fi
      module="`echo $line | cut -d '=' -f 1 | sed 's|_load||g'`"
      # Try to locate module now
      if [ -e "${fMnt}/boot/kernel/${module}.ko" ] ; then
        mPath="kernel"
      elif [ -e "${fMnt}/boot/modules/${module}.ko" ] ; then
        mPath="modules"
      else
	# This isn't a module that we can see, lets set it as a variable
	#echo "No such module $line, setting as a variable" >&2
        echo "$line" >> /tmp/.lSysCtls.$$
        continue
      fi

      # Are we loading GELI module?
      if [ "$module" = "geom_eli" ] ; then haveGELI="true" ; fi

      echo "    kfreebsd_module_elf ${loadPrefix}/@/boot/${mPath}/${module}.ko"
    done < /tmp/.lRObjs.$$
  fi

  # Add our compat kenv lines
  echo "    set kFreeBSD.bootfile=\"kernel\""
  echo "    set kFreeBSD.kernel=\"kernel\""
  echo "    set kFreeBSD.kernel_options=\"\""
  echo "    set kFreeBSD.kernelname=\"/boot/kernel/kernel\""
  echo "    set kFreeBSD.module_path=\"/boot/kernel;/boot/modules\""

  # Any kenv to set?
  if [ "$haveSysCtls" = "1" ] ; then
    while read line
    do
      # Strip out the vfs.root.mountfrom, we set that elsewhere
      echo "$line" | grep -q "vfs.root.mountfrom"
      if [ $? -eq 0 ] ; then continue ; fi

      line="`echo $line | sed 's|"||g'`"
      key="`echo $line | cut -d '=' -f 1`"
      val="`echo $line | cut -d '=' -f 2`"
      if [ -z "$key" -o -z "$val" ] ; then continue ; fi
      echo "    set kFreeBSD.${key}=\"${val}\""
    done < /tmp/.lSysCtls.$$
  fi

  rm /tmp/.lRObjs.$$
  rm /tmp/.lSysCtls.$$
  if [ -n "$1" ] ; then
     umount /mnt.$$ >/dev/null
     rmdir /mnt.$$ >/dev/null
  fi

  # Set the grub.platform kenv variable
  echo "    set kFreeBSD.grub.platform=\"\$grub_platform\""

  # See if we need to do GELI passphrase passthrough
  if [ "$haveGELI" = "true" ] ; then
    echo "    set kFreeBSD.kern.geom.eli.passphrase=\"\$pass\""
  fi
}

detect_beadm()
{
   /usr/bin/which -s beadm >/dev/null 2>/dev/null
   if [ $? -ne 0 ] ; then return 0; fi

   # Check if we are running from the installer and use its beadm
   if [ -e "/root/beadm.install" ] ; then
      BEADM="/root/beadm.install"

      # Check if this is valid
      testBE=`$BEADM list`
      if [ -z "$testBE" ] ; then
        # No BE's, lets switch back to regular mode
        BEADM="`/usr/bin/which beadm`"
      fi
   else
      BEADM="`/usr/bin/which beadm`"
   fi

   ${BEADM} list >/dev/null 2>/dev/null
   if [ $? -ne 0 ] ; then return 0; fi


   if [ -e "/etc/defaults/pcbsd" ] ; then
      NICK="PC-BSD"
   else
      NICK="TrueOS"
   fi

   # Now we will list all the non-default BEs
   $BEADM list -H >/tmp/.grub-beadm.$$ 2>/dev/null
   while read line
   do
     be=`echo $line | awk '{print $1}'`
     flags=`echo $line | awk '{print $2}'`

     # Skip the active BE
     if [ "$flags" = "NR" -o "$flags" = "R" ] ; then
	continue
     fi
     BE_LIST="${be} ${BE_LIST}"
   done < /tmp/.grub-beadm.$$
   rm /tmp/.grub-beadm.$$

   if [  -z "$BE_LIST" ] ; then return; fi

   echo "submenu \"Boot Environment Menu\" {"

   # Get list of beadm datasets
   for b in $BE_LIST
   do 
      # Got a beadm snapshot, lets get the complete dataset name
      beLine=`${BEADM} list -a | grep "/$BEDS/${b}"`
      cdataset=`echo $beLine | awk '{print $1}'`
      cdatadate=`echo $beLine | awk '{print $5}'`
      cdatatime=`echo $beLine | awk '{print $6}'`
      ztank=`echo $cdataset | cut -d '/' -f 1`
      shortdataset="/`echo $cdataset | cut -d '/' -f 2-5`"
      # Get the BE nickname if it exists
      beNickName=`${BEADM} list | grep "^${b} " | tr -s ' ' | cut -d ' ' -f 7-12`
      if [ -z "$beNickName" ] ; then
         beNickName="$b"
      fi

      # First part of this dataset
      cat > /tmp/.grubdataset.$$.1 << EOF
    insmod zfs
    search --no-floppy -s -l $ztank
EOF

      # Second part of loader to save
      cat > /tmp/.grubdataset.$$.2 << EOF
    kfreebsd_module ${shortdataset}/@/boot/zfs/zpool.cache type=/boot/zfs/zpool.cache
    set kFreeBSD.vfs.root.mountfrom=zfs:$cdataset
EOF
      
      # Now lets look for options in loader.conf to load
      loadPrefix="${shortdataset}"
      # If this is the current mounted dataset, we can skip mounting it
      mount | grep -q -e "$cdataset on / (" -e "$cdataset on /mnt ("
      if [ $? -eq 0 ] ; then
        display_loaderopts >>/tmp/.grubdataset.$$.2
      else
        display_loaderopts $cdataset >> /tmp/.grubdataset.$$.2
      fi

      # Lets start a submenu for each BE
cat << EOF

submenu "${NICK} (${beNickName}) - ${cdatadate} ${cdatatime}" {
EOF
  
  # Lets do the default entry first
  #################################
  cat << EOF
  menuentry "Normal Bootup" {
EOF
  # Get the dataset guts
  cat /tmp/.grubdataset.$$.1
  echo "    kfreebsd ${shortdataset}/@/boot/kernel/kernel ${out}"
  echo "    kfreebsd_loadenv ${shortdataset}@/boot/device.hints"
  cat /tmp/.grubdataset.$$.2

  # Set any options
  cat << EOF
  }
EOF

  # Next lets do single user mode
  #################################
  cat << EOF
  menuentry "Single User Mode" {
EOF
  # Get the dataset guts
  cat /tmp/.grubdataset.$$.1
  echo "    kfreebsd ${shortdataset}/@/boot/kernel/kernel ${out} -s"
  echo "    kfreebsd_loadenv ${shortdataset}@/boot/device.hints"
  cat /tmp/.grubdataset.$$.2

  # Set any options
  cat << EOF
  }
EOF

  # Next lets do verbose mode
  #################################
  cat << EOF
  menuentry "Verbose Mode" {
EOF
  # Get the dataset guts
  cat /tmp/.grubdataset.$$.1
  echo "    kfreebsd ${shortdataset}/@/boot/kernel/kernel ${out} -v"
  echo "    kfreebsd_loadenv ${shortdataset}@/boot/device.hints"
  cat /tmp/.grubdataset.$$.2

  # Set any options
  cat << EOF
  }
EOF

  if [ -e "/etc/defaults/pcbsd" ] ; then
    # Next lets do display wizard
    #################################
    cat << EOF
  menuentry "Run the Display Wizard" {
EOF
    # Get the dataset guts
    cat /tmp/.grubdataset.$$.1
    echo "    kfreebsd ${shortdataset}/@/boot/kernel/kernel ${out}"
    echo "    kfreebsd_loadenv ${shortdataset}@/boot/device.hints"
    cat /tmp/.grubdataset.$$.2

    # Set any options
    cat << EOF
    set kFreeBSD.runwiz=YES
  }
EOF
    # Now for vesa mode
    #################################
    cat << EOF
  menuentry "Run X in vesa mode" {
EOF
    # Get the dataset guts
    cat /tmp/.grubdataset.$$.1
    echo "    kfreebsd ${shortdataset}/@/boot/kernel/kernel ${out}"
    echo "    kfreebsd_loadenv ${shortdataset}@/boot/device.hints"
    cat /tmp/.grubdataset.$$.2

    # Set any options
    cat << EOF
    set kFreeBSD.xvesa=YES
  }
EOF
    # Now for non-xorg mode
    #################################
    cat << EOF
  menuentry "Boot to console (Disable X)" {
EOF
    # Get the dataset guts
    cat /tmp/.grubdataset.$$.1
    echo "    kfreebsd ${shortdataset}/@/boot/kernel/kernel ${out}"
    echo "    kfreebsd_loadenv ${shortdataset}@/boot/device.hints"
    cat /tmp/.grubdataset.$$.2

    # Set any options
    cat << EOF
    set kFreeBSD.noxorg=YES
  }
EOF
  fi


      # Lastly lets close the submenu section
      cat << EOF
}

EOF
   done

   echo "}"
   echo ""

   # Cleanup after ourselves
   if [ -e "/tmp/.grubdataset.$$.1" ] ; then
     rm /tmp/.grubdataset.$$.1
     rm /tmp/.grubdataset.$$.2
   fi
}

# Detect our types of disk layouts
detect_beadm

