omap-usb-install

#!/bin/bash

# From https://github.com/networkimprov/omap3-usb-boot-install
#
# Install script to cold flash over USB using an image made of
# subimages for MLO, u-boot, script, dtb, zimage and initrd.
#
# If a sub-image is not specified, an empty one is created.
#
# This tool uses modified omap34usbboot.tar.gz that's available at:
# https://groups.google.com/forum/#!topic/pandaboard/9z5ebHEnuqs
#
# The MLO is generated by u-boot when it's compiled with:
# $ CONFIG_SPL=1 CROSS_COMPILE=arm-linux-gnueabi- make
#
# The u-boot install script is created from install.src with:
# $ mkimage -T script -C none -n "Install script" -d install.src u-boot.script
#
# The mkimage tool comes with u-boot or can be installed with
# u-boot-tools package.

SZ_512K=524288
SZ_1M=1048576
SZ_2M=2097152
SZ_8M=8388608
SZ_10M=10485760
SZ_20M=20971520

update=0
loop=0
tests=0
install=0
install_image=""
dieid=""
rootfs_file="" # user input

main() {
  check_options $@

  cd $(dirname $0)

  if [ "$update" == "1" ]; then
    echo "Update cold flash scripts..."
    update_scripts
  fi

  generate_install_image

  while [ "$install" == "1" ]; do
    if [ "$rootfs_file" == "erase" ]; then
      echo "Wiping u-boot in NAND..."
      cold_flash $install_image
    elif [ "$rootfs_file" == "kernel" ]; then
      echo "Booting kernel without flashing..."
      cold_flash $install_image
    else
      echo "Cold flashing over USB..."
      cold_flash $install_image

      # Valid dieid length must be 32
      if [ ${#dieid} != "32" ]; then
        error_exit "ERROR: Bad dieid $dieid ${#dieid}"
      fi

      disk="/dev/disk/by-id/usb-Linux_File-Stor_Gadget_"$dieid"-0:0"
      uart="/dev/serial/by-id/usb-Anvl_Multi_Gadget_"$dieid"-if02"

      wait_for_device 45 b $disk

      echo "Debug UART at $uart"

      if [ "$rootfs_file" == "" ]; then
        echo "WARNING: No rootfs specified, assuming manual install"
      else
        if [ "$rootfs_file" != "initramfs" ] && [ "$rootfs_file" != "recovery" ] && 
	    [ ! -f "$rootfs_file" ]; then
          error_exit "ERROR: Could not find file (or absolute path) for rootfs $rootfs_file"
        fi
        install_on_emmc $dieid $disk
      fi

      if [ "$tests" == "1" ]; then
        if [ ! -c $uart ]; then
          error_exit "ERROR: Could not find UART $uart"
        fi
        echo "Rebooting..."
        echo "reboot -f" > $uart
        wait_for_device 60 c "/dev/serial/by-id/usb-Anvl_Multi_Gadget-if02" # look for network device?

        while ! ifconfig usb0 | grep -q RUNNING; do sleep 1; done
        sudo ifconfig usb0 192.168.99.1/30 up # move to host config
        echo "Running test scripts..."
        if ! ssh root@192.168.99.2 < scripts/led_colors.sh; then
          if [ "$loop" == "1" ]; then
            echo "Set aside this device for repair, and connect the next device"
            continue
          else
            break
          fi
        fi
      fi
    fi

    if [ "$loop" == "1" ]; then
      echo "Done, please connect the next device"
    else
      echo "Done, please power cycle the device"
      break
    fi
  done

  rm -f $install_image

  exit 0
}

options_exit() {
  echo "usage: $0 [-u | --update] [-i | --install[=][rootfs.tar.gz | erase | kernel | initramfs | recovery]] [-t | --tests] [-l | --loop]"
  exit 1
}

error_exit() {
  test -f $install_image && rm -f $install_image
  echo $1
  exit 1
}

check_options() {
  if [ $# -lt 1 ]; then
    options_exit
  fi

  if ! options=$(getopt -o ui::lt -l update,install::,loop,tests -- "$@"); then
    options_exit
  fi

  eval set -- "$options"

  while [ $# -gt 0 ]; do
    case $1 in
    -u|--update)
      update=1
    ;;
    -i|--install)
      install=1
      rootfs_file=$2
      shift
    ;;
    -l|--loop)
      loop=1
    ;;
    -t|--tests)
      tests=1
    ;;
    (--)
      shift
      break
    ;;
    (-*)
      error_exit "Unknown option $1"
    ;;
    *)
      options_exit
    ;;
    esac
    shift
  done
}

