#!/bin/bash
# linaro-media-create - Create an installation media from a Linaro image

# Copyright 2010 Robert Nelson <robertcnelson@gmail.com>
# Copyright 2010 Linaro
# Based on rcn's setup_sdcard.sh script.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 3 as published
# by the Free Software Foundation.
#
# 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

set -e

MLO_FILE="MLO.omap"
UBOOT_FILE="u-boot.bin.omap"

unset MMC MMC1 MMC2 MMC3 IMAGE_FILE

ensure_command() {
# ensure_command foo foo-package
which "$1" 2>/dev/null 1>/dev/null || ( echo "Installing required command $1 from package $2" && sudo apt-get install "$2" )
}

ensure_command uuidgen uuid-runtime
ensure_command parted parted
ensure_command sfdisk util-linux
ensure_command fdisk util-linux
ensure_command wget wget
ensure_command mkfs.ext3 e2fsprogs
ensure_command mkfs.ext4 e2fsprogs
ensure_command mkimage uboot-mkimage
ensure_command md5sum coreutils
ensure_command realpath realpath

#Defaults
RFS=ext3
BINARY_TARBALL='binary-tar.tar.gz'
BOOT_LABEL=boot
RFS_LABEL=rootfs
CODENAME=Chessy
RFS_UUID=`uuidgen -r`
IS_LIVE=
FAT_SIZE=32
IMAGE_SIZE=2G

DIR=$PWD

