#!/bin/sh

# Copyright (c) 2019
# Nio Wiklund alias sudodus <nio dot wiklund at gmail dot com>
# Thomas Schmitt <scdbackup@gmx.net>
# Provided under GPL version 2 or later.

# Check whether we are on GNU/Linux
if uname -s | grep -v '^Linux' >/dev/null
then
  echo "This program is entirely specialized on Linux kernel device names." >&2
  echo "Found to be on: '$(uname -s)'" >&2
  exit 2
fi

# Accept sudo-executable commands only in well known directories.
# (Listed with increasing priority.)
lsblk_cmd=
dd_cmd=
if test "$(whoami)" = "root"
then
  sudo_x_dir_list="/usr/bin /bin /usr/sbin /sbin"
else
  sudo_x_dir_list="/usr/sbin /sbin /usr/bin /bin"
fi
for i in $sudo_x_dir_list
do
  if test -x "$i"/lsblk
  then
    lsblk_cmd="$i"/lsblk
  fi
  if test -x "$i"/dd
  then
    dd_cmd="$i"/dd
  fi
  if test -x "$i"/umount
  then
    umount_cmd="$i"/umount
  fi
done
if test -z "$lsblk_cmd"
then
  echo "No executable program lsblk found in: $sudo_x_dir_list" >&2
  exit 5
fi

print_usage() {
  echo "usage:  $0 [options] [device_name [device_name ...]]"
  echo
  echo "Looks on GNU/Linux for USB and Memory Card devices and evaluates"
  echo "whether the found devices are plausible targets for image copying."
  echo "If no device names and no -list_all are given, then a plain list of"
  echo "advisable device names is printed to stdout. One per line."
  echo "Device names must not begin by '-' and must be single words. They must"
  echo "not contain '/'. E.g. 'sdc' is valid, '/dev/sdc' is not valid."
  echo "If devices names are given, then they get listed with advice shown."
  echo "If one of the given device names gets not advised, the exit value is 1."
  echo
  echo "Only if option -DO_WRITE is given and -list_all is not, and if exactly"
  echo "one advisable device is listed, it really gets overwritten by the"
  echo "file content of the given -image_file. In this case the exit value"
  echo "is zero if writing succeeded, non-zero else."
  echo
  echo "Options:"
  echo " -list_all          Print list of all found devices with advice, vendor"
  echo "                    and model, One per line. Ignore any device names."
  echo "                    Ignore -DO_WRITE."
  echo " -with_vendor_model Print vendor and model with each submitted device."
  echo
  echo " -max_size n[M|G|T] Set upper byte size limit for advisable devices."
  echo "                    Plain numbers get rounded down to full millions."
  echo "                    Suffix: M = million, G = billion, T = trillion."
  echo "                    Be generous to avoid problems with GB < GiB."
  echo " -min_size n[M|G|T] Set lower byte size limit for advisable devices."
  echo "                    After processing like with -max_size, one million"
  echo "                    gets added to the size limit."
  echo " -look_for_iso      Demand presence of an ISO 9660 filesystem. If so,"
  echo "                    any further filesystem type is acceptable on that"
  echo "                    device. Else only ISO 9660 and VFAT are accepted."
  echo " -with_sudo         Run '$lsblk_cmd -o FSTYPE' by sudo."
  echo "                    If so, consider to enable this without password."
  echo "                    If no filesystems are detected and the program"
  echo "                    has no superuser power, the device is not advised."
  echo "                    If -DO_WRITE is given, run umount and dd by sudo." 
  echo " -image_file PATH   Set the path of the image file which shall be"
  echo "                    written to a device. Its size will be set as"
  echo "                    -min_size."
  echo " -DO_WRITE          Write the given -image_file to the one advisable"
  echo "                    device that is found. If more than one such device"
  echo "                    is found, then they get listed but no writing"
  echo "                    happens. In this case, re-run with one of the"
  echo "                    advised device names to get a real write run."
  echo
  echo " -help              Print this text to stdout and then end the program."
  echo "Examples:"
  echo " $0 -with_sudo -list_all"
  echo " $0 sdc"
  echo " $0 -with_sudo -image_file debian-live-10.0.0-amd64-xfce.iso -DO_WRITE"
  echo
}

# Roughly convert human readable sizes and plain numbers to 1 / million
round_down_div_million() {
  sed -e 's/\.[0-9]*//' -e 's/[0-9][0-9][0-9][0-9][0-9][0-9]$//' \
      -e 's/[Mm]//' -e 's/[Gg]/000/' -e 's/[Tt]/000000/'
}

### Assessing arguments and setting up the job

# Settings
reset_job() {
  list_all=
  show_reasons=
  look_for_iso=
  devs=
  devs_named=
  max_size=
  with_vendor_model=
  with_sudo=
  image_file=
  do_write=

  # Status
  sudo_cmd=
  have_su_power=
}

