#!/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 device 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 "The option -plug_test can determine the desired target device by"
  echo "inquiring the system with unplugged device and then with plugged one."
  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 "Option -dummy prevents this kind of real action and rather shows the"
  echo "unmount and write commands on stdout."
  echo
  echo "Options:"
  echo " -plug_test         Find the target device by asking the user to press"
  echo "                    Enter when the desired target is _not_ plugged in,"
  echo "                    to then plug it in, and to press Enter again."
  echo "                    This overrides device names and option -list_all."
  echo "                    The found device is then shown with advice, vendor,"
  echo "                    and model. Option -DO_WRITE is obeyed if given."
  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 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 " -dummy             Report the -DO_WRITE actions but do not perform"
  echo "                    them."
  echo " -dummy_force       If a single device name is given, do a run of"
  echo "                    -dummy -DO_WRITE even against the advice of"
  echo "                    this program. This probably shows you ways to"
  echo "                    shoot your own foot."
  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 " $0 -with_sudo -image_file debian-live-10.0.0-amd64-xfce.iso -DO_WRITE -plug_test"
  echo
}

# Roughly convert human readable sizes and plain numbers to 1 / million
round_down_div_million() {
  sed \
    -e 's/^[0-9][0-9][0-9][0-9][0-9][0-9]$/0/' \
    -e 's/^[0-9][0-9][0-9][0-9][0-9]$/0/' \
    -e 's/^[0-9][0-9][0-9][0-9]$/0/' \
    -e 's/^[0-9][0-9][0-9]$/0/' \
    -e 's/^[0-9][0-9]$/0/' \
    -e 's/^[0-9]$/0/' \
    -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=
  dummy_run=
  dummy_force=
  do_plug_test=

  # 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 "FAILURE: Cannot obtain size of -image_file '$i'" >&2
        exit 13
      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" = "-plug_test"
    then
      do_plug_test=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" = "-dummy"
    then
      dummy_run=y
    elif test "$i" = "-dummy_force"
    then
      dummy_run=y
      do_write=y
      dummy_force=y
    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
      num=$(echo "$i" | wc -w)
      if test "$num" = 1
      then
        devs_named=y
        devs="$devs $i"
        show_reasons=y
      else
        echo "$0 : Given device name is not a single word: '$i'" >&2
        exit 12
      fi
    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
    echo "Testing sudo to possibly get password prompting done now:" >&2
    if sudo "$lsblk_cmd" -h >/dev/null
    then
      echo "sudo $lsblk_cmd seems ok." >&2
      echo >&2
      sudo_cmd=sudo
      have_su_power=y
    else
      echo "FAILURE: Cannot execute program $lsblk_cmd by sudo" >&2
      exit 11
    fi
  fi
}

## Obtain a blank separated list of top-level names which do not look like
## CD, floppy, RAM dev, or loop device.
collect_devices() {
  "$lsblk_cmd" -d -n -o NAME \
    | grep -v '^sr[0-9]' \
    | grep -v '^fd[0-9]' \
    | grep -v '^zram[0-9]' \
    | grep -v '^loop[0-9]' \
    | tr '\n\r' '  '
}

