#!/usr/bin/env bash

FLAG_OVERWRITE=1
FLAG_REGISTER_KEY=2
FLAG_TPMTOOL_SUPPORTS_SRK_WELL_KNOWN=4
FLAG_SRK_WELL_KNOWN=8
FLAG_TPM2=16

TSS_TCSD_HOSTNAME_DEFAULT=localhost
TSS_TCSD_PORT_DEFAULT=30003

logit()
{
	if [ -z "$LOGFILE" ]; then
		echo "$@" >&1
	else
		echo "$@" >> "$LOGFILE"
	fi
}

logerr()
{
	if [ -z "$LOGFILE" ]; then
		echo "Error: $*" >&2
	else
		echo "Error: $*" >> "$LOGFILE"
	fi
}

# Get the size of a file in bytes
#
# @1: filename
function get_filesize()
{
	if [[ "$(uname -s)" =~ (Linux|CYGWIN_NT-) ]]; then
		stat -c%s "$1"
	else
		# OpenBSD
		stat -f%z "$1"
	fi
}

# Create a config value by escaping the proper characters
#
# @param 1: The string to escape
function escape_pkcs11_url()
{
	echo "$1" | sed 's/;/\\;/g'
}

# Use expect for automating the interaction with the tpmtool
#
# @param 1...: parameters to pass to tpmtool command line
#
# TPM_SRK_PASSWORD and TPM_KEY_PASSWORD global variables are used
# for the SRK and key passwords respectively.
run_tpmtool() {
	local prg out rc

	prg="spawn tpmtool "$@"
		expect {
			\"Enter SRK password:\" {
				send \"${TPM_SRK_PASSWORD}\n\"
				exp_continue
			}
			\"Enter key password:\" {
				send \"${TPM_KEY_PASSWORD}\n\"
				exp_continue
			}
			\"tpmkey:\" {
				send_user \"\n\"
			}
			eof {
				exit
			}
		}
		catch wait result
		exit [lindex \$result 3]
	"
	out=$(expect -c "${prg}")
	rc=$?
	echo "${out}"
	return $rc
} #run_tpmtool