arg_interpreter() {
  next_is=
  for i in "$@"
  do
    # The next_is option parameter readers get programmed by the -options
    if test "$next_is" = "max_size"
    then
      max_size="$(echo "$i" | round_down_div_million)"
      next_is=
    elif test "$next_is" = "min_size"
    then
      min_size="$(echo "$i" | round_down_div_million)"
      min_size="$(expr $min_size + 1)"
      next_is=
    elif test "$next_is" = "image_file"
    then
      image_file="$i"
      min_size="$(stat -c '%s' "$i" | round_down_div_million)"
      if test -z "$min_size"
      then
        echo "WARNING: Cannot obtain size of -image_file '$i'" >&2
      else
        min_size="$(expr $min_size + 1)"
      fi
      next_is=
    elif test "$i" = "-list_all"
    then
      list_all=y
      with_vendor_model=y
      show_reasons=y
    elif test "$i" = "-max_size"
    then
      next_is="max_size"
    elif test "$i" = "-min_size"
    then
      next_is="min_size"
    elif test "$i" = "-with_vendor_model"
    then
      with_vendor_model=y
    elif test "$i" = "-look_for_iso"
    then
      look_for_iso=y
    elif test "$i" = "-with_sudo"
    then
      with_sudo=y
    elif test "$i" = "-image_file"
    then
      next_is="image_file"
    elif test "$i" = "-DO_WRITE"
    then
      do_write=y
    elif test "$i" = "-help"
    then
      print_usage
      exit 0
    elif echo "$i" | grep -v '^-' >/dev/null
    then
      devs_named=y
      devs="$devs $i"
      show_reasons=y
    else
      echo "$0 : Unknown option: $i" >&2
      echo >&2
      print_usage >&2
      exit 1
    fi
  done

  # Predict superuser power. Possibly enable sudo with lsblk -o FSTYPE and dd.
  if test "$(whoami)" = "root"
  then
    have_su_power=y
  elif test -n "$with_sudo"
  then
    sudo_cmd=sudo
    have_su_power=y
  fi
}

## Evaluation of available devices and suitability
list_devices() {
  if test -n "$list_all"
  then
    devs=
  fi
  if test -z "$devs"
  then
    # Obtain list of top-level names which do not look like CD or floppy.
    devs=$($lsblk_cmd -d -n -o NAME  | grep -v '^sr[0-9]' | grep -v '^fd[0-9]')
  fi
  
  not_advised=0
  for name in $devs
  do
    # Collect reasons
    tasty=
    yucky=
    reasons=
    good_trans=
    good_fs=
    bad_fs=

    # Unwanted device name patterns
    if (echo "$name" | grep '^sd[a-z][1-9]' >/dev/null) \
       || (echo "$name" | grep '^mmcblk.*p[0-9]' >/dev/null) \
       || (echo "$name" | grep '^nvme.*p[0-9]' >/dev/null)
    then
      yucky=y
      reasons="${reasons}looks_like_disk_partition- "
    fi
    if echo "$name" | grep '^sr[0-9]' >/dev/null
    then
      yucky=y
      reasons="${reasons}looks_like_cd_drive- "
    fi

    # >>> recognize the device from which Debian Live booted

    # Connection type. Normally by lsblk TRAN, but in case of mmcblk artificial.
    if echo "$name" | grep '^mmcblk[0-9]' >/dev/null
    then
      transports="mmcblk"
    elif echo "$name" | fgrep "/" >/dev/null
    then
      echo "NOTE: The device name must not contain '/' characters" >&2 
      transports=not_an_expected_name
      reasons="${reasons}name_with_slash- "
    else
      transports=$($lsblk_cmd -n -o TRAN /dev/"$name")
    fi
    if test -z "$transports"
    then
      transports=lsblk_failed_tran
    fi
    for trans in $transports
    do
      if test "$trans" = "usb" -o "$trans" = "mmcblk"
      then
        good_trans="${trans}+"
        tasty=y
      elif test -n "$trans"
      then
        yucky=y
        if test "$transports" = "not_an_expected_name"
        then
          dummy=dummy
        else
          if echo "$reasons" | fgrep -v "not_usb" >/dev/null
          then
            reasons="${reasons}not_usb- "
          fi
        fi
      fi
    done

    # Wanted or unwanted filesystem types
    fstypes=$($sudo_cmd $lsblk_cmd -n -o FSTYPE /dev/"$name")
    if test "$?" -gt 0
    then
      fstypes="lsblk_fstype_error"
    fi
    # Get overview of filesystems
    has_iso=
    has_vfat=
    has_other=
    for fstype in $fstypes
    do
      if test "$fstype" = "iso9660"
      then
        has_iso=y
        if echo "$good_fs" | fgrep -v "has_$fstype" >/dev/null
        then
          good_fs="${good_fs}has_${fstype}+ "
        fi
      elif test "$fstype" = "vfat"
      then
        has_vfat=y
        if echo "$good_fs" | fgrep -v "has_$fstype" >/dev/null
        then
          good_fs="${good_fs}has_${fstype}+ "
        fi
      elif test -n "$fstype"
      then
        has_other=y
        if echo "$bad_fs" | fgrep -v "has_$fstype" >/dev/null
        then
          bad_fs="${bad_fs}has_${fstype}- "
        fi
      fi
    done
    # Decide whether the found filesystems look dispensible enough
    reasons="${reasons}${good_fs}${bad_fs}"
    if test "${bad_fs}${good_fs}" = "" -a -z "$have_su_power"
    then
      yucky=y
      reasons="${reasons}no_fs_while_not_su- "
    elif test -n "$look_for_iso"
    then
      if test -n "$has_iso"
      then
        tasty=y
        reasons="${reasons}look_for_iso++ "
      else
        yucky=y
        reasons="${reasons}no_iso9660- "
      fi
    elif test -n "$has_other"
    then
      yucky=y
    else
      tasty=y
    fi
  
    # Optional tests for size
    if test -n "$max_size" -o -n "$min_size"
    then
      size=$($lsblk_cmd -n -b -o SIZE /dev/"$name" | head -1 | round_down_div_million)
      if test -z "$size"
      then
        yucky=y
        reasons="${reasons}lsblk_no_size- "
      fi
    fi
    if test -n "$max_size" -a -n "$size"
    then
      if test "$size" -gt "$max_size"
      then
        yucky=y
        reasons="${reasons}size_too_large- "
      fi
    fi
    if test -n "$min_size" -a -n "$size"
    then
      if test "$size" -lt "$min_size"
      then
        yucky=y
        reasons="${reasons}size_too_small- "
      fi
    fi
  
    # Now decide overall and report
    descr=
    if test -n "$with_vendor_model"
    then
      descr=": "$($lsblk_cmd -n -o VENDOR,MODEL /dev/"$name" | tr '\n\r' '  ' | tr -s ' ')
    fi
    if test -n "$yucky"
    then
      if test -n "$show_reasons"
      then
        echo "$name : NO  : $reasons$descr"
      fi
      not_advised=1
    else
      if test -n "$show_reasons"
      then
        echo "$name : YES : $good_trans $reasons$descr"
      else
        echo "$name"
      fi
    fi
  done
  return 0;
}

