#!/bin/bash
### rootfsinstall: create partitions and fileysytems
#
# Based on rpi-install by: Davide Madrisan
# Modified by: Silvan Calarco <silvan.calarco@mambasoft.it>
#
SYSTEMFSTYPE=ext4

function usage() {
   echo "Usage: $0 device rootfs-archive [-y] [-rpi] [-fs fstype] [-a archive] [-c command] [-k kernel] [-p password]"
   echo
   echo "Options:"
   echo "-a archive  : additional archive to unpack in rootfs root"
   echo "-cp command : command to run before chroot. Use @ROOT@ as reference to rootfs path"
   echo "-c command  : command to run through chroot in rootfs"
   echo "-fs fstype  : filesystem type for system partition (default: ext4)"
   echo "-h hostname : set given hostname"
   echo "-k kernel   : extract/use kernel file from ROOTFS/boot/ for boot"
   echo "-p password : set given password for root user"
   echo "-rpi        : make a RaspberryPi SD with boot partition"
   echo "-y          : assume yes to all questions (script mode)"
   echo
   echo "Example: $0 /dev/sdb /tmp/openmamba-rootfs-light.arm.tar.xz"
   exit 1
}

[ $UID -eq 0 ] || {
   echo "ERROR: this script must be called as root; aborting."
   exit 1
}

while [ "$1" ]; do
   case $1 in
      -y) ASSUMEYES=1 ;;
      -rpi) RPIMODE=1 ;;
      -fs) [ "$2" ] || {
              echo "Error: -fs option requires an argument; aborting."
              exit 1
           }
           SYSTEMFSTYPE=$2
           shift
           ;;
      -a) [ "$2" ] || {
              echo "Error: -a option requires an archive as argument; aborting."
              exit 1
           }
           ADDITIONALARCHIVE=`realpath $2`
           [ -r "$ADDITIONALARCHIVE" ] || {
              echo "Error: file $ADDITIONALARCHIVE not found; aborting."
              exit 1
           }
           shift
           ;;
      -c) [ "$2" ] || {
              echo "Error: -c option requires a command string as argument; aborting."
              exit 1
           }
           CHROOTCOMMAND=$2
           shift
           ;;
      -cp) [ "$2" ] || {
              echo "Error: -p option requires a command string as argument; aborting."
              exit 1
           }
           PRECHROOTCOMMAND=$2
           shift
           ;;
      -k) [ "$2" ] || {
              echo "Error: -k option requires a kernel file name as argument; aborting."
              exit 1
           }
           USEKERNEL=$2
           shift
           ;;
      -h) [ "$2" ] || {
              echo "Error: -h option requires a hostname as argument; aborting."
              exit 1
           }
           SETHOSTNAME=$2
           shift
           ;;
      -p) [ "$2" ] || {
              echo "Error: -p option requires a password as argument; aborting."
              exit 1
           }
           ROOTPASSWORD=$2
           shift
           ;;
      -*) echo "Error: invalid option $1; aborting."
          exit 1
          ;;
       *) if [ ! "$DRIVE" ]; then
             DRIVE=$1
          elif [ ! "$ROOTFSARCHIVE" ]; then
             ROOTFSARCHIVE=`realpath $1`
          else
             echo "Error: invalid argument $1; aborting."
             exit 1
          fi
          ;;
    esac
    shift
done
[ "$ROOTFSARCHIVE" ] || { usage; exit 1; }

if [ "${DRIVE:0:11}" = "/dev/mmcblk" ]; then
   PARTITION1=${DRIVE}p1
   PARTITION2=${DRIVE}p2
else
   PARTITION1=${DRIVE}1
   PARTITION2=${DRIVE}2
fi

# FIXME: in non rpi mode we only have system partition
[ "$RPIMODE" ] || PARTITION2=$PARTITION1

[ "$ASSUMEYES" ] || {
   echo -n "All current data in ${DRIVE} will be lost, do you want to proceed [y/N]?"
   read ans
   [ "$ans" = "y" -o "$ans" = "Y" ] || exit 0
}

