#!/usr/bin/env bash
# Script to bind the VFIO stub driver to PCI devices
# By Trent Arcuri, 2017
#

tmpfile=$(mktemp)
genfile='/etc/modprobe.d/vfioselect.conf'
online=false
declare -A lspciid

usage()
{
  echo """Usage: $(basename $0) [<arguments>]
  -t,--tmpfile  <path>   Specify path to temporary file
  -g,--genfile  <path>   Specify path to modprobe rule
  -o,--online            Attempt to bind/unbind without the need to reboot
  -h,--help              Display this help dialogue"""
}

cleanup()
{
  rm $tmpfile

  if [ -e $genfile.old ]; then
    rm $genfile.old
  fi
}

format()
{
  lspci -x > $tmpfile

  IFS=$'\n' 	# make newlines the only separator
  for line in $(lspci -F $tmpfile -n)
  do
    lspcimach+=("$line")
    lspcibus+=("$(echo $line|awk '{print $1}')")
  done

  for line in "${lspcibus[@]}"
  do
    lspciid[$line]="$(lspci -F $tmpfile -ns $line|awk '{print $3}')"
  done

  for line in $(lspci -F $tmpfile -mm)
  do
    lspcihuma+=("line")
    lspcifm+=("$(echo $line|awk -F\" '{print $4" "$2": "$6}')")
  done
  > $tmpfile
}

makelist()
{
  if [ -e $genfile ]; then
    cp $genfile $genfile.old
  else
    touch $genfile
  fi

  IFS=$'\n' 	# make newlines the only separator
  store=0
  for pci in "${lspcibus[@]}"
  do
    if grep -Eqi "${lspciids[$store]}" "$genfile" && grep -Eqi $pci "$genfile"
    then
      menuitems+=("$pci" "${lspcifm[$store]}" "on")
    else
      menuitems+=("$pci" "${lspcifm[$store]}" "off")
    fi
    let store++
  done

  if dialog --keep-tite --checklist \
      "Choose which devices to bind:" 30 180 30 \
      "${menuitems[@]}" 2>$tmpfile; then
    return 0
  else
    return 1
  fi
}


makefile()
{
  printf "#$(cat $tmpfile)\n" > $genfile
  printf "options vfio-pci ids=" >> $genfile
  IFS=" "
  for line in $(cat $tmpfile); do
    printf "${lspciid[$line]}," >> $genfile
  done
  truncate -s-1 $genfile #Fix trailing comma
  echo "$genfile written successfully."
}

bind()
{
  #Unbind from device drivers, and bind to VFIO
  for id in ${bindids[@]};do
    echo "Binding $id"
    #Get pci slot
    pci="0000:$id"
    pcipath="/sys/bus/pci/devices/$pci"
    #Rebinding process
    if [ -e $pcipath ]; then
      echo "$pci" > "$pcipath/driver/unbind"
    fi
    echo "${lspciid[$id]}"|sed 's/:/ /'> /sys/bus/pci/drivers/vfio-pci/new_id
  done
}

unbind()
{
  # Unbind from VFIO, and bind to device drivers
  for id in ${unbindids[@]};do
    echo "Unbinding $id"
    echo 1 > /sys/bus/pci/devices/0000:$id/remove
  done
  echo 1 > /sys/bus/pci/rescan
}

online()
{
  IFS=$'\n'
  #Get list of new and old ID's
  for id in $(grep -Eoh "[0-9a-f]{2}:[0-9a-f]{2}\.[0-9]" $genfile); do
    newids+=("$id")
  done
  for id in $(grep -Eoh "[0-9a-f]{2}:[0-9a-f]{2}\.[0-9]" $genfile.old); do
    oldids+=("$id")
  done

  #Get list of different ID's, determine what to do with them
  for id in $(echo ${newids[@]} ${oldids[@]} | sed 's/ /\n/g'| sort | uniq -u)
  do
    diffids+=("$id")
  done
  for id in $(echo ${diffids[@]} ${oldids[@]} | sed 's/ /\n/g'| sort | uniq -d)
  do
    unbindids+=("$id")
  done
  for id in $(echo ${diffids[@]} ${newids[@]} | sed 's/ /\n/g'| sort | uniq -d)
  do
    bindids+=("$id")
  done

  echo "Bind:   ${bindids[@]}"
  echo "Unbind: ${unbindids[@]}"
  echo

  #Bind/Unbind the ID's respectively
  if [[ ${bindids[@]} ]]; then
    bind
  fi
  if [[ ${unbindids[@]} ]]; then
    unbind
  fi
}

main()
{
  #Check dependencies
  if which dialog > /dev/null ; then
    :
  else
    dialog
    exit 1
  fi

  #Argument formatting
	local argv=("$@");
	# argc is the count of arguments
	local argc=${#argv[@]};

	# this is important to ensure globbing is active
	shopt -s extglob;

	# Handle compressed short options
	re="(^| )\\-[[:alnum:]]{2,}"; # regex to detect shortoptions
	# we evaluate this as a long string, thus ${argv[*]}, instead of ${argv[@]}
	if [[ "${argv[*]}" =~ $re ]]; then
		local compiled_args=();
		for ((i=0; i<argc; i++)); do
			if [[ "${argv[$i]}" =~ $re ]]; then
				local compressed_args="${argv[$i]#*-}";
				for ((r=0; r<${#compressed_args}; r++)); do
					compiled_args+=("-${compressed_args:$r:1}");
				done
				shift;
				compiled_args+=("$@");
				## recurse
				main "${compiled_args[@]}";
				## we "pass" the exit code back up the recursions to the OS
				exit $?;
			fi
			compiled_args+=("${argv[$i]}");
			shift;
		done
		exit;
	fi

	#Check arguments
	while [ "$1" != "" ]; do
		case $1 in
			-t | --tmpfile )
				shift
				cleanup
				tmpfile=$1
				;;
			-g | --genfile )
				shift
				genfile=$1
				;;
			-o | --online )
				online=true
				;;
			-h | --help )
				usage
        exit
        ;;
      -?*)
        printf 'WARN: Unknown option (ignored): %s\n' "$1" >&2
		esac
		shift
	done

	#Check for root access
	if [[ $EUID -ne 0 ]]; then
		echo "This program requires root access."
		exit 1
	else
    format
		if makelist; then
			makefile
		else
			cleanup
			exit 130
		fi
		if $online; then
			online
		fi
		cleanup
	fi
}

main "$@"