write_image() {

  echo "--- not activated yet ---"

  if test -z "$umount_cmd"
  then
    echo "No executable program umount found in: $sudo_x_dir_list" >&2
    return 6
  fi
  partitions=$($lsblk_cmd -n -p -o NAME /dev/"$2" | grep -v '^'/dev/"$2"'$' \
               | sed -e 's/[^a-zA-Z0-9_+@:.,/-]//g' | tr '\n\r' '  ')
  echo "Would do:   $sudo_cmd $umount_cmd /dev/"$2" $partitions"

  if test -z "$dd_cmd"
  then
    echo "No executable program dd found in: $sudo_x_dir_list" >&2
    return 6
  fi
  echo "Would do:   $sudo_cmd $dd_cmd if='${1}' bs=1M of=/dev/${2} ; sync"

  echo "--- not activated yet ---"

  return 0
}

reset_job
arg_interpreter "$@"
list_devices
if test -n "$list_all"
then
  dummy=dummy
elif test -n "$do_write"
then
  with_vendor_model=
  show_reasons=
  candidates=$(list_devices | tr '\n\r' '  ')
  num=$(echo $candidates | wc -w)
  if test "$num" -eq 1
  then
    if test -n "$image_file"
    then
      write_image "$image_file" $candidates
      exit $?
    else
      echo "Would write to /dev/$candidates if an -image_file were given." >&2
      exit 0
    fi
  elif test "$num" -gt 1
  then
    echo "WILL NOT WRITE ! More than one candidate found for target device:" >&2
    show_reasons=y
    with_vendor_model=y
    devs="$candidates"
    list_devices >&2
    echo "HINT: Re-run with one of the names {$(echo $candidates | sed -e 's/ /,/g')} as additional argument." >&2
    exit 3
  else
    if test -n "$devs_named"
    then
      echo "NO CANDIDATE FOR TARGET DEVICE AMONG THE GIVEN NAMES !" >&2
      echo "Overall available devices:"
    else
      echo "NO CANDIDATE FOR TARGET DEVICE FOUND !" >&2
    fi
    list_all=y
    show_reasons=y
    with_vendor_model=y
    list_devices >&2
    exit 4
  fi
fi

if test -n "$devs"
then
  exit $not_advised
fi