ldto

#!/bin/bash
# SPDX-License-Identifier: GPL-2.0
# Copyright (C) 2021 Da Xue <da@libre.computer>
# PURPOSE: Manage the application and removal of device tree overlays in the Linux kernel.

set -e

if [ -z "$VENDOR" ]; then
	if [ ! -e /sys/class/dmi/id/board_vendor ]; then
		echo "No vendor found!" >&2
		exit 1
	fi
	VENDOR=$(tr -d '\0' < /sys/class/dmi/id/board_vendor)
fi


if [ -z "$BOARD" ]; then
	if [ ! -e /sys/class/dmi/id/board_name ]; then
		echo "No board name found!" >&2
		exit 1
	fi
	BOARD=$(tr -d '\0' < /sys/class/dmi/id/board_name)
fi

PATH_OF_CONFIG=/sys/kernel/config/device-tree/overlays
PATH_OF_DT=/sys/firmware/fdt
PATH_OF_DTB_OVERRIDE=
PATH_OF_EFI_DTB=

if [ ! -d "$PATH_OF_CONFIG" ]; then
	echo "The running kernel does not support this tool." >&2
	exit 1
fi

cd $(dirname $(readlink -f "${BASH_SOURCE[0]}"))

PATH_DTBO=$VENDOR/$BOARD/dt
PATH_DTCFG=$VENDOR/$BOARD/dt.config
MAP_FILE="$VENDOR/$BOARD/dt.map"
TAB_CHAR='	'
declare -A MAP_DT_ALIAS
declare -A MAP_ALIAS_DT