cold_flash() {
  image=$1

  # Check for a locally compiled omap34usbboot binary first
  if [ -f src/omap34usbboot/out/usbboot ]; then
    prefix=src/omap34usbboot/out
  else
    arch=$(uname -m)
    case "$arch" in
    i?86)
      prefix=arch/i386
    ;;
    x86_64)
      prefix=arch/x86_64
    ;;
    *)
      error_exit "Unsupported architecture, please build src/omap34usbboot manually"
    ;;
    esac
  fi

  while read -r output; do
    echo $output
    if echo $output | grep "No such file or directory" > /dev/null; then
      exit 2
    fi 
    if echo $output | grep "command not found" > /dev/null; then
      exit 3
    fi 
    if echo $output | grep "unexpected 2ndstage response" > /dev/null; then
      exit 4
    fi 
    dieid=$(echo $output | awk '{ split($0,a,"dieid: "); print substr(a[2], 1, 32) }')
  done < <(sudo $prefix/usbboot $prefix/omap3_aboot $image 2>&1)

  if [ "$dieid" = "" ]; then
    error_exit "ERROR: failed to get dieid"
  fi

  return 0
}

append_image() {
  file="$1"
  pad_size="$2"

  if [ ! -f "$file" ]; then
    echo "WARNING: file $file missing, creating empty $file"
    touch "$file"
  fi

  size=$(du -b "$file" | cut -f 1)
  if [ "$size" -gt "$pad_size" ]; then
    error_exit "ERROR: $file too big"
  fi

  imgsize=$(du -b $install_image | cut -f 1)
  let imgsize+=pad_size

  if ! cat "$file" >> $install_image; then
    error_exit "ERROR: could not add $file to $install_image"
  fi

  if ! truncate -s $imgsize $install_image; then
    error_exit "ERROR: truncate failed for $file"
  fi
}

generate_install_image() {
  if ! install_image=$(mktemp); then
    error_exit "ERROR: Could not create install image file"
  fi

  append_image images/u-boot.bin $SZ_512K    # at 0x80800000
  append_image images/MLO $SZ_512K           # at 0x80880000
  append_image images/u-boot.img $SZ_512K    # at 0x80900000

  # One of these depending of the command line options
  if [ "$rootfs_file" == "erase" ]; then
    append_image images/erase.scr $SZ_512K   # at 0x80980000
    return 0
  fi

  if [ "$rootfs_file" == "kernel" ]; then
    append_image images/kernel.scr $SZ_512K  # at 0x80980000
  elif [ "$rootfs_file" == "recovery" ]; then
    append_image images/recovery.scr $SZ_512K  # at 0x80980000
  else
    append_image images/boot.scr $SZ_512K    # at 0x80980000
  fi

  append_image images/dtb $SZ_2M             # at 0x80a00000
  append_image images/zimage $SZ_10M         # at 0x80c00000
  append_image images/initramfs $SZ_20M      # at 0x81600000
}

update_scripts() {
  if ! mkimage -A arm -T script -C none -n "Install script" -d src/scripts/boot.script images/boot.scr; then
    error_exit "ERRROR: could not mkimage install script"
  fi

  if ! mkimage -A arm -T script -C none -n "Erase script" -d src/scripts/erase.script images/erase.scr; then
    error_exit "ERRROR: could not mkimage erase script"
  fi

  if ! mkimage -A arm -T script -C none -n "Kernel boot script" -d src/scripts/kernel.script images/kernel.scr; then
    error_exit "ERRROR: could not mkimage kernel boot script"
  fi

  if ! mkimage -A arm -T script -C none -n "Recovery install script" -d src/scripts/recovery.script images/recovery.scr; then
    error_exit "ERRROR: could not mkimage recovery install script"
  fi

  return 0
}

wait_for_device() {
  timeout=$1
  type=$2
  path=$3
  echo "Waiting for device $path..."
  while [ ! -$type $path ]; do
    let timeout--
    if [ "$timeout" -lt "1" ]; then
      error_exit "ERROR: Timed out waiting for $path"
    fi
    sleep 1
  done

  return 0
}