create_localca_cert() {
	local flags=$1
	local dir="$2"
	local outfile="$3"
	local owner="$4"
	local pid="$5" # TPM2 parameter

	local cakey=${dir}/swtpm-localca-rootca-privkey.pem
	local cacert=${dir}/swtpm-localca-rootca-cert.pem
	local tpmkey=${dir}/swtpm-localca-tpmca-privkey.pem
	local tpmpubkey=${dir}/swtpm-localca-tpmca-pubkey.pem
	local tpmca=${dir}/swtpm-localca-tpmca-cert.pem
	local template=${dir}/template
	local tpmkeyurl
	local msg output

	if ! [ -r "${cakey}" ] || ! [ -r "${cacert}" ]; then
		msg=$("${CERTTOOL}" \
			--generate-privkey \
			${SWTPM_ROOTCA_PASSWORD:+--password "${SWTPM_ROOTCA_PASSWORD}"} \
			--outfile "${cakey}" \
			2>&1)
		[ $? -ne 0 ] && {
			logerr "Could not create root-CA key ${cakey}."
			logerr "${msg}"
			return 1
		}
		chmod 640 "${cakey}"

		echo "cn=swtpm-localca-rootca" > "${template}"
		echo "ca" >> "${template}"
		echo "cert_signing_key" >> "${template}"
		echo "expiration_days = 3650" >> "${template}"

		msg=$(GNUTLS_PIN="${SWTPM_ROOTCA_PASSWORD}" ${CERTTOOL} \
			--generate-self-signed \
			--template "${template}" \
			--outfile "${cacert}" \
			--load-privkey "${cakey}" \
			2>&1)

		if [ $? -ne 0 ]; then
			logerr "Could not create root CA."
			logerr "${msg}"
			rm -f "${cakey}" "${template}"
			return 1
		fi
	else
		logit "Reusing existing root CA"
	fi

	rm -f "${tpmkey}" "${tpmpubkey}" "${tpmca}"

	if [ $((flags & FLAG_TPM2)) -ne 0 ]; then
		local tokenurl tpmkeyurl
		local token="swtpm-tpmca-${pid}"
		local label="${token}" # must be same
		local keylabel="swtpm-tpmca-key"
		local userpin="${SWTPM_PKCS11_PIN:-swtpm-tpmca}"

		tokenurl=$(p11tool --list-tokens 2>&1 | \
				grep -E ";token=${token}\$" | \
				sed -n "s/.*URL: //p")
		if [ -z "${tokenurl}" ]; then
			if [ -z "${SWTPM_PKCS11_SO_PIN}" ]; then
				logerr "The env. variable SWTPM_PKCS11_SO_PIN must be set to create token ${label}."
				return 1
			fi
			msg=$(tpm2_ptool addtoken \
				--pid "${pid}" \
				--sopin "${SWTPM_PKCS11_SO_PIN}" \
				--userpin "${userpin}" \
				--label "${label}" 2>&1)
			if [ $? -ne 0 ]; then
				logerr "Error: Could not create pkcs11 token"
				logerr "${msg}"
				return 1
			fi
			tokenurl=$(p11tool --list-tokens 2>&1 | \
					grep -E ";token=${token}\$" | \
					sed -n "s/.*URL: //p")
			if [ -z "${tokenurl}" ]; then
				logerr "Error: Could not get token URL for token '${token}'"
				logerr "${msg}"
				return 1
			fi
			msg=$(tpm2_ptool config \
				--key tcti \
				--value tabrmd \
				--label "${label}")
			if [ $? -ne 0 ]; then
				logerr "Error: Could not set config value for tcti key"
				logerr "${msg}"
				return 1
			fi
		fi

		export GNUTLS_PIN="${userpin}"
		# GNUTLS_SO_PIN not needed at this point

		msg="$(p11tool --login --list-keys "${tokenurl}" 2>&1)"
		if [ $? -eq 0 ]; then
			tpmkeyurl=$(echo "${msg}" | \
				grep ";object=${keylabel}" | \
				sed -n "s/.*URL: //p")
		fi
		if [ -z "${tpmkeyurl}" ]; then
			msg=$(tpm2_ptool addkey \
				"--label=${label}" \
				"--userpin=${userpin}" \
				--algorithm=rsa2048 \
				"--key-label=${keylabel}" \
				--id 1 2>&1)
			if [ $? -ne 0 ]; then
				logerr "Error: Could not create create key under pkcs11 token ${token}"
				logerr "${msg}"
				return 1
			fi
			msg="$(p11tool --login --list-keys "${tokenurl}" 2>&1)"
			if [ $? -ne 0 ]; then
				logerr "Error: Could not get TPM key URL for ${tokenurl}"
				logerr "${msg}"
				return 1
			fi
			tpmkeyurl=$(echo "${msg}" | \
					grep ";object=${keylabel}" | \
					sed -n "s/.*URL: //p")
			if [ -z "${tpmkeyurl}" ]; then
				logerr "Error: Could not get TPM key URL for ${tokenurl}"
				logerr "${msg}"
				return 1
			fi
		fi
		rm -f "${tpmpubkey}"

		msg=$(p11tool --export-pubkey "${tpmkeyurl}" --login --outfile "${tpmpubkey}" 2>&1)
		if [ $? -ne 0 ] || \
		   [ ! -r "${tpmpubkey}" ] || [ $(get_filesize "${tpmpubkey}") -eq 0 ]; then
			logerr "Error: Could not get TPM public key"
			logerr "${msg}"
			rm -f "${tpmkey}" "${tpmpubkey}"
			return 1
		fi
	else
		local params=""

		if [ $((flags & FLAG_SRK_WELL_KNOWN)) -ne 0 ]; then
			unset GNUTLS_PIN
			params="--srk-well-known"
		else
			export GNUTLS_PIN="${TPM_SRK_PASSWORD}"
		fi

		if [ $((flags & FLAG_REGISTER_KEY)) -ne 0 ]; then
			msg="$(run_tpmtool --generate-rsa --signing --register ${params})"
			if [ $? -ne 0 ]; then
				logerr "Could not generate registered signing key with tpmtool"
				logerr "${msg}"
				return 1
			fi
			tpmkeyurl=$(echo "${msg}" | sed -n 's/\(tpmkey:uuid=[^;]*\);.*/\1/p')
			if [ -z "${tpmkeyurl}" ]; then
				logerr "Could not parse tpmkey URL"
				logerr "${msg}"
				return 1
			fi
		else
			rm -f "${tpmkey}"
			msg="$(run_tpmtool --generate-rsa --signing --outfile \"${tpmkey}\" ${params})"
			if [ $? -ne 0 ]; then
				logerr "Could not create signing key with tpmtool"
				logerr "${msg}"
				rm -f "${tpmkey}"
				return 1
			fi
			if [ ! -r "${tpmkey}" ] || [ $(get_filesize "${tpmkey}") -eq 0 ]; then
				logerr "The TPM key file ${tpmkey} was not written properly"
				logerr "${msg}"
				rm -f "${tpmkey}"
				return 1
			fi
			chmod 640 "${tpmkey}"
			tpmkeyurl="tpmkey:file=${tpmkey}"
		fi

		rm -f "${tpmpubkey}"
		msg=$(run_tpmtool "--pubkey=${tpmkeyurl}" --outfile \"${tpmpubkey}\" ${params})
		if [ $? -ne 0 ] || \
		   [ ! -r "${tpmpubkey}" ] || [ $(get_filesize "${tpmpubkey}") -eq 0 ]; then
			logerr "Error: Could not get TPM public key"
			logerr "${msg}"
			rm -f "${tpmkey}" "${tpmpubkey}"
			return 1
		fi
	fi

	echo "cn=swtpm-localca" > "${template}"
	echo "ca" >> "${template}"
	echo "cert_signing_key" >> "${template}"
	echo "expiration_days = 3650" >> "${template}"

	msg=$(${CERTTOOL} \
		--generate-certificate \
		--template "${template}" \
		--outfile "${tpmca}" \
		--load-ca-privkey "${cakey}" \
		--load-ca-certificate "${cacert}" \
		--load-privkey "${tpmkeyurl}" \
		--load-pubkey "${tpmpubkey}" \
		2>&1)

	if [ $? -ne 0 ]; then
		logerr "Could not create TPM CA"
		logerr "${msg}"
		rm -f "${template}"
		return 1
	fi

	output="statedir = ${dir}
signingkey = $(escape_pkcs11_url ${tpmkeyurl})
issuercert = ${tpmca}
certserial = ${dir}/certserial"

	if [ $((flags & FLAG_TPM2)) -eq 0 ]; then
		output+="$(echo -e "\nTSS_TCSD_HOSTNAME = ${TSS_TCSD_HOSTNAME}")"
		output+="$(echo -e "\nTSS_TCSD_PORT = ${TSS_TCSD_PORT}")"
	else
		output+="$(echo -e "\nSWTPM_PKCS11_PIN = ${SWTPM_PKCS11_PIN}")"
		# output+="$(echo -e "\nSWTPM_PKCS11_SO_PIN = ${SWTPM_PKCS11_SO_PIN}")"
	fi
	if [ -n "${TPM_KEY_PASSWORD}" ]; then
		output+="$(echo -e "\nsigningkey_password = ${TPM_KEY_PASSWORD}")"
	fi
	if [ -n "${TPM_SRK_PASSWORD}" ]; then
		output+="$(echo -e "\nparentkey_password = ${TPM_SRK_PASSWORD}")"
	fi

	if [ -n "${outfile}" ]; then
		echo "${output}" > "${outfile}"
		chmod 640 "${outfile}"
	fi
	echo "${output}"

	if [ "$(id -u)" -eq 0 ]; then
		chown "${owner}:${group}" "${dir}"

		pushd "${dir}" &>/dev/null
		if [ $? -eq 0 ]; then
			chown "${owner}:${group}" ./*
			popd &>/dev/null
		fi

		if [ -n "${outfile}" ]; then
			chown "${owner}:${group}" "${outfile}"
		fi
	fi

	rm -f "${template}"

	return 0
} #create_localca_cert

usage() {
	local flags=$2

	local tpmtool_note=" use 'well known' password if not
                   given"
	[ $((flags & FLAG_TPMTOOL_SUPPORTS_SRK_WELL_KNOWN)) -eq 0 ] && \
		tpmtool_note="
                   Note: the well known password of 20 zero bytes is not
                         supported by tpmtool"

	cat << _EOF_
Create a TPM-based CA for signing EK and platform certificates.

Usage: $(basename "$1") [options]

THIS SCRIPT IS EXPERIMENTAL

The following options are supported:

--dir directory    Directory where to write the CA files into; must not exist
                   unless --overwrite is passed
--overwrite        Overwrite any data in an existing directory; tries to
                   reuse a root CA if one is found there
--register         Create a registered TPM 1.2 key rather than a file that
                   contains the key; this option has no effect if --tpm2 is
                   used
--key-password s   Password for the newly created TPM key; required if
                   --register is not passed
                   Note: use the same as the --srk-password (bug in certtool)
--srk-password s   Password for the TPM's SRK;${tpmtool_note}
--outfile file     File to write the configuration to; if not passed it will be
                   written to stdout only
--owner owner      The owner of the directory and the files; only set if this
                   script is run as root; recommended to be 'tss'
--group group      The group owning the directory and the files;
                   recommended to be 'tss'
--tss-tcsd-hostname hostname
                   The name of the host where tcsd (TrouSerS daemon) is running
                   on; default is '${TSS_TCSD_HOSTNAME_DEFAULT}'
--tss-tcsd-port p  The TCP port on which tcsd is listening for connections;
                   default is ${TSS_TCSD_PORT_DEFAULT}
--tpm2             Setup a CA that uses a TPM 2.0
--pid <pid>        Pimary object Id used by tpm2_ptool; only valid if --tpm2
                   is used
--help, -h, -?     Display this help screen and exit


The following environment variables are supported:

SWTPM_ROOTCA_PASSWORD  The root CA's private key password

_EOF_
} #usage

# Check whether tpmtool supports --srk-well-known
tpmtool_supports_srk_well_known()
{
	local tmp

	tmp=$(tpmtool --help | grep "srk-well-known")
	[ -z "${tmp}" ] && return 1
	return 0
}

main() {
	local flags=0
	local dir outfile owner group msg pid

	if tpmtool_supports_srk_well_known; then
		flags=$((flags | FLAG_TPMTOOL_SUPPORTS_SRK_WELL_KNOWN | FLAG_SRK_WELL_KNOWN))
	fi

	CERTTOOL=certtool
	export TSS_TCSD_HOSTNAME=${TSS_TCSD_HOSTNAME_DEFAULT}
	export TSS_TCSD_PORT=${TSS_TCSD_PORT_DEFAULT}

	while [ $# -ne 0 ]; do
		case "$1" in
		--dir)
			shift
			dir="$1"
			;;
		--overwrite)
			flags=$((flags | FLAG_OVERWRITE))
			;;
		--register)
			flags=$((flags | FLAG_REGISTER_KEY))
			;;
		--srk-password)
			shift
			TPM_SRK_PASSWORD="$1"
			flags=$((flags & ~FLAG_SRK_WELL_KNOWN))
			;;
		--key-password)
			shift
			TPM_KEY_PASSWORD="$1"
			;;
		--outfile)
			shift
			outfile="$1"
			;;
		--owner)
			shift
			owner="$1"
			;;
		--group)
			shift
			group="$1"
			;;
		--tss-tcsd-hostname)
			shift
			TSS_TCSD_HOSTNAME="$1"
			;;
		--tss-tcsd-port)
			shift
			TSS_TCSD_PORT="$1"
			;;
		--tpm2)
			flags=$((flags | FLAG_TPM2))
			;;
		--pid)
			shift
			pid="$1"
			;;
		--help|-h|-?)
			usage "$0" "${flags}"
			exit 0
			;;
		*)
			logerr "Unsupported option $1"
			exit 1
			;;
		esac
		shift
	done
	if [ -z "${dir}" ]; then
		logerr "Missing --dir option."
		return 1
	fi
	# strip trailing '/' from dir
	dir="$(echo "${dir}" | sed -n 's|[/]*$||p')"

	if [ -d "${dir}" ] && [ $((flags & FLAG_OVERWRITE)) -eq 0 ]; then
		logerr "Refusing to overwrite existing directory ${dir}."
		return 1
	fi

	if [ -z "${TPM_SRK_PASSWORD}" ] && [ $((flags & FLAG_TPM2)) -eq 0 ] &&
	   [ $((flags & FLAG_TPMTOOL_SUPPORTS_SRK_WELL_KNOWN)) -eq 0 ]; then
		logerr "SRK password must be provided"
		return 1
	fi

	if [ -z "${TPM_KEY_PASSWORD}" ] && \
	   [ $((flags & FLAG_REGISTER_KEY)) -eq 0 ] && \
	   [ $((flags & FLAG_TPM2)) -eq 0 ]; then
		logerr "Key password is required"
		return 1
	fi

	if [ $((flags & FLAG_TPM2)) -ne 0 ] && [ -z "${pid}" ]; then
		logerr "--pid is required for TPM 2"
		return 1
	fi

	if [ "$(id -u)" -eq 0 ]; then
		if [ -n "${owner}" ]; then
			msg="$(id -u "${owner}" 2>&1)"
			if [ $? -ne 0 ]; then
				logerr "User ${owner} cannot be used: ${msg}"
				return 1
			fi
		else
			owner="root"
		fi
		if [ -n "${group}" ]; then
			msg="$(id -g "${group}" 2>&1)"
			if [ $? -ne 0 ]; then
				logerr "Group ${group} cannot be used: ${msg}"
				return 1
			fi
		else
			group="root"
		fi
	fi

	mkdir -p "${dir}"
	if [ $? -ne 0 ]; then
		logerr "Could not create directory ${dir}."
		return 1
	fi

	create_localca_cert "${flags}" "${dir}" "${outfile}" "${owner}" "${pid}"
	return $?
} #main

main "$@"
exit $?