sudo parted -s ${DRIVE} mktable msdos
if [ "$?" != "0" ]; then
   echo "ERROR: parted exited with error $? while creating partition table on ${DRIVE}; aborting."
fi

if [ "$RPIMODE" ]; then
   # FAT partition
   sudo parted -s -a cyl -- ${DRIVE} mkpart primary fat32 0 40M
   if [ "$?" != "0" ]; then
      echo "ERROR: parted exited with error $? while creating boot partition on ${DRIVE}; aborting."
   fi
   # rootfs partition
   sudo parted -s -a cyl -- ${DRIVE} mkpart primary $SYSTEMFSTYPE 40M -1
   if [ "$?" != "0" ]; then
      echo "ERROR: parted exited with error $? while creating system partition on ${DRIVE}; aborting."
   fi
else
   # rootfs partition
   sudo parted -s -a cyl -- ${DRIVE} mkpart primary $SYSTEMFSTYPE 0 -1
   if [ "$?" != "0" ]; then
      echo "ERROR: parted exited with error $? while creating system partition on ${DRIVE}; aborting."
   fi
fi
sync

if [ "$RPIMODE" ]; then
   echo "Formatting $PARTITION1 (vfat) ..."
   mkfs.vfat -s 1 -F 32 -n "BOOT" $PARTITION1 || {
      echo "ERROR: unable to format $PARTITION1; aborting."
      exit 1
   }
fi

echo "Formatting $PARTITION2 ($SYSTEMFSTYPE)..."
mke2fs -q -t $SYSTEMFSTYPE -j -L "openmamba" $PARTITION2 || {
   echo "ERROR: unable to format $PARTITION2; aborting."
   exit 1
}

BOOTMP=`mktemp -d -t rootfsinstall.XXXXXXXX`
ROOTMP=`mktemp -d -t rootfsinstall.XXXXXXXX`

[ -d $BOOTMP -a -d $ROOTMP ] || {
   echo "Error: unable to create temporary mount points; aborting."
   exit 1
}

mount $PARTITION2 $ROOTMP || {
   echo "ERROR: unable to mount partition $PARTITION2; aborting."
   exit 1
}
trap "[ -e $ROOTMP ] && umount $ROOTMP;[ -e $BOOTMP ] && umount $BOOTMP" 0 HUP INT QUIT ABRT KILL TERM

