#!/usr/bin/env bash
#
# zrec-age - Encrypted-recfile wrapper around zrec.
#
# Decrypts an age-encrypted recfile to a temporary plaintext file, runs zrec
# against it, then re-encrypts the result back to the original location.
#
# Usage:
#   zrec-age [options] <encrypted-recfile> [record-type]
#
# Options:
#   -i <identity-file>     age identity (private key) file for DECRYPTION.
#                          May be repeated. (Implies key-based mode.)
#   -r <recipient>         age recipient (public key) for ENCRYPTION.
#                          May be repeated. (Implies key-based mode.)
#   -R <recipients-file>   File containing one recipient per line.
#   -p                     Use passphrase (symmetric) mode for both
#                          decrypt and encrypt. (Default if no -i/-r/-R given.)
#   -z <path-to-zrec>      Path to the zrec script (default: looks on PATH,
#                          then alongside this script).
#   -h                     Show this help.
#
# Environment:
#   ZREC_BIN               Overrides the zrec location (same as -z).
#
# Notes:
#   * In passphrase mode you will be prompted up to three times by age
#     (decrypt, then encrypt). This is unavoidable with symmetric age.
#   * Plaintext only ever lives in a 0700 temp directory and is removed
#     (best-effort shredded) on exit.
#
set -o pipefail

# ---------------------------------------------------------------------------
# Globals
# ---------------------------------------------------------------------------
PROG="$(basename "$0")"
WORKDIR=""
PLAINTEXT=""
PASSPHRASE_MODE="auto" # auto | yes | no
declare -a IDENTITIES=()
declare -a RECIPIENTS=()
declare -a RECIPIENT_FILES=()
ZREC_BIN="${ZREC_BIN:-}"

# ---------------------------------------------------------------------------
# Logging helpers (all to stderr; stdout stays clean)
# ---------------------------------------------------------------------------
log() { printf '%s\n' "$*" >&2; }
err() { printf 'Error: %s\n' "$PROG" "$*" >&2; }
die() {
	err "$*"
	exit 1
}

usage() {
	sed -n '3,40p' "$0" | sed 's/^# \{0,1\}//' >&2
	exit "${1:-1}"
}

# ---------------------------------------------------------------------------
# Dependency checks
# ---------------------------------------------------------------------------
check_deps() {
	command -v age >/dev/null 2>&1 || die "'age' not found. Install it (brew install age / apt install age)."

	if [ -z "$ZREC_BIN" ]; then
		if command -v zrec >/dev/null 2>&1; then
			ZREC_BIN="$(command -v zrec)"
		else
			# Fall back to a zrec sitting next to this wrapper.
			local self_dir
			self_dir="$(cd "$(dirname "$0")" && pwd)"
			if [ -x "$self_dir/zrec" ]; then
				ZREC_BIN="$self_dir/zrec"
			fi
		fi
	fi
	[ -n "$ZREC_BIN" ] && [ -x "$ZREC_BIN" ] ||
		die "Could not find an executable 'zrec' (use -z or set ZREC_BIN)."
}

# ---------------------------------------------------------------------------
# Cross-platform secure temp dir creation.
# ---------------------------------------------------------------------------
make_workdir() {
	# mktemp -d differs slightly across platforms but -d <template> is portable.
	WORKDIR="$(mktemp -d "${TMPDIR:-/tmp}/zrec-age.XXXXXX")" ||
		die "Failed to create temporary directory."
	chmod 700 "$WORKDIR"
}

# ---------------------------------------------------------------------------
# Best-effort secure delete (cross-platform).
# ---------------------------------------------------------------------------
secure_rm() {
	local f="$1"
	[ -f "$f" ] || return 0
	if command -v shred >/dev/null 2>&1; then
		shred -u "$f" 2>/dev/null && return 0
	fi
	# macOS has no shred; overwrite once then unlink as a fallback.
	if command -v dd >/dev/null 2>&1; then
		local size
		size="$(wc -c <"$f" 2>/dev/null | tr -d ' ')"
		if [ -n "$size" ] && [ "$size" -gt 0 ]; then
			dd if=/dev/urandom of="$f" bs="$size" count=1 conv=notrunc 2>/dev/null || true
		fi
	fi
	rm -f "$f"
}

# ---------------------------------------------------------------------------
# Cleanup trap: shred plaintext, remove workdir. Runs on ANY exit.
# ---------------------------------------------------------------------------
cleanup() {
	[ -n "$PLAINTEXT" ] && secure_rm "$PLAINTEXT"
	[ -n "$WORKDIR" ] && [ -d "$WORKDIR" ] && rm -rf "$WORKDIR"
}
trap cleanup EXIT INT TERM HUP

# ---------------------------------------------------------------------------
# Build the age recipient argument array for encryption.
# Echoes nothing; populates a global array via name. We just assemble here.
# ---------------------------------------------------------------------------
build_recipient_args() {
	AGE_ENC_ARGS=()
	local r f
	for r in "${RECIPIENTS[@]}"; do
		AGE_ENC_ARGS+=(-r "$r")
	done
	for f in "${RECIPIENT_FILES[@]}"; do
		[ -f "$f" ] || die "Recipients file not found: $f"
		AGE_ENC_ARGS+=(-R "$f")
	done
}

