#!/usr/bin/env bash

# ksigner - Secure Boot kernel signing utility
# Version: @VERSION@

# Source configuration
CONFIG_FILE="/etc/ksigner/ksigner.conf"
if [[ -f "$CONFIG_FILE" ]]; then
	source "$CONFIG_FILE"
fi

# Key lifetime in days (default: 100 years)
KEY_LIFETIME_DAYS=$((365 * 100))

# Directory paths for keys
KEY_PUB_DIR="/etc/pki/sbsign/certs/"
KEY_PRIV_DIR="/etc/pki/sbsign/private/"

# Key filenames
KEY_PUB="MOK.pem"
KEY_PRIV="MOK.priv"
KEY_DER="MOK.der"

KEY_SUBJ="/O=ksigner/OU=KSigner Secure Boot/CN=ksigner"

REQUIRED_BINARIES=(
	"openssl"
	"mokutil"
	"sbsign"
	"sbverify"
	"sha512hmac"
)

panic() {
	echo "ERROR: $1" >&2
	exit 1
}

log() {
	echo "$(date '+%Y-%m-%d %H:%M:%S') - $1" >&2
}

usage() {
	echo "Usage: $0 {setup|sign|sign-all|version|status} [vmlinuz_kernel_filepath]"
	echo "  setup		- Create and install signing keys"
	echo "  sign		- Sign a kernel file (optional: vmlinuz_kernel_filepath)"
	echo "  sign-all	- Sign all available kernels"
	echo "  version	- Show version information"
	echo "  status	- Show signing key status"
	exit 1
}

version() {
	echo "ksigner version @VERSION@"
	echo "Copyright (C) @YEAR@"
	echo "This is free software; see the source for copying conditions."
}

status() {
	echo "Kernel Signer Status:"
	echo "====================="
	echo "Public key file: $KEY_PUB_DIR$KEY_PUB"
	if [[ -f "$KEY_PUB_DIR$KEY_PUB" ]]; then
		echo "  Status: Found"
		echo "  Details:"
		openssl x509 -in "$KEY_PUB_DIR$KEY_PUB" -noout -subject -dates 2>/dev/null || echo "  Error reading certificate"
	else
		echo "  Status: Not found"
	fi

	echo "Private key file: $KEY_PRIV_DIR$KEY_PRIV"
	[[ -f "$KEY_PRIV_DIR$KEY_PRIV" ]] && echo "  Status: Found" || echo "  Status: Not found"

	echo "DER key file: $KEY_PUB_DIR$KEY_DER"
	[[ -f "$KEY_PUB_DIR$KEY_DER" ]] && echo "  Status: Found" || echo "  Status: Not found"

	echo "MOK keys enrolled:"
	mokutil --list-enrolled 2>/dev/null | grep -A 3 -B 1 "Kernel Signing" || echo "  No custom MOK keys found"

	echo "Configuration file: $CONFIG_FILE"
	[[ -f "$CONFIG_FILE" ]] && echo "  Status: Found" || echo "  Status: Using defaults"
}

req_check() {
	local missing_binaries=()
	for binary in "${REQUIRED_BINARIES[@]}"; do
		if ! command -v "$binary" >/dev/null 2>&1; then
			missing_binaries+=("$binary")
		fi
	done

	if [[ ${#missing_binaries[@]} -gt 0 ]]; then
		echo "Missing required binaries: ${missing_binaries[*]}" >&2
		echo "Please install the following packages:" >&2
		for binary in "${missing_binaries[@]}"; do
			case "$binary" in
			"openssl") echo "  - openssl" >&2 ;;
			"mokutil") echo "  - mokutil" >&2 ;;
			"sbsign") echo "  - sbsigntools" >&2 ;;
			"sha512hmac") echo "  - hmaccalc" >&2 ;;
			esac
		done
		exit 1
	fi
	log "All required binaries are present"
}

version_greater() {
	local ver1="$1"
	local ver2="$2"
	[ "$ver1" = "$(printf '%s\n%s' "$ver1" "$ver2" | sort -V | tail -n1)" ]
}

check_kernel_signature() {
	local kernel_file="$1"

	if [[ -z "$kernel_file" ]]; then
		echo "Error: File '$kernel_file' not provided"
		return 1
	fi

	if [[ ! -f "$kernel_file" ]]; then
		echo "Error: File '$kernel_file' not found"
		return 1
	fi

	local output
	output=$(sbverify --list "$kernel_file" 2>&1)

	if [[ -n "$output" ]] && ([[ "$output" == *"signature"* ]] || [[ "$output" == *"issuer"* ]] || [[ "$output" == *"Certificate"* ]]); then
		return 0
	fi

	return 1
}

find_all_kernels() {
	local all_files=()
	for file in /boot/vmlinuz-*; do
		[ -f "$file" ] || continue
		if [[ "$file" == *"rescue"* ]]; then
			continue
		fi
		all_files+=("$file")
	done
	echo "${all_files[@]}"
}

find_latest_kernel() {
	local latest_file=""
	local latest_version=""
	for file in /boot/vmlinuz-*; do
		[ -f "$file" ] || continue
		if [[ "$file" == *"rescue"* ]]; then
			continue
		fi
		local kver="${file#*vmlinuz-}"
		kver="${kver%%-*}"
		if [ -z "$latest_version" ] || version_greater "$kver" "$latest_version"; then
			latest_version="$kver"
			latest_file="$file"
		fi
	done
	echo "$latest_file"
}

