#!/usr/bin/env bash
# ============================================================
#  CML Certificate Refresh
#  Author  : Mike Dent
#  Email   : mike@mikedent.io
#  Version : 1.0
#  Date    : 2026-05-02
#  Desc    : Installs or refreshes wildcard TLS certificates on
#            a Cisco Modeling Labs (CML) appliance for both
#            nginx (CML web UI :443) and Cockpit (:9090) from
#            a single cert+key pair staged in $HOME/certs.
# ============================================================
#
# Usage:
#   ~/certs/cml_cert_refresh.sh [--cert PATH --key PATH] [--rollback]
#                               [--dry-run] [--yes] [--help]
#
# Inputs (precedence high → low):
#   1. --cert PATH / --key PATH flags  (explicit, can be anywhere)
#   2. CERT_FILE / KEY_FILE env vars   (filenames within CERT_DIR)
#   3. Auto-discovery in CERT_DIR      (interactive picker if multiple
#                                       certs/keys are present)
#
#   CERT_DIR defaults to $HOME/certs.
#
# Workflow:
#   1. sftp the cert + key + this script to $HOME/certs on the CML host.
#   2. ssh in and run:  ~/certs/cml_cert_refresh.sh
#   3. Confirm the discovered cert/key at the prompt.
#   4. For renewal, replace the cert/key in $HOME/certs and re-run.
#

set -euo pipefail

# ---- Configuration ---------------------------------------------------------
CERT_DIR="${CERT_DIR:-$HOME/certs}"
CERT_FILE="${CERT_FILE:-}"
KEY_FILE="${KEY_FILE:-}"

# Resolved full paths to the cert and key. Populated by resolve_certs() from
# (in order) --cert/--key flags, CERT_FILE/KEY_FILE env vars, or auto-discovery
# in CERT_DIR.
CERT_PATH=""
KEY_PATH=""

NGINX_PUBKEY="/etc/nginx/pubkey.pem"
NGINX_PRIVKEY="/etc/nginx/privkey.pem"

COCKPIT_DIR="/etc/cockpit/ws-certs.d"
# We hijack cockpit.cert / cockpit.key — the files cockpit-certificate-ensure
# auto-generates on this CML build. Replacing them in place with symlinks to
# the staged cert+key is stable: cockpit-certificate-ensure leaves existing
# files alone, and cockpit serves the lex-last *.cert in this directory.
COCKPIT_CERT="${COCKPIT_DIR}/cockpit.cert"
COCKPIT_KEY="${COCKPIT_DIR}/cockpit.key"

DRY_RUN=0
ASSUME_YES=0
MODE="install"

# ---- Logging ---------------------------------------------------------------
info() { printf '==> %s\n' "$*"; }
warn() { printf 'WARN: %s\n' "$*" >&2; }
err()  { printf 'ERROR: %s\n' "$*" >&2; }

run() {
  if (( DRY_RUN )); then
    printf '    [dry-run] %s\n' "$*"
  else
    "$@"
  fi
}

usage() {
  cat <<EOF
Usage: $(basename "$0") [--cert PATH --key PATH] [--rollback]
                              [--dry-run] [--yes] [--help]

Install or refresh wildcard TLS certificates on a CML appliance for both
nginx (CML web UI on :443) and Cockpit (:9090).

Options:
  --cert PATH  Path to the fullchain certificate (overrides discovery)
  --key  PATH  Path to the private key           (overrides discovery)
  --rollback   Restore the original self-signed certificates from .bak files
  --dry-run    Print actions without executing them (skips confirmation)
  -y, --yes    Skip the cert-confirmation prompt (for automation)
  -h, --help   Show this help

Cert/key resolution (high → low precedence):
  1. --cert and --key flags
  2. CERT_FILE and KEY_FILE env vars (filenames within CERT_DIR)
  3. Auto-discovery in CERT_DIR (openssl-based; interactive picker if
     multiple certs or keys are present)

Environment overrides:
  CERT_DIR    Directory containing cert files     (default: \$HOME/certs)
  CERT_FILE   Filename of the fullchain cert      (default: auto-discover)
  KEY_FILE    Filename of the private key         (default: auto-discover)
EOF
}