if [ "$RPIMODE" ]; then
   ### SD-Card: install software

   if [ ! "$USEKERNEL" ]; then
      SWHOME=`mktemp -d`
      echo "Cloning firmware repository from git..."
      cd $SWHOME
      git clone git://github.com/raspberrypi/firmware.git --depth=1
   else
      echo "Fetching firmware boot files from git..."
   fi

   mount $PARTITION1 $BOOTMP || {
      echo "ERROR: unable to mount partition $PARTITION1; aborting."
      exit 1
   }

   # bootcode.bin: 2nd stage bootloader, starts with SDRAM disabled
   # loader.bin  : 3rd stage bootloader, starts with SDRAM enabled
   # start.elf   : The GPU binary firmware image, provided by the foundation
   # kernel.img  : The OS kernel to load on the ARM processor.
   for f in bootcode.bin fixup_x.dat kernel.img start_x.elf LICENCE.broadcom; do
      if [ ! "$USEKERNEL" ]; then
         cp $SWHOME/firmware/boot/$f $BOOTMP/
      else
         curl -s -L -k https://github.com/raspberrypi/firmware/raw/master/boot/$f -o $BOOTMP/$f || {
            echo "ERROR: could not fetch $f from raspberrypi git; aborting."
         }
      fi
   done

   # cmdline.txt : parameters passed to the kernel on boot
   echo "\
dwc_otg.lpm_enable=0 \
root=/dev/mmcblk0p2 \
rootfstype=$SYSTEMFSTYPE \
quiet \
rootwait \
splash" > $BOOTMP/cmdline.txt

   # Create the kernel configuration file
   echo "\
arm_freq=800
force_turbo=1
gpu_mem=100
disable_overscan=1
start_file=start_x.elf
fixup_file=fixup_x.dat
framebuffer_depth=32
framebuffer_ignore_alpha=1" > $BOOTMP/config.txt

   ## Enable HDMI mode rather than DVI, if required
   #echo "\
   #hdmi_drive=2" >> $BOOTMP/config.txt

   # optional files:
   #   config.txt    : A configuration file read by the GPU.
   #                   Use this to override set the video mode, alter system clock
   #                   speeds, voltages, etc.
   #   vlls directory: Additional GPU code, e.g. extra codecs. 
   #                   Not present in the initial release.

   if [ ! "$USEKERNEL" ]; then
      # copy kernel modules
      mkdir -p $ROOTMP/lib/modules
      cp -a $SWHOME/firmware/modules/* $ROOTMP/lib/modules/
   fi
   [ "$SWHOME" -a "${SWHOME:0:4}" = "/tmp" ] && rm -rf $SWHOME

fi # rpimode

(
   cd $ROOTMP >/dev/null || {
      echo "ERROR: unable to access system partition mount point; aborting."
      umount $PARTITION2
      rmdir $ROOTMP
      exit 1
   }

   echo "Uncompressing $ROOTFSARCHIVE..."
   echo -n "This can take a while, please wait..."
   tar xf $ROOTFSARCHIVE || {
      echo
      echo "ERROR: unable to uncompress rootfs archive; aborting."
      exit 1
   }
   echo "DONE"

   if [ "$ADDITIONALARCHIVE" ]; then
      echo "Uncompressing $ADDITIONALARCHIVE..."
      echo -n "This can take a while, please wait..."
      tar xf $ADDITIONALARCHIVE || {
         echo
         echo "ERROR: unable to uncompress additional archive; aborting."
         exit 1
      } 
      echo "DONE"
   fi
)

[ $? -eq 0 ] || exit $?

if [ "$USEKERNEL" ]; then
   if [ "$RPIMODE" ]; then
      KERNELDEST=$BOOTMP/kernel.img
   else
      KERNELDEST=/tmp/$USEKERNEL
   fi
   if [ -r $ROOTMP/boot/$USEKERNEL ]; then
      cp $ROOTMP/boot/$USEKERNEL $KERNELDEST || {
         echo "ERROR: unable to find or copy kernel $USEKERNEL; aborting."
         exit 1
      }
      if [ -e /tmp/$USEKERNEL ]; then
         echo "Extracted kernel is available as $KERNELDEST"
      fi
   else
      echo "ERROR: can't find kernel file $ROOTMP/boot/$USERKERNEL; aborting."
      exit 1
   fi
fi

if [ "$SETHOSTNAME" ]; then
   echo "$SETHOSTNAME" > $ROOTMP/etc/hostname
   sed -i "s|HOSTNAME=.*|HOSTNAME=$SETHOSTNAME|" $ROOTMP/etc/sysconfig/network
   sed -i "s|\(127\.0\.0\.1\W*\)\(.*\)|\1 $SETHOSTNAME.localdomain $SETHOSTNAME localhost.localdomain localhost|" $ROOTMP/etc/hosts
fi

if [ "$ROOTPASSWORD" ]; then
   echo "$ROOTPASSWORD" | chroot $ROOTMP passwd root --stdin || {
      echo "ERROR: unable to set password for root; aborting."
      exit 1
   }
fi

if [ "$PRECHROOTCOMMAND" ]; then
   PRECHROOTCOMMAND=`echo $PRECHROOTCOMMAND | sed "s|@ROOT@|$ROOTMP|g"`
   $PRECHROOTCOMMAND || {
      echo "ERROR: error running pre-chroot command; aborting."
      exit 1
   }
fi

if [ "$CHROOTCOMMAND" ]; then
   chroot $ROOTMP $CHROOTCOMMAND || {
      echo "ERROR: error running chroot command; aborting."
      exit 1
   }
fi

sync
if [ "$RPIMODE" ]; then
   umount $PARTITION1
   rmdir $BOOTMP
fi
umount $PARTITION2
rmdir $ROOTMP

echo "Done."
exit 0
