#!/bin/bash

set -o nounset
set -o errexit

REPO_URI=${REPO_URI:-"https://github.com/Hexxeh/rpi-firmware"}

UPDATE_SELF=${UPDATE_SELF:-1}
UPDATE_URI="https://raw.githubusercontent.com/Hexxeh/rpi-update/master/rpi-update"

if [[ ${BOOT_PATH:-"unset"} == "unset" && ${ROOT_PATH:-"unset"} != "unset" ]] ||
[[ ${BOOT_PATH:-"unset"} != "unset" && ${ROOT_PATH:-"unset"} == "unset" ]]; then
	echo " *** You need to specify both ROOT_PATH and BOOT_PATH, or neither"
	exit 1
fi

BRANCH=${BRANCH:-"master"}
ROOT_PATH=${ROOT_PATH:-"/"}
BOOT_PATH=${BOOT_PATH:-"/boot"}
WORK_PATH=${WORK_PATH:-"${ROOT_PATH}/root"}
SKIP_KERNEL=${SKIP_KERNEL:-0}
SKIP_SDK=${SKIP_SDK:-0}
SKIP_REPODELETE=${SKIP_REPODELETE:-0}
SKIP_BACKUP=${SKIP_BACKUP:-0}
SKIP_DOWNLOAD=${SKIP_DOWNLOAD:-0}
WANT_SYMVERS=${WANT_SYMVERS:-0}
RPI_UPDATE_UNSUPPORTED=${RPI_UPDATE_UNSUPPORTED:-0}
FW_REPO="${REPO_URI}.git"
FW_REPOLOCAL=${FW_REPOLOCAL:-"${WORK_PATH}/.rpi-firmware"}
FW_PATH="${BOOT_PATH}"
FW_MODPATH="${ROOT_PATH}/lib/modules"
FW_REV=${1:-""}
FW_REVFILE="${FW_PATH}/.firmware_revision"
[ "${RPI_UPDATE_UNSUPPORTED}" -ne 0 ] && echo -e "You appear to be trying to update firmware on an incompatible distribution. To force update, run the following:\nsudo -E RPI_UPDATE_UNSUPPORTED=0 rpi-update" && exit 1

function update_self() {
	echo " *** Performing self-update"
	_tempFileName="$0.tmp"

	if ! curl -L --output "${_tempFileName}" "${UPDATE_URI}"; then
		echo " !!! Failed to download update for rpi-update!"
		echo " !!! Make sure you have ca-certificates installed and that the time is set correctly"
		exit 1
	fi

	OCTAL_MODE=$(stat -c '%a' "$0")
	if ! chmod ${OCTAL_MODE} "${_tempFileName}" ; then
		echo " !!! Failed: Error while trying to set mode on ${_tempFileName}"
		exit 1
	fi

	cat > "${WORK_PATH}/.updateScript.sh" << EOF
	if mv "${_tempFileName}" "$0"; then
		rm -- "\$0"
		exec env UPDATE_SELF=0 /bin/bash "$0" "${FW_REV}"
	else
		echo " !!! Failed!"
	fi
EOF

	echo " *** Relaunching after update"
	exec /bin/bash "${WORK_PATH}/.updateScript.sh"
}

function update_modules {
	if [[ ${SKIP_KERNEL} -eq 0 ]]; then
		echo " *** Updating kernel modules"
		cp -R "${FW_REPOLOCAL}/modules/"* "${FW_MODPATH}/"
		find "${FW_REPOLOCAL}/modules" -mindepth 1 -maxdepth 1 -type d | while read DIR; do
			BASEDIR=$(basename "${DIR}")
			echo " *** depmod ${BASEDIR}"
			depmod -b "${ROOT_PATH}" -a "${BASEDIR}"
		done
	else
		echo " *** As requested, not updating kernel modules"
	fi
}

function update_vc_libs {
	echo " *** Updating VideoCore libraries"

	ELFOUTPUT=$(readelf -a "${ROOT_PATH}/bin/bash")
	if [ "${ELFOUTPUT}" != "${ELFOUTPUT/VFP_args/}" ]; then
		echo " *** Using HardFP libraries"
		cp -R "${FW_REPOLOCAL}/vc/hardfp/"* "${ROOT_PATH}/"
	else
		echo " *** Using SoftFP libraries"
		cp -R "${FW_REPOLOCAL}/vc/softfp/"* "${ROOT_PATH}/"
	fi
}

function update_sdk {
	if [[ ${SKIP_SDK} -eq 0 ]]; then
		echo " *** Updating SDK"
		cp -R "${FW_REPOLOCAL}/vc/sdk/"* "${ROOT_PATH}/"
	else
		echo " *** As requested, not updating SDK"
	fi
}

