#!/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 [-rpi|-cubox] [-a archive] [-c command] [-fs fstype] [-k kernel] [-p password] [-y]"
   echo
   echo "Target platform (optional):"
   echo "-rpi        : make a RaspberryPi SD with boot partition"
   echo "-cubox      : make a Cubox SD with boot partition"
   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 and use specified kernel file from ROOTFS/boot/ for boot"
   echo "-p password : set given password for root user"
   echo "-r mode     : set given default video mode where applicable (e.g. -r 1920x1080-32@60)"
   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 ;;
      -cubox) TARGETDEVICE=cubox ;;
      -rpi) TARGETDEVICE=rpi ;;
      -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
           ;;
      -r) [ "$2" ] || {
              echo "Error: -r option requires a video mode as argument; aborting."
              exit 1
           }
           VIDEOMODE=$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

case $TARGETDEVICE in
   rpi) BOOTFSTYPE=fat32
        BOOTFSSIZE=40M
        KERNELDEST="@BOOT@/kernel.img"
        ;;
   cubox) BOOTFSTYPE=ext3
          BOOTFSSIZE=20M
          KERNELDEST="@BOOT@/boot/uImage"
          ;;
   *) BOOTFSTYPE=
      if [ "$USEKERNEL" ]; then
         KERNELDEST="/tmp/$USEKERNEL"
      else
         KERNELDEST=
      fi
      PARTITION2=$PARTITION1 ;;
esac

[ -r "$ROOTFSARCHIVE" ] || {
   echo "ERROR: unable to read rootfs archive $ROOTFSARCHIVE; aborting."
   exit 1
}

echo "Rootfs archive : $ROOTFSARCHIVE"
echo "Target SD card : ${DRIVE}"
[ "$TARGETDEVICE" ] && echo "Target platform: $TARGETDEVICE"
[ "$BOOTFSTYPE" ] &&   echo "Boot fs        : $BOOTFSTYPE($BOOTFSSIZE)"
[ "$BOOTFSTYPE" ] &&   echo "System fs      : $SYSTEMFSTYPE"
[ "$USEKERNEL" ] &&    echo "Kernel         : $USEKERNEL -> $KERNELDEST"
echo

[ "$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 [ "$BOOTFSTYPE" ]; then
   # FAT partition
   sudo parted -s -a cyl -- ${DRIVE} mkpart primary $BOOTFSTYPE 0 $BOOTFSSIZE
   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 $BOOTFSSIZE -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 [ "$BOOTFSTYPE" ]; then
   echo "Formatting $PARTITION1 ($BOOTFSTYPE) ..."
   case "$BOOTFSTYPE" in
      vfat) mkfs.vfat -s 1 -F 32 -n "BOOT" $PARTITION1 || {
               echo "ERROR: unable to format $PARTITION1; aborting."
               exit 1
            } ;;
      ext*) mke2fs -q -t $BOOTFSTYPE -j -L "BOOT" $PARTITION1 || {
               echo "ERROR: unable to format $PARTITION1; aborting."
               exit 1
            } ;;
   esac
fi

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

[ "$BOOTFSTYPE" ] && BOOTMP=`mktemp -d -t rootfsinstall.XXXXXXXX`
ROOTMP=`mktemp -d -t rootfsinstall.XXXXXXXX`

[ -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 "[ "$ROOTMP" -a -e $ROOTMP ] && umount $ROOTMP && rmdir $ROOTMP;[ "$BOOTMP" -a -e $BOOTMP ] && umount $BOOTMP && rmdir $BOOTMP" 0 HUP INT QUIT ABRT KILL TERM

if [ "$BOOTFSTYPE" ]; then
   mount $PARTITION1 $BOOTMP || {
      echo "ERROR: unable to mount partition $PARTITION1; aborting."
      exit 1
   }
fi

if [ "$TARGETDEVICE" = "rpi" ]; 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

   # 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

elif [ "$TARGETDEVICE" = "cubox" ]; then
   [ "$VIDEO_MODE" ] || VIDEO_MODE="1920x1080-32@60"
   mkdir -p $BOOTMP/boot
   echo "Creating u-boot boot.scr..."
   cat > $BOOTMP/boot/boot.txt << _EOF
setenv bootargs 'console=ttyS0,115200n8 vmalloc=384M root=/dev/mmcblk0p2 video=dovefb:lcd0:$VIDEO_MODE-edid clcd.lcd0_enable=1 clcd.
lcd1_enable=0'
echo ======== Loading kernel ========
ext2load mmc 0:1 0x00200000 /boot/uImage
echo ======== Booting kernel ========
bootm
_EOF
   mkimage -A arm -O linux -T script -C none -n "u-boot commands" -d $BOOTMP/boot/boot.txt $BOOTMP/boot/boot.scr || {
      echo "ERROR: mkimage: error creating boot.scr; aborting."
      exit 1
   }
fi

(
   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 $?

KERNELDEST=`echo $KERNELDEST | sed "s|@BOOT@|$BOOTMP|g"`
if [ "$USEKERNEL" ]; then
   if [ -r $ROOTMP/boot/$USEKERNEL ]; then
      cp $ROOTMP/boot/$USEKERNEL $KERNELDEST || {
         echo "ERROR: unable to find or copy kernel $USEKERNEL; aborting."
         exit 1
      }
      echo "Kernel image available as $KERNELDEST"
   else
      echo "WARNING: can't find kernel file $ROOTMP/boot/$USERKERNEL"
   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

[ -e /usr/bin/qemu-arm.static ] && cp /usr/bin/qemu-arm.static $ROOTMP/usr/bin/qemu-arm

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

echo "Done."
exit 0