## Trying to find the desired device by watching plug-in effects
plug_in_watcher() {
  found_devices=
  echo >&2
  echo "Caused by option -plug_test: Attempt to find the desired device" >&2
  echo "by watching it appear after being plugged in." >&2
  echo >&2
  echo "Step 1:" >&2
  echo "Please make sure that the desired target device is plugged _out_ now." >&2
  echo "If it is currently plugged in, make sure to unmount all its fileystems" >&2
  echo "and then unplug it." >&2
  echo "Press the Enter key when ready." >&2
  read dummy
  old_device_list=' '$(collect_devices)' '

# <<< Mock-up to save USB socket wear-off by erasing items from old_device_list
# <<< Their presence in new_device_list will let them appear as fresh plugs
# old_device_list=' '$(echo -n $old_device_list | sed -e 's/sdc//')' '

  echo "Found and noted as _not_ desired: $old_device_list" >&2
  echo >&2
  echo "Step 2:" >&2
  echo "Please plug in the desired target device and then press the Enter key." >&2
  read dummy
  echo -n "Waiting up to 10 seconds for a new device to be listed ..." >&2
  end_time="$(expr $(date +'%s') + 10)"
  while test $(date +'%s') -le "$end_time"
  do
    new_device_list=' '$(collect_devices)' '
    if test "$old_device_list" = "$new_device_list"
    then
      sleep 1
      echo -n '.' >&2
    else
      for i in $new_device_list
      do
        if echo "$old_device_list" | grep -F -v ' '"$i"' ' >/dev/null
        then
          found_devices="$found_devices $i"
        fi
        # Break the waiting loop
        end_time=0
      done
    fi
  done
  echo >&2
  if test -z "$found_devices"
  then
    echo "SORRY: No new candidate device was found." >&2
    return 8
  fi
  num=$(echo $found_devices | wc -w)
  if test "$num" -gt 1
  then
    echo "SORRY: More than one new candidate device appeared: $found_devices" >&2
    return 9
  fi
  echo "Found and noted as desired device: $found_devices" >&2
  if test -n "$devs"
  then
    echo "(-plug_test is overriding device list given by arguments: $devs )" >&2
  fi
  if test -n "$list_all"
  then
    echo "(-plug_test is overriding -list_all)" >&2
    list_all=
  fi
  devs_named=y
  with_vendor_model=y
  show_reasons=y
  devs=$(echo -n $found_devices)
  echo >&2
  return 0
}

## 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, floppy, RAM dev
    devs=$(collect_devices)
  fi
  
  not_advised=0
  for name in $devs
  do
    # Collect reasons
    yucky=
    reasons=
    good_trans=
    good_fs=
    bad_trans=
    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- "
    elif echo "$name" | grep '^sr[0-9]' >/dev/null
    then
      yucky=y
      reasons="${reasons}looks_like_cd_drive- "
    elif echo "$name" | grep '^fd[0-9]' >/dev/null
    then
      yucky=y
      reasons="${reasons}looks_like_floppy- "
    elif echo "$name" | grep '^loop[0-9]' >/dev/null
    then
      yucky=y
      reasons="${reasons}looks_like_loopdev- "
    elif echo "$name" | grep '^zram[0-9]' >/dev/null
    then
      yucky=y
      reasons="${reasons}looks_like_ramdev- "
    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" | grep -F "/" >/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
    for trans in $transports
    do
      if test "$trans" = "usb" -o "$trans" = "mmcblk"
      then
        good_trans="${trans}+"
      elif test -n "$trans"
      then
        bad_trans="$trans"
        yucky=y
        if test "$transports" = "not_an_expected_name"
        then
          dummy=dummy
        else
          if echo "$reasons" | grep -F -v "not_usb" >/dev/null
          then
            reasons="${reasons}not_usb- "
          fi
        fi
      fi
    done
    if test -z "$good_trans" -a -z "$bad_trans"
    then
      yucky=y
      reasons="${reasons}no_bus_info- "
    elif test -z "$bad_trans"
    then
      reasons="${reasons}$good_trans "
    fi

    # 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" | grep -F -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" | grep -F -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" | grep -F -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
        reasons="${reasons}look_for_iso++ "
      else
        yucky=y
        reasons="${reasons}no_iso9660- "
      fi
    elif test -n "$has_other"
    then
      yucky=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 : $reasons$descr"
      else
        echo "$name"
      fi
    fi
  done
  return 0;
}

## Puts list of mounted (sub-)devices of $1 into $mounted_devs
list_mounted_of() {
  partitions=$("$lsblk_cmd" -l -n -p -o NAME /dev/"$1" \
               | grep -v '^'/dev/"$1"'$' \
               | tr '\n\r' '  ')
  mounted_devs=
  for i in /dev/"$1" $partitions
  do
    # Show the found mount lines and add their device paths to list
    mount_line=$(mount | grep '^'"$i"' ')
    if test -n "$mount_line"
    then
      echo "  $mount_line"
      mounted_devs="$mounted_devs $i"
    fi
  done
}