function update_firmware {
	echo " *** Updating firmware"
	rm -rf "${FW_PATH}/"*.elf
	rm -rf "${FW_PATH}/"*.bin
	cp "${FW_REPOLOCAL}/"*.elf "${FW_PATH}/"
	cp "${FW_REPOLOCAL}/"*.bin "${FW_PATH}/"
	cp "${FW_REPOLOCAL}/"*.dat "${FW_PATH}/"
	if [[ ${SKIP_KERNEL} -eq 0 ]]; then
		cp "${FW_REPOLOCAL}/"*.img "${FW_PATH}/"
		if [[ -n $(shopt -s nullglob; echo "${FW_REPOLOCAL}/"*.dtb) ]]; then
			cp "${FW_REPOLOCAL}/"*.dtb "${FW_PATH}/"
		fi
	else
		echo " *** As requested, not updating kernel"
	fi
	if [[ ${WANT_SYMVERS} -ne 0 ]]; then
		if [[ -f "${FW_REPOLOCAL}/"Module.symvers ]]; then
			cp "${FW_REPOLOCAL}/"Module.symvers "${FW_PATH}/"
		fi
		if [[ -f "${FW_REPOLOCAL}/"git_hash ]]; then
			cp "${FW_REPOLOCAL}/"git_hash "${FW_PATH}/"
		fi
	fi
}

function finalise {
	if [[ -f "${FW_PATH}/arm192_start.elf" ]]; then
		echo " *** Setting 192M ARM split"
		cp "${FW_PATH}/arm192_start.elf" "${FW_PATH}/start.elf"
	fi
	echo " *** Running ldconfig"
	ldconfig -r "${ROOT_PATH}"
	echo " *** Storing current firmware revision"
	echo "${FW_REV}" > "${FW_REVFILE}"
}

function do_backup {
	if [[ ${SKIP_BACKUP} -eq 0 ]]; then
		echo " *** Backing up files (this will take a few minutes)"
		if [[ -f "${FW_PATH}.bak" ]]; then
			echo " *** Remove old firmware backup"
			rm -rf "${FW_PATH}.bak"
		fi
		echo " *** Backing up firmware"
		cp -a "${FW_PATH}" "${FW_PATH}.bak"
		if [[ ${SKIP_KERNEL} -eq 0 ]]; then
			if [[ -f "${FW_MODPATH}.bak" ]]; then
				echo " *** Remove old modules backup"
				rm -rf "${FW_MODPATH}.bak"
			fi
			echo " *** Backing up modules $(uname -r)"
			mkdir -p "${FW_MODPATH}.bak" && cp -a "${FW_MODPATH}/$(uname -r)" "${FW_MODPATH}.bak"
		fi
	fi
}

function do_update {
	download_rev
	if [[ -f "${FW_REPOLOCAL}/pre-install" ]]; then
		echo " *** Running pre-install script"
		source "${FW_REPOLOCAL}/pre-install"
	fi
	update_firmware
	update_modules
	update_vc_libs
	update_sdk
	finalise
	if [[ -f "${FW_REPOLOCAL}/post-install" ]]; then
		echo " *** Running post-install script"
		source "${FW_REPOLOCAL}/post-install"
	fi
	remove_rev
	echo " *** Syncing changes to disk"
	sync
	echo " *** If no errors appeared, your firmware was successfully updated to ${FW_REV}"
	if [[ "${ROOT_PATH}" == "/" ]]; then
		echo " *** A reboot is needed to activate the new firmware"
	fi
}

function download_rev {
	if [[ ${SKIP_DOWNLOAD} -eq 0 ]]; then
		echo " *** Downloading specific firmware revision (this will take a few minutes)"
		rm -rf "${FW_REPOLOCAL}"
		mkdir -p "${FW_REPOLOCAL}"
		curl -L "${REPO_URI}/tarball/${FW_REV}" | tar xzf - -C "${FW_REPOLOCAL}" --strip-components=1
	fi
}

function remove_rev {
	echo " *** Deleting downloaded files"
	if [[ ${SKIP_REPODELETE} -eq 0 ]]; then
		rm -rf "${FW_REPOLOCAL}"
	fi
}

if [[ ${EUID} -ne 0 ]]; then
	echo " !!! This tool must be run as root"
	exit 1
fi

echo " *** Raspberry Pi firmware updater by Hexxeh, enhanced by AndrewS and Dom"

if [[ ${UPDATE_SELF} -ne 0 ]]; then
	update_self
fi

if [[ ! -d "${FW_PATH}" ]]; then
	echo " !!! ${FW_PATH} doesn't exist"
	exit 1
fi

if [[ ${SKIP_KERNEL} -eq 0 ]] && [[ ! -d "${FW_MODPATH}" ]]; then
	echo " !!! ${FW_MODPATH} doesn't exist"
	exit 1
fi

command -v readelf >/dev/null 2>&1 || {
	echo " !!! This tool requires you have readelf installed, please install it first"
	echo "     In Debian, try: sudo apt-get install binutils"
	echo "     In Arch, try: pacman -S binutils"
	exit 1
}

# ask github for latest version hash
REPO_API=${REPO_URI/github.com/api.github.com\/repos}/git/refs/heads/${BRANCH}
GITREV=$(curl -s ${REPO_API} | awk '{ if ($1 == "\"sha\":") { print substr($2, 2, 40) } }')
FW_REV=${FW_REV:-${GITREV}}

if [[ "${FW_REV}" == "" ]]; then
	echo " *** No hash received from github: ${REPO_API}"
  # run again with errors not suppressed
  curl ${REPO_API}
  exit 1
fi

if [[ ! -f "${FW_REVFILE}" ]]; then
	echo " *** We're running for the first time"
	do_backup
else
	if [[ $(cat "${FW_REVFILE}") == "$FW_REV" ]]; then
		echo " *** Your firmware is already up to date"
		exit 0
	fi
fi

do_update
