#!/bin/bash

# Copyright (C) 2016 Paul Kocialkowski <contact@paulk.fr>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

REGEXP="\([^[:space:]]*\)[[:space:]]\(.*\)"

KEYBLOCK="keyblock"
VBPRIVK="vbprivk"
VBPUBK="vbpubk"
KEYB="keyb"
PEM="pem"
CRT="crt"

KEYS="ec_root_key ec_data_key root_key firmware_data_key kernel_subkey kernel_data_key recovery_key recovery_kernel_data_key installer_kernel_data_key"
KEYBLOCKS="firmware ec recovery_kernel kernel installer_kernel"
SUBKEYS="firmware_data_key root_key ec_data_key ec_root_key recovery_kernel_data_key recovery_key kernel_data_key kernel_subkey installer_kernel_data_key recovery_key"
ALGORITHMS="7 7 11 7 7 4 11 11 11"
MODES="7 7 11 7 10"

usage() {
	printf "$executable [action] [keys directory] (firmware image path)\n" >&2

	printf "\nActions:\n" >&2
	printf "  generate - Generate a set of keys\n" >&2
	printf "  sign - Sign a firmware image\n" >&2
	printf "  verify - Verify keyblocks\n" >&2

	printf "\nOutput files:\n" >&2
	printf "  sign - Generates a firmware images with the \"-signed\" suffix\n" >&2

	printf "\nEnvironment variables:\n" >&2
	printf "  KEYS_VERSION - Version to give the keys\n" >&2
	printf "  VBOOT_TOOLS_PATH - Path to vboot tools\n" >&2
}

generate() {
	local keys_directory=$1

	local algorithms=$ALGORITHMS
	local subkeys=$SUBKEYS
	local modes=$MODES
	local keyblock
	local algorithm
	local pubkey
	local privkey
	local mode

	for key in $KEYS
	do
		algorithm=$( echo "$algorithms" | sed "s/$REGEXP/\1/g" )
		algorithms=$( echo "$algorithms" | sed "s/$REGEXP/\2/g" )

		key_length=$(( 1 << (10 + ($algorithm / 3)) ))

		openssl genrsa -F4 -out "$keys_directory/$key.$PEM" "$keys_directory/$key_length"
		openssl req -batch -new -x509 -key "$keys_directory/$key.$PEM"
		openssl req -batch -new -x509 -key "$keys_directory/$key.$PEM" -out "$key.$CRT"
		dumpRSAPublicKey -cert "$keys_directory/$key.$CRT" > "$keys_directory/$key.$KEYB"
		futility vbutil_key --pack "$keys_directory/$key.$VBPUBK" --key "$keys_directory/$key.$KEYB" --version "$KEYS_VERSION" --algorithm "$algorithm"
		futility vbutil_key --pack "$keys_directory/$key.$VBPRIVK" --key "$keys_directory/$key.$PEM" --algorithm "$algorithm"

		rm -f "$keys_directory/$key.$PEM" "$keys_directory/$key.$CRT" "$keys_directory/$key.$KEYB"
	done

	for keyblock in $KEYBLOCKS
	do
		pubkey=$( echo "$subkeys" | sed "s/$REGEXP/\1/g" )
		subkeys=$( echo "$subkeys" | sed "s/$REGEXP/\2/g" )
		privkey=$( echo "$subkeys" | sed "s/$REGEXP/\1/g" )
		subkeys=$( echo "$subkeys" | sed "s/$REGEXP/\2/g" )

		mode=$( echo "$modes" | sed "s/$REGEXP/\1/g" )
		modes=$( echo "$modes" | sed "s/$REGEXP/\2/g" )

		futility vbutil_keyblock --pack "$keys_directory/$keyblock.$KEYBLOCK" --flags "$mode" --datapubkey "$keys_directory/$pubkey.$VBPUBK" --signprivate "$keys_directory/$privkey.$VBPRIVK"
		futility vbutil_keyblock --unpack "$keys_directory/$keyblock.$KEYBLOCK" --signpubkey "$keys_directory/$privkey.$VBPUBK"
	done
}

sign() {
	local keys_directory=$1
	local firmware_image_path=$2

	futility sign --signprivate="$keys_directory/firmware_data_key.$VBPRIVK" --keyblock "$keys_directory/firmware.$KEYBLOCK" --kernelkey "$keys_directory/kernel_subkey.$VBPUBK" -v "$KEYS_VERSION" --infile "$firmware_image_path"
	futility gbb_utility -s --recoverykey="$keys_directory/recovery_key.$VBPUBK" --rootkey="$keys_directory/root_key.$VBPUBK" "$firmware_image_path" "$firmware_image_path"
}

verify() {
	local keys_directory=$1

	local subkeys=$SUBKEYS
	local pubkey
	local privkey

	for keyblock in $KEYBLOCKS
	do
		pubkey=$( echo "$subkeys" | sed "s/$REGEXP/\1/g" )
		subkeys=$( echo "$subkeys" | sed "s/$REGEXP/\2/g" )
		privkey=$( echo "$subkeys" | sed "s/$REGEXP/\1/g" )
		subkeys=$( echo "$subkeys" | sed "s/$REGEXP/\2/g" )

		futility vbutil_keyblock --unpack "$keys_directory/$keyblock.$KEYBLOCK" --signpubkey "$keys_directory/$privkey.$VBPUBK"
	done
}

verify_firmware() {
	local keys_directory=$1
	local firmware_image_path=$2

	futility verify -k "$keys_directory/root_key.$VBPUBK" --type bios "$firmware_image_path" || printf "\nBad firmware image signature!\n" >&2 && return 1
}

requirements() {
	local requirement
	local requirement_path

	for requirement in "$@"
	do
		requirement_path=$( which "$requirement" || true )

		if [ -z "$requirement_path" ]
		then
			printf "Missing requirement: $requirement\n" >&2
			exit 1
		fi
	done
}

setup() {
	root=$( realpath "$( dirname "$0" )" )
	executable=$( basename "$0" )

	if [ -z "$KEYS_VERSION" ]
	then
		KEYS_VERSION=1
	fi

	if ! [ -z "$VBOOT_TOOLS_PATH" ]
	then
		PATH="$PATH:$VBOOT_TOOLS_PATH"
	fi
}

cros_boot_keys() {
	local action=$1
	local keys_directory=$2
	local firmware_image_path=$3

	set -e

	setup "$@"

	if [ -z "$action" ] || ! [ -d "$keys_directory" ]
	then
		usage
		exit 1
	fi

	case $action in
		"generate")
			requirements "openssl" "dumpRSAPublicKey" "futility"
			generate "$keys_directory"
			;;
		"sign")
			if ! [ -f "$firmware_image_path" ]
			then
				usage
				exit 1
			fi

			requirements "futility"
			sign "$keys_directory" "$firmware_image_path"
			;;
		"verify")
			requirements "futility"
			verify "$keys_directory"

			if [ -f "$firmware_image_path" ]
			then
				verify_firmware "$keys_directory" "$firmware_image_path"
			fi
			;;
		*)
			usage
			exit 1
			;;
	esac
}

cros_boot_keys "$@"