#!/bin/bash
#
# air-ssh-keys — privileged helper for cast-server to manage SSH
# authorized_keys for the `air` user WITHOUT pulling the SD card.
#
# cast-server runs with ProtectHome=true and cannot write to
# /home/air directly. This wrapper is the single sudoers-allowlisted
# path. The air-vpn sudoers pattern is deliberately copied: a thin
# wrapper with a hard subcommand allowlist is safer than trying to
# enumerate argv forms in /etc/sudoers.d/.
#
# Subcommands:
#   list            — print authorized_keys, one per line
#   add             — read one pubkey line from stdin, append if unique
#   remove <fpr>    — delete line matching the given SHA256 fingerprint
#
# Input validation: `ssh-keygen -lf -` rejects anything that isn't a
# valid public key (returns non-zero). A pubkey that parses cleanly
# through ssh-keygen cannot smuggle command= / environment= /
# from= options because we strip everything before the algorithm
# token via awk before appending.
set -euo pipefail

HOME_DIR=/home/air
AUTH_DIR="${HOME_DIR}/.ssh"
AUTH_FILE="${AUTH_DIR}/authorized_keys"
MAX_KEYS=20

ensure_layout() {
    install -d -m 0700 -o air -g air "${AUTH_DIR}"
    if [ ! -f "${AUTH_FILE}" ]; then
        install -m 0600 -o air -g air /dev/null "${AUTH_FILE}"
    fi
}

case "${1:-}" in
    list)
        # Listing is read-only. On freshly flashed or read-only-home systems,
        # the file may not exist yet; that means "no uploaded keys", not an
        # error that should fill Recent Activity.
        if [ ! -r "${AUTH_FILE}" ]; then
            exit 0
        fi
        # Keep order; strip comments/blank lines.
        grep -Ev '^\s*(#|$)' "${AUTH_FILE}" || true
        ;;

    add)
        ensure_layout
        raw="$(head -c 8193)"
        if [ "${#raw}" -gt 8192 ]; then
            echo "air-ssh-keys: key line > 8192 bytes" >&2
            exit 2
        fi
        # Accept only the FIRST non-empty line, trimmed.
        line="$(printf '%s' "$raw" | awk 'NF {print; exit}')"
        if [ -z "$line" ]; then
            echo "air-ssh-keys: empty input" >&2
            exit 2
        fi
        # Hard-cap the line length so a pathological paste can't DoS the
        # ssh-keygen validator or bloat authorized_keys. 8 kB is larger
        # than any legitimate pubkey including a long comment.
        if [ "${#line}" -gt 8192 ]; then
            echo "air-ssh-keys: key line > 8192 bytes" >&2
            exit 2
        fi
        # Validate via ssh-keygen. The output contains the SHA256
        # fingerprint we'll use as the key's identifier.
        tmp="$(mktemp "${TMPDIR:-/tmp}/air-ssh-keys.XXXXXX")"
        trap 'rm -f "$tmp"' EXIT
        printf '%s\n' "$line" > "$tmp"
        if ! fpr_line="$(ssh-keygen -lf "$tmp" 2>/dev/null)"; then
            echo "air-ssh-keys: not a valid SSH public key" >&2
            exit 2
        fi
        # ssh-keygen -lf prints: "<bits> SHA256:abc...  comment (ED25519)"
        fpr="$(printf '%s' "$fpr_line" | awk '{print $2}')"
        # Dedupe by fingerprint — don't add the same key twice.
        if grep -v '^\s*#' "${AUTH_FILE}" | \
           while IFS= read -r existing; do
               [ -n "$existing" ] || continue
               existing_fpr="$(printf '%s\n' "$existing" | ssh-keygen -lf - 2>/dev/null | awk '{print $2}' || true)"
               [ "$existing_fpr" = "$fpr" ] && exit 7
           done; [ $? -ne 7 ]; then
            # Count enforcement
            current="$(grep -cEv '^\s*(#|$)' "${AUTH_FILE}" || echo 0)"
            if [ "$current" -ge "$MAX_KEYS" ]; then
                echo "air-ssh-keys: max ${MAX_KEYS} keys already present" >&2
                exit 3
            fi
            # Append a clean, options-free line. Keep only type+key+comment.
            # awk tokens: $1=type, $2=b64, $3..=comment (optional).
            cleaned="$(printf '%s' "$line" | awk '{
                type=$1; key=$2;
                comment="";
                for (i=3;i<=NF;i++) comment=(comment=="")?$i:comment" "$i;
                if (comment=="") print type" "key;
                else            print type" "key" "comment;
            }')"
            printf '%s\n' "$cleaned" >> "${AUTH_FILE}"
            chown air:air "${AUTH_FILE}"
            chmod 0600 "${AUTH_FILE}"
        fi
        printf '%s\n' "$fpr"
        ;;

    remove)
        ensure_layout
        if [ "$#" -ne 2 ]; then
            echo "air-ssh-keys: remove <fingerprint>" >&2
            exit 2
        fi
        fpr="${2:-}"
        if [ -z "$fpr" ]; then
            echo "air-ssh-keys: remove <fingerprint>" >&2
            exit 2
        fi
        # OpenSSH SHA256 fingerprints are unpadded base64-ish strings.
        # Use a real regex rather than a shell glob so trailing junk,
        # whitespace, or metacharacters cannot be accepted accidentally.
        if ! [[ "$fpr" =~ ^SHA256:[A-Za-z0-9+/]{32,64}$ ]]; then
            echo "air-ssh-keys: unexpected fingerprint format" >&2
            exit 2
        fi
        tmp="$(mktemp "${TMPDIR:-/tmp}/air-ssh-keys.XXXXXX")"
        trap 'rm -f "$tmp"' EXIT
        removed=0
        while IFS= read -r existing || [ -n "$existing" ]; do
            case "$existing" in
                ''|\#*)
                    printf '%s\n' "$existing" >> "$tmp"
                    continue
                    ;;
            esac
            existing_fpr="$(printf '%s\n' "$existing" | ssh-keygen -lf - 2>/dev/null | awk '{print $2}' || true)"
            if [ "$existing_fpr" = "$fpr" ]; then
                removed=1
                continue
            fi
            printf '%s\n' "$existing" >> "$tmp"
        done < "${AUTH_FILE}"
        install -m 0600 -o air -g air "$tmp" "${AUTH_FILE}"
        if [ "$removed" -eq 0 ]; then
            echo "air-ssh-keys: no matching key" >&2
            exit 4
        fi
        ;;

    *)
        echo "usage: $0 list | add | remove <fingerprint>" >&2
        exit 2
        ;;
esac