## Does the work of unmounting and dd-ing
write_image() {

  if test -z "$umount_cmd"
  then
    echo "No executable program umount found in: $sudo_x_dir_list" >&2
    return 6
  fi
  echo "Looking for mount points of $2:"
  mounted_devs=
  list_mounted_of "$2"

  if test -n "$dummy_force"
  then
    echo "AGAINST THE ADVICE BY THIS PROGRAM, a daring user could do:"
    dummy_run=y
  elif test -n "$dummy_run"
  then
    echo "Would do if not -dummy:"
  fi
  if test -n "$mounted_devs"
  then
    for i in $mounted_devs
    do
      if test -n "$dummy_run"
      then
        echo "  $sudo_cmd $umount_cmd $i"
      else
        if $sudo_cmd "$umount_cmd" "$i"
        then
          echo "Unmounted: $i"
        else
          echo "FAILURE: Non-zero exit value with:  $sudo_cmd $umount_cmd $i" >&2
          return 7
        fi
      fi
    done

    # Check again if any mount points still exist
    if test -z "$dummy_run"
    then
      list_mounted_of "$2"
      if test -n "$mounted_devs"
      then
        echo "FAILURE: $sudo_cmd $umount_cmd could not remove all mounts: $mounted_devs" >&2
        return 7
      fi
    fi
  fi

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

  # >>> ??? Erase possible GPT backup table at end of device ?

  if test -n "$dummy_force"
  then
    echo "BE SMART. BE CAUTIOUS. BEWARE."
  fi
  return 0
}

# main()

reset_job
arg_interpreter "$@"

if test -n "$do_plug_test"
then
  plug_in_watcher
  ret=$?
  if test "$ret" -ne 0
  then
    exit $ret
  fi
fi

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_cand=$(echo $candidates | wc -w)
  num_devs=$(echo $devs| wc -w)
  if test -n "$dummy_force" -a "$num_devs" -ne 1
  then
    echo "SORRY: Refusing -dummy_force with not exactly one device given." >&2
    exit 10
  fi
  if test -n "$dummy_force" -a -n "$dummy_run" -a "$num_cand" -ne 1 
  then
    # -dummy_force in a situation where the program would normally refuse
    echo
    echo "Overriding any advice because of -dummy_force"
    candidates="$devs"
    num_cand=1
  elif test -n "$dummy_force"
  then
    # Downgrade -dummy_force to -dummy in order to avoid the ugly warning
    dummy_force=
    dummy_run=y
  fi
  if test "$num_cand" -eq 1
  then
    if test -n "$image_file"
    then
      if test -n "$do_plug_test"
      then
        echo >&2
        echo "Step 3:" >&2
        if test -n "$dummy_run"
        then
          echo "This would be the last chance to abort. Enter the word 'yes' to see -dummy report." >&2
        else
          echo "Last chance to abort. Enter the word 'yes' to start REAL WRITING." >&2
        fi
        read dummy
        if test "$dummy" = "yes" -o "$dummy" = "'yes'" -o "$dummy" = '"yes"'
        then
          dummy=dummy
        else
          echo "WRITE RUN PREVENTED by user input '$dummy'." >&2
          exit 12  
        fi
      fi
      write_image "$image_file" $candidates
      exit $?
    else
      if test -n "$dummy_run"
      then
        echo "Would simulate writing to /dev/$candidates if an -image_file were given."
      else
        echo "Would write to /dev/$candidates if an -image_file were given."
      fi
      exit 0
    fi
  elif test "$num_cand" -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: Unplug the unwanted devices from the machine," >&2
    echo "      or work with option -plug_test," >&2
    echo "      or add the desired name out of {$(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
    else
      echo "NO CANDIDATE FOR TARGET DEVICE FOUND !" >&2
    fi
    echo "Overall available devices:" >&2
    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