# ---------------------------------------------------------------------------
# Decrypt encrypted file -> $PLAINTEXT
# ---------------------------------------------------------------------------
decrypt_file() {
	local encfile="$1"

	if [ "$PASSPHRASE_MODE" = "yes" ]; then
		log "Decrypting (passphrase mode)..."
		age --decrypt -o "$PLAINTEXT" "$encfile" ||
			die "Decryption failed (wrong passphrase or corrupt file?)."
	else
		local -a id_args=()
		local id
		for id in "${IDENTITIES[@]}"; do
			[ -f "$id" ] || die "Identity file not found: $id"
			id_args+=(-i "$id")
		done
		[ "${#id_args[@]}" -gt 0 ] || die "Key-based decryption requires at least one -i <identity>."
		log "Decrypting (key mode)..."
		age --decrypt "${id_args[@]}" -o "$PLAINTEXT" "$encfile" ||
			die "Decryption failed (wrong identity or corrupt file?)."
	fi
}

# ---------------------------------------------------------------------------
# Encrypt $PLAINTEXT -> encrypted file (atomic replace via temp + mv).
# ---------------------------------------------------------------------------
encrypt_file() {
	local encfile="$1"
	local tmp_enc="$WORKDIR/out.age"

	if [ "$PASSPHRASE_MODE" = "yes" ]; then
		log "Re-encrypting (passphrase mode)..."
		age --passphrase -o "$tmp_enc" "$PLAINTEXT" ||
			die "Re-encryption failed; original file left UNCHANGED."
	else
		build_recipient_args
		[ "${#AGE_ENC_ARGS[@]}" -gt 0 ] ||
			die "Key-based encryption requires at least one -r/-R recipient."
		log "Re-encrypting (key mode)..."
		age --encrypt "${AGE_ENC_ARGS[@]}" -o "$tmp_enc" "$PLAINTEXT" ||
			die "Re-encryption failed; original file left UNCHANGED."
	fi

	# Preserve original permissions where possible.
	if [ -f "$encfile" ]; then
		# Copy mode from the original encrypted file onto the new ciphertext.
		if command -v stat >/dev/null 2>&1; then
			local mode
			# GNU stat vs BSD stat.
			mode="$(stat -c '%a' "$encfile" 2>/dev/null || stat -f '%Lp' "$encfile" 2>/dev/null)"
			[ -n "$mode" ] && chmod "$mode" "$tmp_enc" 2>/dev/null || true
		fi
	fi

	# Atomic replace.
	mv -f "$tmp_enc" "$encfile" ||
		die "Failed to move new ciphertext into place; original may be intact, check $tmp_enc."
}

# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------
parse_args() {
	while getopts ":i:r:R:pz:h" opt; do
		case "$opt" in
		i)
			IDENTITIES+=("$OPTARG")
			PASSPHRASE_MODE="no"
			;;
		r)
			RECIPIENTS+=("$OPTARG")
			PASSPHRASE_MODE="no"
			;;
		R)
			RECIPIENT_FILES+=("$OPTARG")
			PASSPHRASE_MODE="no"
			;;
		p) PASSPHRASE_MODE="yes" ;;
		z) ZREC_BIN="$OPTARG" ;;
		h) usage 0 ;;
		\?) die "Unknown option: -$OPTARG" ;;
		:) die "Option -$OPTARG requires an argument." ;;
		esac
	done
	shift $((OPTIND - 1))

	[ "$#" -ge 1 ] || usage 1

	ENCFILE="$1"
	shift
	# Remaining args (record-type) get forwarded to zrec verbatim.
	ZREC_EXTRA_ARGS=("$@")

	# Resolve auto mode -> passphrase if no keys supplied.
	[ "$PASSPHRASE_MODE" = "auto" ] && PASSPHRASE_MODE="yes"

	# Sanity: in key mode, ensure we can both decrypt and re-encrypt.
	if [ "$PASSPHRASE_MODE" = "no" ]; then
		[ "${#IDENTITIES[@]}" -gt 0 ] ||
			die "Key mode: provide -i <identity> for decryption."
		if [ "${#RECIPIENTS[@]}" -eq 0 ] && [ "${#RECIPIENT_FILES[@]}" -eq 0 ]; then
			die "Key mode: provide -r/-R recipient(s) for re-encryption."
		fi
	fi
}

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
main() {
	parse_args "$@"
	check_deps

	[ -f "$ENCFILE" ] || die "Encrypted recfile '$ENCFILE' does not exist."
	[ -r "$ENCFILE" ] || die "Encrypted recfile '$ENCFILE' is not readable."

	make_workdir
	PLAINTEXT="$WORKDIR/recfile.rec"

	decrypt_file "$ENCFILE"

	# Snapshot to detect whether zrec actually changed anything.
	local before_sum after_sum
	before_sum="$(cksum <"$PLAINTEXT")"

	log ""
	log "Launching zrec on decrypted copy..."
	log "----------------------------------------"

	# Run zrec interactively against the plaintext temp file.
	if "$ZREC_BIN" "$PLAINTEXT" "${ZREC_EXTRA_ARGS[@]}"; then
		after_sum="$(cksum <"$PLAINTEXT")"
		if [ "$before_sum" = "$after_sum" ]; then
			log "No changes detected; skipping re-encryption."
		else
			encrypt_file "$ENCFILE"
			log "Encrypted recfile updated: $ENCFILE"
		fi
	else
		local rc=$?
		err "zrec exited with status $rc; original encrypted file left UNCHANGED."
		exit "$rc"
	fi
}

main "$@"