# ---- Args ------------------------------------------------------------------
while (( $# > 0 )); do
  case "$1" in
    --cert)     [[ $# -ge 2 ]] || { err "--cert requires a path"; exit 2; }
                CERT_PATH="$2"; shift ;;
    --cert=*)   CERT_PATH="${1#*=}" ;;
    --key)      [[ $# -ge 2 ]] || { err "--key requires a path"; exit 2; }
                KEY_PATH="$2"; shift ;;
    --key=*)    KEY_PATH="${1#*=}" ;;
    --rollback) MODE="rollback" ;;
    --dry-run)  DRY_RUN=1 ;;
    -y|--yes)   ASSUME_YES=1 ;;
    -h|--help)  usage; exit 0 ;;
    *) err "Unknown argument: $1"; usage; exit 2 ;;
  esac
  shift
done

# ---- Helpers ---------------------------------------------------------------

require_sudo() {
  if ! sudo -v; then
    err "This script requires sudo privileges."
    exit 1
  fi
}

cert_pubkey_hash() {
  openssl x509 -in "$1" -noout -pubkey 2>/dev/null | openssl md5 | awk '{print $NF}'
}

key_pubkey_hash() {
  openssl pkey -in "$1" -pubout 2>/dev/null | openssl md5 | awk '{print $NF}'
}

cert_subject() {
  openssl x509 -in "$1" -noout -subject 2>/dev/null | sed 's/subject= *//'
}

cert_not_after() {
  openssl x509 -in "$1" -noout -enddate 2>/dev/null | sed 's/notAfter=//'
}

cert_san() {
  openssl x509 -in "$1" -noout -ext subjectAltName 2>/dev/null \
    | grep -oE 'DNS:[^,[:space:]]+' | sed 's/DNS://'
}

# ---- Cert/key resolution ---------------------------------------------------

# Read a single line from /dev/tty if available, else stdin. $1 = prompt.
# Returns 1 on EOF / read failure so callers can distinguish "user typed
# nothing" from "no input is coming" and avoid spin-looping.
prompt_read() {
  local reply
  printf '%s' "$1" >&2
  if [[ -r /dev/tty ]]; then
    read -r reply </dev/tty 2>/dev/null || return 1
  else
    read -r reply 2>/dev/null || return 1
  fi
  printf '%s' "$reply"
}