get_mmcs_by_id() {
if [ ! ${IMAGE_FILE} ]; then
  for device in /dev/disk/by-id/*; do
    if [ `realpath $device` = $MMC ]; then
      if echo $device | grep -q -- "-part[0-9]*$"; then
        echo "device $MMC must not be a partition part ($device)" 1>&2
        exit 1
      fi
      for part_id in `ls $device-part*`; do
        part=`realpath $part_id`
        part_no=`echo $part_id | sed -e 's/.*-part//g'`
        # echo "part $part_no found: $part_id" 1>&2
        if test "$part_no" = 1; then
          MMC1=$part
        elif test "$part_no" = 2; then
          MMC2=$part
        elif test "$part_no" = 3; then
          MMC3=$part
        fi
      done
      break
    fi
  done
fi
}

prepare_sources() {
  if [ "$IS_LIVE" ]; then
    parts_dir=casper
    boot_snippet='boot=casper'
    [ "$IS_LOWMEM" ] && lowmem_opt=only-ubiquity
  else
    parts_dir=boot
    boot_snippet='root=UUID='${RFS_UUID}
  fi

  rm -rf ${DIR}/${parts_dir}/ || true
  rm -rf ${DIR}/binary/ || true
  rm -rf binary/initrd.img-* || true
  rm -rf binary/vmlinuz-* || true

  sudo tar -xf $BINARY_TARBALL binary/${parts_dir}

  ln -sf binary/${parts_dir}/initrd.img-* .
  ln -sf binary/${parts_dir}/vmlinuz-* .

  if [ "${IMAGE_FILE}" ]; then
    cat > binary/boot.cmd << BOOTCMD
setenv bootcmd 'mmc init; fatload mmc 0:1 0x80000000 uImage; bootm 0x80000000'
setenv bootargs '${serial_opts} ${splash_opts} earlyprintk fixrtc nocompcache ${lowmem_opt} root=/dev/mmcblk0p2 rootwait rw vram=12M omapfb.debug=y omapfb.mode=dvi:1280x720MR-16@60'
boot
BOOTCMD
  else
    cat > binary/boot.cmd << BOOTCMD
setenv bootcmd 'mmc init; fatload mmc 0:1 0x80000000 uImage; fatload mmc 0:1 0x81600000 uInitrd; bootm 0x80000000 0x81600000'
setenv bootargs '${serial_opts} ${splash_opts} earlyprintk fixrtc nocompcache ${lowmem_opt} ${boot_snippet} rootwait ro vram=12M omapfb.debug=y omapfb.mode=dvi:1280x720MR-16@60'
boot
BOOTCMD
  fi
}

cleanup_sd() {

 echo ""
 echo "Umounting Partitions"
 echo ""

 if test -n "$MMC1"; then
   sudo umount ${MMC1} &> /dev/null || true
 fi
 if test -n "$MMC2"; then
   sudo umount ${MMC2} &> /dev/null || true
 fi
 if [ "${MMC}" ]; then
   sudo parted -s ${MMC} mklabel msdos
 fi
}

create_partitions() {
if [ "${IMAGE_FILE}" ]; then
partdev=${IMAGE_FILE}
else
partdev=${MMC}
fi

if [ "$FAT_SIZE" = "32" ]; then
    PARTITION_TYPE="0x0C"
else
    PARTITION_TYPE="0x0E"
fi

# Create a VFAT or FAT16 partition of 9 cylinders which is about 64M
# and a linux partition of the rest

sudo sfdisk -D -H $HEADS -S $SECTORS $CYLINDER_ARG $partdev << THEEND
,9,$PARTITION_TYPE,*
,,,-
THEEND

if [ "${IMAGE_FILE}" ]; then
VFATOFFSET=$(($(fdisk -l -u $IMAGE_FILE | grep FAT | awk '{print $3}')*512))
VFATSIZE=$(($(fdisk -l -u $IMAGE_FILE | grep FAT | awk '{print $5}')*1024))
ROOTOFFSET=$(($(fdisk -l -u $IMAGE_FILE | grep Linux | awk '{print $2}')*512))
ROOTSIZE2=$(($(fdisk -l -u $IMAGE_FILE | grep Linux | awk '{print $3}')))
ROOTSIZE1=$(($(fdisk -l -u $IMAGE_FILE | grep Linux | awk '{print $2}')))
ROOTSIZE=$((((ROOTSIZE2-ROOTSIZE1)/2)*1024))
MMC1=$(sudo losetup -f --show $IMAGE_FILE --offset $VFATOFFSET --sizelimit $VFATSIZE)
MMC2=$(sudo losetup -f --show $IMAGE_FILE --offset $ROOTOFFSET --sizelimit $ROOTSIZE)
fi
}


prepare_partitions() {

echo ""
echo "Formating Boot Partition"
echo ""

sudo mkfs.vfat -F ${FAT_SIZE} ${MMC1} -n ${BOOT_LABEL}

echo ""
echo "Formating ${RFS} Partition"
echo ""
sudo mkfs.${RFS} -U "$RFS_UUID" ${MMC2} -L ${RFS_LABEL}
}

populate_boot() {
 echo ""
 echo "Populating Boot Partition"
 echo ""

 echo ""
 echo "Installing Boot Loader"
 echo ""

 if [ "$IS_LIVE" ]; then
   parts_dir=casper
 else
   parts_dir=boot
 fi

 mkdir -p ${DIR}/disk || true
 sudo mount ${MMC1} ${DIR}/disk
 case "$DEVIMAGE" in
   beagle|igep)
     if [ "$DEVIMAGE" = "beagle" ]; then
          if [ -e binary/${parts_dir}/${MLO_FILE} ] && [ -e binary/${parts_dir}/${UBOOT_FILE} ]; then
              sudo cp -v binary/${parts_dir}/${MLO_FILE}   ${DIR}/disk/MLO
              sudo cp -v binary/${parts_dir}/${UBOOT_FILE} ${DIR}/disk/u-boot.bin
          fi
     fi
     sync
     cd ${DIR}
     echo "done"

     sudo cp -f ${DIR}/binary/${parts_dir}/uImage.omap ${DIR}/disk/uImage
     sudo cp -f ${DIR}/binary/${parts_dir}/uInitrd.omap ${DIR}/disk/uInitrd

     sudo mkimage -A arm -O linux -T script -C none -a 0 -e 0 -n "$CODENAME 10.05" -d ${DIR}/binary/boot.cmd ${DIR}/disk/boot.scr
     sudo cp -v ${DIR}/disk/boot.scr ${DIR}/disk/boot.ini
     ;;
   vexpress)
     sudo tar --strip-components=5 -C ${DIR}/disk/ -xf $BINARY_TARBALL binary/usr/lib/u-boot/ca9x4_ct_vxp/u-boot.bin
     sudo mkimage -A arm -O linux -T kernel -C none -a 0x60008000 -e 0x60008000 \
          -n "Linux" -d ${DIR}/binary/${parts_dir}/vmlinuz*-vexpress ${DIR}/disk/uImage
     sudo mkimage -A arm -O linux -T ramdisk -C none -a 0x81000000 \
          -e 0x81000000 -n "initramfs" -d \
          ${DIR}/binary/${parts_dir}/initrd.img-*-linaro-vexpress \
          "${DIR}"/disk/uInitrd
     ;;
   *)
     echo "Set --dev parameter: --dev <beagle|igep|vexpress>" 1>&2
     cleanup_workdir
     exit
     ;;
 esac

 cd ${DIR}/disk/
 sync
 sync
 cd ${DIR}/

 sudo umount ${DIR}/disk || true
}

populate_rootfs() {
 echo ""
 echo "Populating rootfs Partition"
 echo "Be patient, this may take a few minutes"
 echo ""
 sudo mount ${MMC2} ${DIR}/disk

 sudo tar -xf $BINARY_TARBALL --strip-components=1 -C disk/

# add fstab entry for rootfs and boot
 echo "UUID=${RFS_UUID} / ${RFS}  errors=remount-ro 0 1 " | sudo tee -a ${DIR}/disk/etc/fstab

 if [ "$CREATE_SWAP" ] ; then

  echo ""
  echo "Creating SWAP File"
  echo ""

  SPACE_LEFT=$(df ${DIR}/disk/ | grep ${MMC2} | awk '{print $4}')

  let SIZE=$SWAP_SIZE*1024

  if [ $SPACE_LEFT -ge $SIZE ] ; then
   sudo dd if=/dev/zero of=${DIR}/disk/SWAP.swap bs=1M count=$SWAP_SIZE
   sudo mkswap ${DIR}/disk/SWAP.swap
   echo "/SWAP.swap  none  swap  sw  0 0" | sudo tee -a ${DIR}/disk/etc/fstab
   else
   echo "SWAP file bigger then whats left on partition"
  fi
 fi

 echo ""
 echo "Creating /etc/flash-kernel.conf"
 echo ""

 echo "UBOOT_PART=/dev/mmcblk0p1" > ${DIR}/disk/etc/flash-kernel.conf

 cd ${DIR}/disk/
 sync
 sync
 cd ${DIR}/

 sudo umount ${DIR}/disk || true
}

check_mmc() {
 FDISK=$(sudo LC_ALL=C sfdisk -l | grep "[Disk] ${MMC}" | awk '{print $2}')

 if test "-$FDISK-" = "-$MMC:-"
 then
  echo ""
  echo "I see..."
  echo "sudo sfdisk -l:"
  sudo LC_ALL=C sfdisk -l | grep "[Disk] /dev/" --color=never
  echo ""
  echo "mount:"
  mount | grep -v none | grep "/dev/" --color=never
  echo ""
  read -p "Are you 100% sure, on selecting [${MMC}] (y/n)? "
  [ "$REPLY" == "y" ] || exit
  echo ""
 else
  echo ""
  echo "Are you sure? I Don't see [${MMC}], here is what I do see..."
  echo ""
  echo "sudo sfdisk -l:"
  sudo LC_ALL=C sfdisk -l | grep "[Disk] /dev/" --color=never
  echo ""
  echo "mount:"
  mount | grep -v none | grep "/dev/" --color=never
  echo ""
  exit
 fi
}

check_fs_type() {
 IN_VALID_FS=1

 if test "-$FS_TYPE-" = "-ext2-"
 then
 RFS=ext2
 unset IN_VALID_FS
 fi

 if test "-$FS_TYPE-" = "-ext3-"
 then
 RFS=ext3
 unset IN_VALID_FS
 fi

 if test "-$FS_TYPE-" = "-ext4-"
 then
 RFS=ext4
 unset IN_VALID_FS
 fi

 if test "-$FS_TYPE-" = "-btrfs-"
 then
 RFS=btrfs
 unset IN_VALID_FS
 fi

 if [ "$IN_VALID_FS" ] ; then
   usage
 fi
}

calculatesize() {
    IMAGE_SIZE=${IMAGE_SIZE/G/M*1024}
    IMAGE_SIZE=${IMAGE_SIZE/M/K*1024}
    IMAGE_SIZE=${IMAGE_SIZE/K/*1024}
    IMAGE_SIZE=$(($IMAGE_SIZE))
}

setup_sizes() {
    HEADS=255
    SECTORS=63
    SECTORSIZE=512
    if [ "${IMAGE_FILE}" ]; then
    calculatesize
    CYLINDERSIZE=$(($HEADS*$SECTORS*$SECTORSIZE))
    CYLINDERS=$(($IMAGE_SIZE/$CYLINDERSIZE))
    CYLINDER_ARG="-C $CYLINDERS"
    IMAGE_SIZE=$(($CYLINDERS*$CYLINDERSIZE))
    else
    CYLINDER_ARG=""
    fi
}

setup_image() {
    sudo qemu-img create -f raw $IMAGE_FILE $IMAGE_SIZE
}

usage() {
    echo "usage: $(basename $0) --mmc /dev/sdd"
    echo "<or>"
    echo "usage: $(basename $0) --image_file mmc.img"
cat <<EOF

required options:
--mmc </dev/sdX>
    Unformated MMC Card
<or>
--image_file <xxx>
    specify name of image file

--dev <board>
    use development boot options; this includes setting up serial ttys as well
    as enabling normal debug options for the target board. Current board values:
    * beagle
    * igep
    * vexpress

Additional/Optional options:
-h --help
    this help

--rootfs <fs_type>
    ext2
    ext3 - <set as default>
    ext4
    btrfs

--boot_label <boot_label>
    boot partition label

--rfs_label <rfs_label>
    rootfs partition label

--swap_file <xxx>
    Creats a Swap file of (xxx)MB's

--live
    Create boot command for casper/live images; if this is not
    provided a UUID for the rootfs is generated and used as the root=
    option

--live-256m
    Create boot command for casper/live images; adds only-ubiquity option
    to allow use of live installer on boards with 256M memory - like beagle

--console <ttyXY>
    add a console to kernel boot parameter; this parameter can be defined
    multiple times.

--image_size nnnG or nnnM
    specify size of SD image to create; use with --image_file only (default: 2G)

--binary <filename>
    specify file used to create the bootable system. Default binary-tar.tar.gz
EOF
exit
}

checkparm() {
    if [ "$(echo $1|grep ^'\-')" ];then
        echo "E: Need an argument"
        usage
    fi
}

cleanup_workdir() {
  if [ -d binary/ ]; then
    sudo rm -rf binary/
  fi
}


consoles=""

# parse commandline options
while [ ! -z "$1" ]; do
    case $1 in
        -h|--help)
            usage
            MMC=1
            ;;
        --mmc)
            checkparm $2
            MMC="$2"
            check_mmc
            ;;
        --image_file)
            checkparm $2
            IMAGE_FILE="$2"
            ;;
        --image_size)
            checkparm $2
            IMAGE_SIZE=$2
            ;;
        --rootfs)
            checkparm $2
            FS_TYPE="$2"
            check_fs_type
            ;;
        --boot_label)
            checkparm $2
            BOOT_LABEL="$2"
            ;;
        --rfs_label)
            checkparm $2
            RFS_LABEL="$2"
            ;;
        --swap_file)
            checkparm $2
            SWAP_SIZE="$2"
            CREATE_SWAP=1
            ;;
        --live)
            IS_LIVE=1
            ;;
        --live-256m)
            IS_LIVE=1
            IS_LOWMEM=1
            ;;
        --console)
            checkparm $2
            consoles="$consoles $2"
            ;;
        --dev)
            checkparm $2
            DEVIMAGE=$2
            ;;
        --binary)
            checkparm $2
            BINARY_TARBALL="$2"
            ;;
    esac
    shift
done

setup_sizes
if [ "${IMAGE_FILE}" ]; then
  setup_image
fi

serial_opts=""
if [ "$consoles" ]; then
  for c in ${consoles}; do
    serial_opts="$serial_opts console=$c"
  done
  if [ "$IS_LIVE" ]; then
    serial_opts="$serial_opts serialtty=ttyS2"
  fi
fi

if [ "$DEVIMAGE" ]; then
  case "$DEVIMAGE" in
    beagle|igep)
      serial_opts="$serial_opts console=tty0 console=ttyS2,115200n8"
      if [ "$IS_LIVE" ]; then
        serial_opts="$serial_opts serialtty=ttyS2"
      fi
      ;;
    vexpress)
      serial_opts="$serial_opts console=tty0 console=ttyAMA0,38400n8"
      if [ "$IS_LIVE" ]; then
        serial_opts="$serial_opts serialtty=ttyAMA0"
      fi
      # ARM Boot Monitor is used to load u-boot, uImage etc. into flash and
      # only allows for FAT16
      FAT_SIZE=16
      ;;
    *)
      echo "unknown --dev paramater: $DEVIMAGE" 1>&2
      ;;
  esac
else
  if [ "$IS_LIVE" ]; then
    splash_opts="quiet splash"
  fi
fi

if [ ! "${MMC}" -a ! "${IMAGE_FILE}" ]; then
    usage
fi

 prepare_sources
 get_mmcs_by_id
 cleanup_sd
 create_partitions
 echo -n "waiting for partitioning to settle ..."
 sync
 sleep 3
 echo "done."
 get_mmcs_by_id
 if test -z "$MMC1" -o -z "$MMC2"; then
   echo "MMC1: $MMC1 nor MMC2: $MMC2 must be empty"
   exit 2
 fi
 prepare_partitions
 populate_boot
 populate_rootfs

 if [ "${IMAGE_FILE}" ]; then
  echo "Create ${IMAGE_FILE}.gz"
  gzip -f ${IMAGE_FILE} > ${IMAGE_FILE}.gz
 fi

 cleanup_workdir