setup_signing_keys() {
	# Step 1: Create the signing keys
	log "[Step 1] Creating signing keys..."
	openssl req -new -x509 -newkey rsa:4096 \
		-keyout "$KEY_PRIV" \
		-outform DER -out "$KEY_DER" \
		-nodes -days "$KEY_LIFETIME_DAYS" \
		-subj "$KEY_SUBJ" ||
		panic "[Step 1] Failed to create signing keys"

	# Step 2: Convert the keys to PEM format
	log "[Step 2] Converting keys to PEM format..."
	openssl x509 -inform der -in $KEY_DER -out $KEY_PUB ||
		panic "[Step 2] Failed to convert keys to PEM format"

	# Step 3: Create the directory for the keys
	log "[Step 3] Creating the directory for the keys..."
	mkdir -p $KEY_PUB_DIR $KEY_PRIV_DIR ||
		panic "[Step 3] Failed to create the directory for the keys"

	# Step 4: Copy the keys to the directory
	log "[Step 4] Copying the keys to the directory..."
	mv -f $KEY_DER $KEY_PUB $KEY_PUB_DIR ||
		panic "[Step 4a] Failed to copy the keys to the directory"
	mv -f $KEY_PRIV $KEY_PRIV_DIR ||
		panic "[Step 4b] Failed to copy the keys to the directory"

	# Step 5: Set the permissions for the keys
	log "[Step 5] Setting the permissions for the keys..."
	chmod -R 600 $KEY_PRIV_DIR ||
		panic "[Step 5] Failed to set the permissions for the keys"

	# Step 6: Import the keys to the MOK
	log "[Step 6] Importing the keys to the MOK..."
	mokutil --import $KEY_PUB_DIR$KEY_DER ||
		panic "[Step 6] Failed to import the keys to the MOK"

	echo
	log "Signing keys have been created and installed"
	echo
	echo "Please reboot your computer and enroll the keys with MOK"
	echo
	echo "Enroll MOK -> Continue -> Yes -> Enter password -> OK"
	echo
}

sign_kernel() {
	local kern_version="$1"
	local kern_file="$2"

	# Step 1: Check if the kernel is already signed
	log "[Step 1] Checking if '$kern_version' is already signed..."
	if check_kernel_signature "$kern_file"; then
		log "'$kern_version' is already signed, skipping"
		return
	fi

	# Step 2: Sign the kernel
	log "[Step 2] Signing '$kern_version'..."
	/usr/bin/sbsign \
		--key "$KEY_PRIV_DIR$KEY_PRIV" \
		--cert "$KEY_PUB_DIR$KEY_PUB" \
		"$kern_file" \
		--output "$kern_file.signed" ||
		panic "[Step 1] Failed to sign '$kern_version'"

	# Step 3: Verify the kernel was signed
	log "[Step 3] Verifying '$kern_file' was signed"
	[ -f "$kern_file.signed" ] ||
		panic "'$kern_file.signed' was not found"

	# Step 4: Move the signed kernel
	log "[Step 4] Moving '$kern_file.signed' to '$kern_file'"
	mv -f "$kern_file.signed" "$kern_file" ||
		panic "Failed to move '$kern_file.signed'"

	# Step 5: Make the kernel executable
	log "[Step 5] Setting permissions for '$kern_file'"
	chmod +x "$kern_file" ||
		panic "Failed to make '$kern_file' executable"

	# Step 6: Create the HMAC
	log "[Step 6] Creating HMAC for '$kern_file'"
	sha512hmac "$kern_file" >"${kern_file/vmlinuz/.vmlinuz}.hmac" ||
		panic "Failed to create HMAC for '$kern_file'"

	echo
	log "Signed '$kern_version' successfully"
	echo
}

file_checks() {
	local kfile="$1"
	[ -e "$kfile" ] ||
		panic "Kernel file '$kfile' was not found"
	[ -e "$KEY_PUB_DIR$KEY_PUB" ] ||
		panic "'$KEY_PUB_DIR$KEY_PUB' was not found, please run 'setup' first"
	[ -e "$KEY_PRIV_DIR$KEY_PRIV" ] ||
		panic "'$KEY_PRIV_DIR$KEY_PRIV' was not found, please run 'setup' first"
}

main() {
	case "$1" in
	"setup")
		req_check
		setup_signing_keys
		;;
	"sign")
		req_check
		shift
		[ -n "$1" ] &&
			kfile="$1" ||
			kfile="$(find_latest_kernel)"
		kver="${kfile#*vmlinuz-}"
		kver="${kver%%-*}"
		file_checks "$kfile"
		sign_kernel "$kver" "$kfile"
		;;
	"sign-all")
		req_check
		all_kernels_str=$(find_all_kernels)
		read -ra all_kernels <<<"$all_kernels_str"
		for kfile in "${all_kernels[@]}"; do
			[ -n "$kfile" ] || continue
			kver="${kfile#*vmlinuz-}"
			kver="${kver%%-*}"
			file_checks "$kfile"
			sign_kernel "$kver" "$kfile"
		done
		;;
	"version")
		version
		;;
	"status")
		status
		;;
	*)
		usage
		;;
	esac
}

# Ensure running as root
[[ $EUID -eq 0 ]] || exec sudo "$0" "$@"

main "$@"