#
# Runs a shell script over the UART using socat.
# Will fail by default and only shows the following lines
# printed by the shell script:
# ERROR: Some message
# SUCCESS: Some message
# Returns 0 only if the shell script prints out a line
# SUCCESS: Some message
# Sorry, no passing of arguments to the script at this point.
#
serial_run_shell_script() {
  uart=$1
  test_script=$2
  verbose=$3
  ret=1

  if [ ! -f $test_script ]; then
    error_exit "Could not find script $test_script"
  fi

  while read -r output; do
    if [ "$verbose" == "1" ]; then
      echo $output
    fi
    if echo $output | grep "ERROR: " > /dev/null; then
      echo $output
      return 1
    fi
    if echo $output | grep "SUCCESS: " > /dev/null; then
      echo $output
      ret=0
    fi
  done < <(cat $test_script | sudo socat - $uart,ispeed=115200,ospeed=9600,raw,echo=0 2>&1)

  return $ret
}

install_on_emmc() {
  dieid=$1
  emmc=$2

  if [ "$rootfs_file" == "initramfs" ] || [ "$rootfs_file" == "recovery" ]; then
    partition=0
    format=0
    untar=0
  else
    partition=1
    format=1
    untar=1
  fi

  if [ "$partition" == "1" ]; then
    echo "Partitioning $emmc..."
    if ! sudo /sbin/parted --script $emmc mklabel msdos; then
      error_exit "ERROR: Could not mklabel"
    fi
    if ! sudo /sbin/parted --script $emmc mkpart primary fat32 2048s 128MB; then
      error_exit "ERROR: Could not add fat32 partition"
    fi
    if ! sudo /sbin/parted --script $emmc set 1 boot on; then
      error_exit "ERROR: Could not set active partition"
    fi
    if ! sudo /sbin/parted --script $emmc mkpart primary ext4 128MB 100%; then
      error_exit "ERROR: Could not add ext4 partition"
    fi
  fi

  sleep 3 # let possible automount finish
  test -b $emmc-part1 && sudo umount $emmc-part1
  test -b $emmc-part2 && sudo umount $emmc-part2
  sleep 3

  if [ "$format" == "1" ]; then
    echo "Formatting new partitions on $emmc..."
    if ! sudo mkfs.vfat -F 32 $(readlink -f $emmc-part1) > /dev/null; then
      error_exit "ERROR: Could not mkfs.vfat"
    fi
    if ! sudo mkfs.ext4 -O ^64bit -q $(readlink -f $emmc-part2); then
      error_exit "ERROR: Could not mkfs.ext4"
    fi
  fi

  if ! vfat=$(mktemp -d /tmp/$dieid-vfat-XXXXXXXX); then
    error_exit "ERROR: Could not mktemp for vfat"
  fi
  if ! sudo mount $emmc-part1 $vfat; then
    error_exit "ERROR: Could not mount vfat"
  fi
  echo "Mounted vfat to temporary location $vfat"

  if ! ext4=$(mktemp -d /tmp/$dieid-ext4-XXXXXXXX); then
    error_exit "ERROR: Could not mktemp for ext4"
  fi
  if ! sudo mount $emmc-part2 $ext4; then
    error_exit "ERROR: Could not mount ext4"
  fi
  echo "Mounted ext4 to temporary location $ext4"

  echo "Populating VFAT partition with u-boot, kernel and initramfs..."
  if ! sudo cp images/{MLO,u-boot.img,u-boot.bin,uEnv.txt,dtb,zimage,initramfs} $vfat/; then
    error_exit "ERROR: Could not copy files to VFAT"
  fi

  if [ "$untar" == "1" ]; then
    echo "Untarring root file system to $ext4..."
    if sudo tar xf "$rootfs_file" -C $ext4 2>&1 | tee /dev/stderr | grep -q ': Cannot '; then
      error_exit "ERROR: Untarring failed for $rootfs_file"
    fi
  fi

  echo "Syncing, unmounting and cleaning up tmp directories..."
  sync
  if ! sudo umount $ext4; then
    error_exit "ERROR: Could not umount ext4"
  fi
  if ! sudo umount $vfat; then
    error_exit "ERROR: Could not umount vfat"
  fi
  test -d $ext4 && rmdir $ext4 # umount may unlink /tmp mountpoint
  test -d $vfat && rmdir $vfat
  return 0
}

#
# Invoke main function above
#
# sudo here and remove other sudo's?

main $@