#!/bin/bash

confdir=/etc/driverctl.d
bus=${SUBSYSTEM:-pci}
probe=1
save=1
debug=0

declare -A devclasses
devclasses=(["all"]=""
            ["storage"]="01"
            ["network"]="02"
            ["display"]="03"
            ["multimedia"]="04"
            ["memory"]="05"
            ["bridge"]="06"
            ["communication"]="07"
            ["system"]="08"
            ["input"]="09"
            ["docking"]="0a"
            ["processor"]="0b"
            ["serial"]="0c"
)

function log()
{
    echo "driverctl: $*" >&2
}

function debug()
{
    [ "$debug" -ne 0 ] && log "$@"
}

function error()
{
    log "$@"
    exit 1
}

function usage()
{
    echo "Usage: driverctl [-v] [--noprobe] [--nosave] set-override <device> <driver>"
    echo "       driverctl [-v] [--noprobe] [--nosave] unset-override <device>"
    echo "       driverctl [-v] [--noprobe] load-override <device>"
    echo "       driverctl [-v] list-devices"
    echo "       driverctl [-v] list-overrides"
    exit 1
}

function unbind()
{
    if [ -L "$syspath/driver" ]; then
        debug "unbinding previous driver $(basename "$(readlink "$syspath/driver")")"
        if ! echo "$dev" > "$syspath/driver/unbind"; then
            error "unbinding $dev failed"
        fi
    else
        debug "device $dev not bound" 
    fi
}

function probe_driver()
{
    debug "reprobing driver for $dev"
    echo "$dev" > "/sys/bus/$bus/drivers_probe"
}

function save_override()
{
    debug "saving driver override for $dev"
    if [ -n "$drv" ]; then
	[ -d "$confdir" ] || mkdir -p "$confdir"
        echo "$drv" > "$confdir/$sddev"
    else
        rm -f "$confdir/$sddev"
    fi
}

function list_devices()
{
    devices=()
    for d in "/sys/bus/$bus/devices"/*; do
        if [ -f "$d/driver_override" ]; then
            override="$(< "$d/driver_override")"
            if [ "$1" -eq 1 ] && [ "$override" == "(null)" ]; then
                continue
            fi
          
            line="$(basename "$d")"
            devices+=("$line")

            if [ -n "$2" ]; then
                class="$(< "$d/class")"
                [ "$2" == "${class:2:2}" ] || continue
            fi
            if [ -L "$d/driver" ]; then
                line+=" $(basename "$(readlink "$d/driver")")"
            else
                line+=" (none)"
            fi
            if [ "$1" -ne 1 ] && [ "$override" != "(null)" ]; then
                line+=" [*]"
            fi

            if [ $debug -ne 0 ]; then
                line+=" ($(udevadm info -q property "$d" | grep ID_MODEL_FROM_DATABASE | cut -d= -f2))"
            fi
            echo "$line"
        fi
    done
    if [ ${#devices[@]} -eq 0 ]; then
        error "No overridable devices found. Kernel too old?"
    fi
}

function set_override()
{
    if [ ! -f "$syspath/driver_override" ]; then
        error "device does not support driver override: $dev"
    fi
    if [ -n "$drv" ] && [ "$drv" != "none" ]; then
        debug "setting driver override for $dev: $drv"
        if [ ! -d "/sys/module/$drv" ]; then
            debug "loading driver $drv"
            /sbin/modprobe -q "$drv" || error "no such module: $drv"
        fi
    else
        debug "unsetting driver override for $dev"
    fi
    unbind
    echo "$drv" > "$syspath/driver_override"

    if [ "$drv" != "none" ] && [ $probe -ne 0 ]; then 
        probe_driver
        if [ ! -L "$syspath/driver" ]; then
            error "failed to bind device $dev to driver $drv"
        fi
    fi
}

while (($# > 0)); do
    case ${1} in
    --noprobe)
        probe=0
        ;;
    --nosave)
        save=0
        ;;
    --debug|--verbose|-v)
        debug=1
        ;;
    -b|--bus)
        bus=${2}
        shift
        ;;
    -h|--help|-*)
        usage
        ;;
    set-override)
        [ $# -ne 3 ] && usage
        cmd=$1
        dev=$2
        drv=$3
        break
        ;;
    load-override|unset-override)
        [ $# -ne 2 ] && usage
        cmd=$1
        dev=$2
        break
        ;;
    list-devices|list-overrides)
        [ $# -gt 2 ] && usage
        if [ -n "$2" ] && [ ! "${devclasses[$2]+_}" ]; then
            error "device type must be one of: ${!devclasses[*]}"
        fi
        cmd=$1
        devtype="${devclasses[${2:-all}]}"
        break
        ;;
    *)
        usage
        ;;
    esac
    shift
done

[ -n "$cmd" ] || usage

if [ -n "$dev" ]; then
        case ${dev} in
        */*)
            bus=${dev%%/*}
            dev=${dev#*/}
            ;;
        esac
        if [ -n "${DEVPATH:-}" ]; then
            devpath="$DEVPATH"
        else
            devlink="/sys/bus/$bus/devices/$dev"
            [ -L "$devlink" ] || error "no such device: $dev"
            devpath=$(realpath "$devlink" | cut -c5-)
        fi
        syspath="/sys/$devpath"
        sddev="$bus-$dev"
fi

case ${cmd} in
    load-override)
        if [ -s "$confdir/$sddev" ]; then
            drv=$(< "$confdir/$sddev")
	    set_override "$dev" "$drv"
        else
            exit 1
        fi
        ;;
    list-devices)
        list_devices 0 "$devtype"
        ;;
    list-overrides)
        list_devices 1 "$devtype"
        ;;
    set-override)
        set_override "$dev" "$drv"
        if [ $save -ne 0 ]; then
           save_override
        fi

        ;;       
    unset-override)
        set_override "$dev" ""
        if [ $save -ne 0 ]; then
           save_override
        fi
        ;;
esac