# Print a numbered list of choices to stderr and read a 1-based selection.
# Echoes the chosen value on stdout. $1 = kind label, $2... = choices.
pick_one() {
  local kind="$1"; shift
  local choices=( "$@" )
  local i choice

  printf 'Multiple %s files found in %s — pick one:\n' "$kind" "$CERT_DIR" >&2
  for i in "${!choices[@]}"; do
    printf '  [%d] %s\n' "$((i+1))" "${choices[$i]}" >&2
  done

  local attempts=0
  while true; do
    if ! choice="$(prompt_read "Enter number [1-${#choices[@]}]: ")"; then
      err "No input available — cannot pick a $kind interactively"
      err "Use --cert PATH --key PATH or CERT_FILE=NAME KEY_FILE=NAME to bypass discovery"
      exit 1
    fi
    if [[ "$choice" =~ ^[0-9]+$ ]] && (( choice >= 1 && choice <= ${#choices[@]} )); then
      printf '%s\n' "${choices[$((choice-1))]}"
      return 0
    fi
    attempts=$((attempts + 1))
    if (( attempts >= 5 )); then
      err "Too many invalid choices — giving up"
      exit 1
    fi
    printf 'Invalid choice. Try again.\n' >&2
  done
}

resolve_certs() {
  # 1. Explicit --cert/--key flags win outright.
  if [[ -n "$CERT_PATH" && -n "$KEY_PATH" ]]; then
    info "Using cert and key from --cert / --key flags"
    return 0
  fi
  if [[ -n "$CERT_PATH" || -n "$KEY_PATH" ]]; then
    err "--cert and --key must be provided together"
    exit 2
  fi

  # 2. CERT_FILE/KEY_FILE env vars (filenames within CERT_DIR).
  if [[ -n "$CERT_FILE" && -n "$KEY_FILE" ]]; then
    CERT_PATH="${CERT_DIR}/${CERT_FILE}"
    KEY_PATH="${CERT_DIR}/${KEY_FILE}"
    info "Using cert and key from CERT_FILE / KEY_FILE env vars"
    return 0
  fi
  if [[ -n "$CERT_FILE" || -n "$KEY_FILE" ]]; then
    err "CERT_FILE and KEY_FILE env vars must be set together"
    exit 2
  fi

  # 3. Auto-discover in CERT_DIR using openssl to identify each candidate.
  info "Discovering cert and key in ${CERT_DIR}"
  if [[ ! -d "$CERT_DIR" ]]; then
    err "${CERT_DIR} does not exist"
    err "Use --cert PATH --key PATH, or set CERT_DIR to a directory containing your cert and key"
    exit 1
  fi

  local certs=() keys=()
  local f
  shopt -s nullglob
  local candidates=( "${CERT_DIR}"/*.pem "${CERT_DIR}"/*.crt "${CERT_DIR}"/*.key )
  shopt -u nullglob

  for f in "${candidates[@]}"; do
    if openssl x509 -in "$f" -noout >/dev/null 2>&1; then
      certs+=( "$f" )
    elif openssl pkey -in "$f" -noout >/dev/null 2>&1; then
      keys+=( "$f" )
    fi
  done

  if (( ${#certs[@]} == 0 )); then
    err "No certificate files found in ${CERT_DIR}"
    err "Looked at *.pem *.crt *.key — none matched 'openssl x509 -noout'"
    err "Use --cert PATH --key PATH to point at files explicitly"
    exit 1
  fi
  if (( ${#keys[@]} == 0 )); then
    err "No private key files found in ${CERT_DIR}"
    err "Looked at *.pem *.crt *.key — none matched 'openssl pkey -noout'"
    err "Use --cert PATH --key PATH to point at files explicitly"
    exit 1
  fi

  if (( ${#certs[@]} == 1 )); then
    CERT_PATH="${certs[0]}"
    info "  Auto-discovered cert: ${CERT_PATH}"
  else
    CERT_PATH="$(pick_one "certificate" "${certs[@]}")"
  fi

  if (( ${#keys[@]} == 1 )); then
    KEY_PATH="${keys[0]}"
    info "  Auto-discovered key:  ${KEY_PATH}"
  else
    KEY_PATH="$(pick_one "private key" "${keys[@]}")"
  fi
}

# ---- Pre-flight ------------------------------------------------------------

# check "<description>" <command> [args...]
# Prints a "  -> <desc> ... ok|FAIL" line and returns the command's exit code.
check() {
  local desc="$1"; shift
  printf '  -> %-44s ' "$desc"
  if "$@" >/dev/null 2>&1; then
    printf 'ok\n'
    return 0
  else
    printf 'FAIL\n'
    return 1
  fi
}

preflight() {
  info "Pre-flight checks"
  info "  Cert: ${CERT_PATH}"
  info "  Key:  ${KEY_PATH}"

  check "cert file exists"         test -f "$CERT_PATH" || { err "Required file not found: $CERT_PATH"; exit 1; }
  check "cert file readable"       test -r "$CERT_PATH" || { err "Cannot read file: $CERT_PATH"; exit 1; }
  check "key file exists"          test -f "$KEY_PATH"  || { err "Required file not found: $KEY_PATH"; exit 1; }
  check "key file readable"        test -r "$KEY_PATH"  || { err "Cannot read file: $KEY_PATH"; exit 1; }
  check "cert is valid PEM"        openssl x509 -in "$CERT_PATH" -noout || { err "Cert is not a valid PEM certificate: $CERT_PATH"; exit 1; }
  check "key is valid PEM"         openssl pkey -in "$KEY_PATH" -noout  || { err "Key is not a valid PEM private key: $KEY_PATH"; exit 1; }

  local cert_hash key_hash
  cert_hash="$(cert_pubkey_hash "$CERT_PATH")"
  key_hash="$(key_pubkey_hash "$KEY_PATH")"
  check "cert/key public-key match" test "$cert_hash" = "$key_hash" || { err "Certificate and private key do not match (public-key hash mismatch)"; exit 1; }

  check "cert is not expired"      openssl x509 -in "$CERT_PATH" -noout -checkend 0 || { err "Certificate is expired"; exit 1; }

  check "/etc/nginx present"       test -d /etc/nginx        || { err "/etc/nginx not found — is this a CML appliance?"; exit 1; }
  check "${COCKPIT_DIR} present"   test -d "${COCKPIT_DIR}"  || { err "${COCKPIT_DIR} not found — is Cockpit installed?"; exit 1; }
  check "systemctl available"      command -v systemctl      || { err "systemctl not available"; exit 1; }
  check "cockpit-ws group exists"  getent group cockpit-ws   || { err "Group 'cockpit-ws' not found — Cockpit may not be installed correctly"; exit 1; }

  info "  Subject: $(cert_subject "$CERT_PATH")"
  info "  Expires: $(cert_not_after "$CERT_PATH")"
  info "  SANs:    $(cert_san "$CERT_PATH" | paste -sd ',' -)"
}

# ---- Confirmation ----------------------------------------------------------

confirm_certs() {
  if (( ASSUME_YES )) || (( DRY_RUN )); then
    return 0
  fi

  local reply
  reply="$(prompt_read 'Proceed with these certs? [y/N] ')" || reply=""

  case "${reply,,}" in
    y|yes) return 0 ;;
    *)
      info "Aborted. To use different files, re-run with one of:"
      info "  $(basename "$0") --cert PATH --key PATH"
      info "  CERT_FILE=NAME KEY_FILE=NAME $(basename "$0")"
      exit 0
      ;;
  esac
}

# ---- Install: nginx --------------------------------------------------------

install_nginx() {
  info "Configuring nginx"
  local cert_path="$CERT_PATH"
  local key_path="$KEY_PATH"

  # pubkey.pem
  if [[ -L "${NGINX_PUBKEY}" ]] && [[ "$(readlink "${NGINX_PUBKEY}")" == "${cert_path}" ]]; then
    info "  ${NGINX_PUBKEY} already linked correctly"
  else
    if [[ -e "${NGINX_PUBKEY}" && ! -e "${NGINX_PUBKEY}.bak" ]]; then
      run sudo mv "${NGINX_PUBKEY}" "${NGINX_PUBKEY}.bak"
    elif [[ -e "${NGINX_PUBKEY}" || -L "${NGINX_PUBKEY}" ]]; then
      run sudo rm -f "${NGINX_PUBKEY}"
    fi
    run sudo ln -sf "${cert_path}" "${NGINX_PUBKEY}"
    info "  ${NGINX_PUBKEY} -> ${cert_path}"
  fi

  # privkey.pem
  if [[ -L "${NGINX_PRIVKEY}" ]] && [[ "$(readlink "${NGINX_PRIVKEY}")" == "${key_path}" ]]; then
    info "  ${NGINX_PRIVKEY} already linked correctly"
  else
    if [[ -e "${NGINX_PRIVKEY}" && ! -e "${NGINX_PRIVKEY}.bak" ]]; then
      run sudo mv "${NGINX_PRIVKEY}" "${NGINX_PRIVKEY}.bak"
    elif [[ -e "${NGINX_PRIVKEY}" || -L "${NGINX_PRIVKEY}" ]]; then
      run sudo rm -f "${NGINX_PRIVKEY}"
    fi
    run sudo ln -sf "${key_path}" "${NGINX_PRIVKEY}"
    info "  ${NGINX_PRIVKEY} -> ${key_path}"
  fi
}

# ---- Install: Cockpit ------------------------------------------------------

install_cockpit() {
  info "Configuring Cockpit"
  local cert_path="$CERT_PATH"
  local key_path="$KEY_PATH"

  # Cockpit reader runs in group `cockpit-ws`. Symlink targets are read with
  # the target's mode/ownership, so the source key must be group-readable by
  # cockpit-ws.
  run sudo chgrp cockpit-ws "${key_path}"
  run sudo chmod 640        "${key_path}"
  run sudo chmod 644        "${cert_path}"

  # Clean up leftover 50-mdlab.{cert,key} from earlier versions of this
  # script that installed at the wrong path. They lose to cockpit.cert via
  # lex-ordering anyway, so removing them is just tidying.
  for stale in "${COCKPIT_DIR}/50-mdlab.cert" "${COCKPIT_DIR}/50-mdlab.key"; do
    if [[ -L "$stale" || -e "$stale" ]]; then
      run sudo rm -f "$stale"
    fi
  done

  # cockpit.cert
  if [[ -L "${COCKPIT_CERT}" ]] && [[ "$(readlink "${COCKPIT_CERT}")" == "${cert_path}" ]]; then
    info "  ${COCKPIT_CERT} already linked correctly"
  else
    if [[ -e "${COCKPIT_CERT}" && ! -L "${COCKPIT_CERT}" && ! -e "${COCKPIT_CERT}.bak" ]]; then
      run sudo mv "${COCKPIT_CERT}" "${COCKPIT_CERT}.bak"
    elif [[ -e "${COCKPIT_CERT}" || -L "${COCKPIT_CERT}" ]]; then
      run sudo rm -f "${COCKPIT_CERT}"
    fi
    run sudo ln -sf "${cert_path}" "${COCKPIT_CERT}"
    info "  ${COCKPIT_CERT} -> ${cert_path}"
  fi

  # cockpit.key
  if [[ -L "${COCKPIT_KEY}" ]] && [[ "$(readlink "${COCKPIT_KEY}")" == "${key_path}" ]]; then
    info "  ${COCKPIT_KEY} already linked correctly"
  else
    if [[ -e "${COCKPIT_KEY}" && ! -L "${COCKPIT_KEY}" && ! -e "${COCKPIT_KEY}.bak" ]]; then
      run sudo mv "${COCKPIT_KEY}" "${COCKPIT_KEY}.bak"
    elif [[ -e "${COCKPIT_KEY}" || -L "${COCKPIT_KEY}" ]]; then
      run sudo rm -f "${COCKPIT_KEY}"
    fi
    run sudo ln -sf "${key_path}" "${COCKPIT_KEY}"
    info "  ${COCKPIT_KEY} -> ${key_path}"
  fi
}

# ---- Rollback --------------------------------------------------------------

rollback_nginx() {
  info "Rolling back nginx"
  for f in "${NGINX_PUBKEY}" "${NGINX_PRIVKEY}"; do
    if [[ -L "$f" || -e "$f" ]]; then
      run sudo rm -f "$f"
    fi
    if [[ -e "${f}.bak" ]]; then
      run sudo mv "${f}.bak" "$f"
      info "  Restored $f"
    else
      warn "  No backup found at ${f}.bak"
    fi
  done
}

rollback_cockpit() {
  info "Rolling back Cockpit"
  for f in "${COCKPIT_CERT}" "${COCKPIT_KEY}"; do
    if [[ -L "$f" || -e "$f" ]]; then
      run sudo rm -f "$f"
    fi
    if [[ -e "${f}.bak" ]]; then
      run sudo mv "${f}.bak" "$f"
      info "  Restored $f"
    else
      warn "  No backup found at ${f}.bak"
    fi
  done

  # Clean up leftover 50-mdlab.{cert,key} from earlier versions
  for stale in "${COCKPIT_DIR}/50-mdlab.cert" "${COCKPIT_DIR}/50-mdlab.key"; do
    if [[ -L "$stale" || -e "$stale" ]]; then
      run sudo rm -f "$stale"
    fi
  done
}

# ---- Restart services ------------------------------------------------------

restart_services() {
  info "Restarting services"
  if ! (( DRY_RUN )); then
    if ! sudo nginx -t; then
      err "nginx config test failed; not restarting services"
      exit 1
    fi
  fi

  for svc in nginx cockpit; do
    if (( DRY_RUN )); then
      printf '    [dry-run] sudo timeout 60 systemctl restart %s\n' "$svc"
      continue
    fi
    info "  Restarting ${svc}"
    if ! sudo timeout 60 systemctl restart "$svc"; then
      err "${svc} restart failed or timed out after 60s"
      err "Last 40 lines of journalctl -u ${svc}:"
      sudo journalctl -u "$svc" -n 40 --no-pager >&2 || true
      exit 1
    fi
  done
}

# ---- Post-flight -----------------------------------------------------------

postflight() {
  if (( DRY_RUN )); then
    info "Skipping post-flight checks (dry-run)"
    return
  fi

  info "Post-flight checks"

  for svc in nginx cockpit; do
    if systemctl is-active --quiet "$svc"; then
      info "  ${svc}: active"
    else
      warn "  ${svc}: not active"
    fi
  done

  local expected_fp
  expected_fp="$(openssl x509 -in "$CERT_PATH" -noout -fingerprint -sha256 | awk -F= '{print $2}')"

  for port in 443 9090; do
    local served_fp
    served_fp="$(echo \
      | openssl s_client -connect "localhost:${port}" -servername "$(hostname -f)" 2>/dev/null \
      | openssl x509 -noout -fingerprint -sha256 2>/dev/null \
      | awk -F= '{print $2}')"
    if [[ -z "$served_fp" ]]; then
      warn "  port ${port}: could not retrieve served cert"
    elif [[ "$served_fp" == "$expected_fp" ]]; then
      info "  port ${port}: served cert matches installed cert"
    else
      warn "  port ${port}: served cert does NOT match installed cert"
      warn "    expected: $expected_fp"
      warn "    served:   $served_fp"
    fi
  done
}

# ---- Main ------------------------------------------------------------------

main() {
  require_sudo

  case "$MODE" in
    install)
      resolve_certs
      preflight
      confirm_certs
      install_nginx
      install_cockpit
      restart_services
      postflight
      info "Done."
      info "  CML web UI: https://$(hostname -f)/"
      info "  Cockpit:    https://$(hostname -f):9090/"
      ;;
    rollback)
      rollback_nginx
      rollback_cockpit
      restart_services
      info "Rollback complete."
      ;;
  esac
}

main "$@"