auto-krokodil/kroko-cli.sh

#!/bin/bash

set -e
set +f
set -u
set -o pipefail
# lspci -nn | grep VGA
# 01:00.0 VGA compatible controller [0300]: NVIDIA Corporation GF108 [GeForce GT 440] [10de:0de0] (rev a1)

# constants
readonly EXIT_ENODRIVER=10
readonly EXIT_ENOGPU=20

# when called from GUI, $KROKO_TMPDIR must be set by GUI
# (to be able ti reuse it)
if [ -z "${KROKO_TMPDIR:-}" ]; then
	#trap 'rm -fr "$KROKO_TMPDIR"' EXIT
	:
fi
KROKO_TMPDIR="${KROKO_TMPDIR:-$(mktemp -d)}"

_echo_err(){
	echo "$@" 1>&2
}

_mktemp(){
	mktemp --tmpdir="$KROKO_TMPDIR" "$@"
}

# $1: string
# $2: symbol
# Count how many times a symbol is inside the string
_count_symbol_in_string(){
	local o
	o="$1"
	local counter
	counter=0
	for (( i = 0; i <= ${#o}; i++ ))
	do
		if [ "${o:$i:1}" = "$2" ]; then
			counter=$((++counter))
		fi
	done
	echo "$counter"
}

# $1: string
# $2: number of block with between [] to extract from
# Example:
# $1: ff [ty] xz [56] yu
# if $2=1; then "ty" is returned
# if $2=2; then "56" is returned
_element_from_string(){
	local o
	o="$1"
	local target
	target="$2"
	local new
	new=""
	local inside_value
	inside_value=0
	local counter
	counter=0
	for (( i = 0; i <= ${#o}; i++ ))
	do
		if [ "${o:$i:1}" = '[' ]; then
			counter=$((++counter))
			if [ "$counter" != "$target" ]; then
				continue
			fi
			inside_value=1
			continue
		fi
		if [ "${o:$i:1}" = ']' ]; then
			inside_value=0
			continue
		fi
		if [ "$inside_value" != 1 ]; then
			continue
		fi
		# processing symbols between [ and ]
		new="$new${o:$i:1}"
	done
	echo "$new"
}

# $1: line from lspci -nn
_extract_device_id(){
	local n
	n="$(_count_symbol_in_string "$1" [)"
	local o
	o="$(_element_from_string "$1" "$n")"
	# 10de:0de0 -> 0de0
	IFS=":" read -a arr <<< "$o"
	echo "${arr[1]}"
}

# $1: line from lspci -nn
_extract_vendor_id(){
	local n
	n="$(_count_symbol_in_string "$1" [)"
	local o
	o="$(_element_from_string "$1" "$n")"
	# 10de:0de0 -> 10de
	IFS=":" read -a arr <<< "$o"
	echo "${arr[0]}"
}

# Input to stdin: output of `lspci -nn`
# $1: path to file with list of NVIDIA vendor IDs
# Outputs lines about GPUs
# Comment from ubuntu-drivers-common:
# Display controllers are device class 03
# There are 4 subclasses:
#   00	VGA compatible controller
#   01	XGA compatible controller
#   02	3D controller
#   80	Display controller
_filter_gpus(){
	while read -r line
	do
		local o
		o="$(_element_from_string "$line" 1)"
		# See if it is a GPU
		if [ "${o:0:2}" != "03" ]; then
			continue
		fi
		# See if it is an NVIDIA GPU
		id="$(_extract_vendor_id "$line")"
		# 10de is NVIDIA's vendor ID
		# XXX It is the only vendor ID?
		if [ "$id" = "10de" ]; then
			echo "$line"
		fi
	done
}

# $1: arch (e.g. x86_64)
# $2: path to file for output
_dnf_mk_file(){
	dnf repoquery \
		--arch "$1" \
		--whatprovides 'nvidia-blob-*' \
		--qf 'NAME %{name}\n%{provides}' \
	> "$2"
}

# $1: input file (output of _dnf_mk_file())
# $2: path to directory with temp files
_sort_provides_by_pkg(){
	local file=""
	while read -r line
	do
		if [[ "$line" =~ ^"NAME " ]]; then
			local arr
			IFS=" " read -a arr <<< "$line"
			file="$2"/provides___"${arr[1]}"
			continue
		fi
		echo "$line" >> "$file"
	done < <(cat "$1")
}

# $1: device id
# $2: directory with files with provides
# output: nvidia390,nvidia470
_get_available_drivers(){
	grep "^nvidia-blob-devid(${1}) =" "$2"/provides___* | \
	awk -F ':' '{print $1}' | awk -F 'provides___' '{print $2}' | \
	sort -u | \
	tr '\n' ',' | sed -e 's/,$//'
}

# $1: device id
# $2: directory with files with provides
# output: nvidia470
_get_best_driver(){
	grep "^nvidia-blob-devid(${1}) =" "$2"/provides___* | \
	awk -F ':' '{print $1}' | awk -F 'provides___' '{print $2}' | \
	sort -u -r | \
	head -n1
}


# $1: string
# Convert string of `lspci -nn` to a human-readable name
_line2name(){
	local o
	o="$1"
	local n
	n="$(_count_symbol_in_string "$1" ])"
	local counter=0
	local started=0
	local human_name=""
	for (( i = 0; i <= ${#o}; i++ ))
	do
		if [ "${o:$i:1}" = ']' ]; then
			counter=$((++counter))
		fi
		if [ "$counter" = 0 ]; then
			continue
		fi
		# Name starts after [0300]:
		if [ "$started" != 1 ] && [ "${o:$i:1}" = ' ' ] && [ "${o:$i-1:1}" = ':' ] && [ "${o:$i-2:1}" = ']' ]; then
			started=1
			continue
		fi
		if [ "$started" != 1 ]; then
			continue
		fi
		if [ "${o:$i:1}" = ' ' ] && [ "${o:$i+1:1}" = '[' ] &&  [ $((counter+1)) = "$n" ]; then
			break
		fi
		human_name="${human_name}${o:$i:1}"
	done
	echo "$human_name"
}

_cli_get_gpus(){
	local arch
	arch="$(rpm -E "%_arch")"
	if [ -z "$arch" ]; then
		_echo_err "Error getting architecture of the host machine"
		return 1
	fi
	local big_file
	big_file="$(_mktemp)"
	_dnf_mk_file "$arch" "$big_file"
	_sort_provides_by_pkg "$big_file" "$KROKO_TMPDIR"
	local vendors_ids_file
	local o
	o="$(lspci -nn | _filter_gpus)"
	if [ "$(echo "$o" | grep -c .)" -le 0 ]; then
		echo "No GPUs found"
		return $EXIT_ENOGPU
	fi
	while read -r line
	do
		local device_id
		device_id="$(_extract_device_id "$line")"
		if [ -z "$device_id" ]; then
			_echo_err "Error extracting device ID"
			continue
		fi
		local human_name
		human_name="$(_line2name "$line")"
		if [ -z "$human_name" ]; then
			_echo_err "Error converting to human readable name"
			return 1
		fi
		local available_drivers
		available_drivers="$(_get_available_drivers "$device_id" "$KROKO_TMPDIR")"
		if [ -z "$available_drivers" ]; then
			_echo_err "No drivers found for $device_id"
			continue
		fi
		local best_driver
		best_driver="$(_get_best_driver "$device_id" "$KROKO_TMPDIR")"
		if [ -z "$best_driver" ]; then
			_echo_err "Error getting the best driver"
			continue
		fi
		echo "${line};${device_id};${human_name};${available_drivers};${best_driver}"
	done <<< "$o"
}

_main(){
	case "$1" in
		"get-gpus" )
			_cli_get_gpus
		;;
	esac
}

if [ "${SOURCING:-0}" != 1 ]; then
	_main "$@"
fi