LDTO_loadMap()
{
	if [ -e "$MAP_FILE" ]; then
		while IFS= read -r mapping; do
			local dt_alias=${mapping%%$TAB_CHAR*}
			local dt_file=${mapping##*$TAB_CHAR}
			MAP_DT_ALIAS[$dt_file]="$dt_alias"
			MAP_ALIAS_DT[$dt_alias]="$dt_file"
		done < <(grep -v '^#' "$MAP_FILE")
	fi
}

LDTO_getDTAlias()
{
	if [[ -v "MAP_DT_ALIAS[$1]" ]]; then
		echo "${MAP_DT_ALIAS[$1]}"
	fi
}

LDTO_getAliasDT()
{
	if [[ -v "MAP_ALIAS_DT[$1]" ]]; then
		echo "${MAP_ALIAS_DT[$1]}"
	fi
}

LDTO_getDTBO()
{
	local dtbo_file="$PATH_DTBO/$1.dtbo"
	if [ -f "$dtbo_file" ]; then
		echo "$dtbo_file"
	fi
}

LDTO_list()
{
	local dtbos="$PATH_DTBO"/*.dtbo
	if [ -z "$(ls $dtbos 2> /dev/null)" ]; then
		echo "No overlays detected. Did you run make?" >&2
		return 1
	fi
	LDTO_loadMap
	echo "#Overlays available" >&2
	for i in $dtbos
	do
		local dto=$(basename $i .dtbo)
		echo "$dto	$(LDTO_getDTAlias $dto)"
	done
}

LDTO_checkOFConfig()
{
	if [ $# -eq 0 ]; then
		if ! test -r "$PATH_OF_CONFIG"; then
			echo "$FUNCNAME: OF path not readable! sudo?" >&2
			return 1
		fi
	else
		if ! test -w "$PATH_OF_CONFIG"; then
			echo "$FUNCNAME: OF path not writable! sudo?" >&2
			return 1
		fi
	fi
}

LDTO_enable()
{
	if [ -z "$1" ]; then
		LDTO_list
		return 1
	fi
	LDTO_checkOFConfig RW
	LDTO_loadMap
	for dto in "$@"; do
		local dto_real=$(LDTO_getAliasDT "$dto")
		if [ ! -z "$dto_real" ]; then
			local dto=$dto_real
		fi
		local dtbo_file=$(LDTO_getDTBO "$dto")
		if [ -z "$dtbo_file" ]; then
			echo "$FUNCNAME: $dto does not exist and cannot be added" >&2
			return 1
		fi
		if [ -e "$PATH_OF_CONFIG/$dto" ]; then
			echo "Overlay $dto: already exists" >&2
			return 1
		fi
		mkdir "$PATH_OF_CONFIG/$dto"
		cat "$dtbo_file" > "$PATH_OF_CONFIG/$dto/dtbo"
		echo "Overlay $dto: `cat $PATH_OF_CONFIG/$dto/status`" >&2
	done
}

LDTO_disable()
{
	if [ -z "$1" ]; then
		LDTO_active
		return 1
	fi
	LDTO_checkOFConfig RW
	LDTO_loadMap
	for dto in "$@"; do
		local dto_real=$(LDTO_getAliasDT "$dto")
		if [ ! -z "$dto_real" ]; then
			local dto=$dto_real
		fi
		if [ ! -e "$PATH_OF_CONFIG/$dto" ]; then
			echo "Overlay $dto: does not exist and cannot be removed" >&2
			return 1
		fi
		rmdir $PATH_OF_CONFIG/$dto
		echo "Overlay $dto: removed" >&2
	done
}

LDTO_status()
{
	echo "#Overlays active:" >&2
	for i in `ls $PATH_OF_CONFIG`;
	do
		basename $i
	done
}

LDTO_active()
{
	LDTO_checkOFConfig
	if [ -z "$1" ]; then
		LDTO_status
	else
		if [ -z "$1" ]; then
			return 1
		fi
		for i in `ls $PATH_OF_CONFIG`;
		do
			if [ "$1" = $(basename $i) ]; then
				return 0
			fi
		done
		return 1
	fi
}

LDTO_importDTConfig()
{
	if [ -f "$PATH_DTCFG" ]; then
		. "$PATH_DTCFG"
	else
		echo "$FUNCNAME: board device tree configuration cannot be found." >&2
		return 1
	fi
}

LDTO_checkFirmwareDT()
{
	if [ ! -e "$PATH_OF_DT" ]; then
		echo "$FUNCNAME: system device tree cannot be found." >&2
		return 1
	fi
	if ! test -r "$PATH_OF_DT"; then
		echo "$FUNCNAME: system device tree cannot be read! sudo?" >&2
		return 1
	fi
}

LDTO_findEFIDTBPath()
{
	local mnt_boot=$(cut -f 2 -d " " /proc/mounts  | grep ^/boot)
	local mnt_boot_count=$(echo "$mnt_boot" | wc -l)
	if [ $mnt_boot_count -lt 1 ]; then
		echo "$FUNCNAME: no mounts found under /boot directory." >&2
		return 1
	fi
	for mnt_dir in $mnt_boot; do
		if [ -d "$mnt_dir"/EFI/BOOT -o -e "$mnt_dir"/uboot.env ]; then
			PATH_OF_EFI_DTB="$mnt_dir"/dtb
			if [ $# -ne 0 ]; then
				if [ -d "$PATH_OF_EFI_DTB" ]; then
					if ! test -w "$PATH_OF_EFI_DTB"; then
						echo "$FUNCNAME: EFI path not writable! sudo?" >&2
						return 1
					fi
				else
					if ! test -w "$mnt_dir"; then
						echo "$FUNCNAME: EFI path not writable! sudo?" >&2
						return 1
					fi
					if [ ! -d "$PATH_OF_EFI_DTB" ]; then
						mkdir -p "$PATH_OF_EFI_DTB"
					fi
				fi
			fi
			return
		fi
	done
	if grep "root=/dev/nfs" /proc/cmdline > /dev/null; then
		PATH_OF_EFI_DTB=/boot/efi/dtb
		if [ ! -d "$PATH_OF_EFI_DTB" ]; then
			mkdir "$PATH_OF_EFI_DTB"
		fi
		return
	fi
	echo "$FUNCNAME: no EFI boot path found." >&2
	return 1
}

LDTO_apply()
{
	LDTO_findEFIDTBPath RW
	local base_dtb="$PATH_OF_DT"
	local target_dtb="$PATH_OF_EFI_DTB/$DT_OVERRIDE"
	local target_dtb_path="${target_dtb%/*}"
	if [ ! -d "$target_dtb_path" ]; then
		mkdir "$target_dtb_path"
	fi
	if [ -f "$target_dtb" ]; then
		local base_dtb="$target_dtb"
	fi
	fdtoverlay -i "$base_dtb" -o "$target_dtb" "$1"
}

LDTO_merge()
{
	LDTO_checkFirmwareDT
	LDTO_importDTConfig
	if [ -z "$1" ]; then
		LDTO_list
		return 1
	fi
	LDTO_loadMap
	for dto in "$@"; do
		local dto_real=$(LDTO_getAliasDT "$dto")
		if [ ! -z "$dto_real" ]; then
			local dto=$dto_real
		fi
		local dtbo_file=$(LDTO_getDTBO "$dto")
		if [ -z "$dtbo_file" ]; then
			echo "$FUNCNAME: $dto does not exist and cannot be added." >&2
			return 1
		fi
		LDTO_apply "$dtbo_file"
		echo "Overlay $dto: merged for next boot" >&2
	done
}

LDTO_checkOverrideDT()
{
	LDTO_importDTConfig
	LDTO_findEFIDTBPath $1
	PATH_OF_DTB_OVERRIDE="$PATH_OF_EFI_DTB/$DT_OVERRIDE"
	if [ ! -f "$PATH_OF_DTB_OVERRIDE" ]; then
		echo "$FUNCNAME: no merged overlays detected on system." >&2
		return 1
	fi
}

LDTO_diff()
{
	LDTO_checkOverrideDT
	LDTO_checkFirmwareDT
	diff -u --color --suppress-common-lines --label=CURRENT <(dtc -I dtb -O dts "$PATH_OF_DT" 2> /dev/null) --label=NEXT <(dtc -I dtb -O dts "$PATH_OF_DTB_OVERRIDE" 2> /dev/null)
}

LDTO_show()
{
	LDTO_checkOverrideDT
	dtc -I dtb -O dts "$PATH_OF_DTB_OVERRIDE" 2> /dev/null | less
}

LDTO_edit(){
	LDTO_checkOverrideDT
	dts=$(mktemp)
	dtc -I dtb -O dts "$PATH_OF_DTB_OVERRIDE" -o $dts 2> /dev/null
	vim $dts
	dtc -I dts -O dtb $dts -o "$PATH_OF_DTB_OVERRIDE"
	rm $dts
}

LDTO_current()
{
	LDTO_checkFirmwareDT
	dtc -I dtb -O dts "$PATH_OF_DT" 2> /dev/null | less
}

LDTO_reset()
{
	LDTO_checkOverrideDT RW
	rm -f "$PATH_OF_DTB_OVERRIDE"
	echo "Overlay: reset for next boot" >&2
}

LDTO_help(){
	echo "$0 list"
	echo "$0 status"
	echo "$0 active [DTBO]"
	echo "$0 enable [DTBO]"
	echo "$0 disable [DTBO]"
	echo "$0 current"
	echo "$0 merge [DTBO]"
	echo "$0 show"
	echo "$0 diff"
	echo "$0 reset"
}

cmd=help
if [ ! -z "$1" ]; then
	cmd=$1
	shift
fi
LDTO_${cmd,,} $@