#! /usr/bin/env bash

set -euo pipefail
 
CTID="${1}"
PHASE="${2}"

# Check for these packages and install if missing
REQUIRED_PACKAGES=(
    locales
    sudo
    curl
    gnupg
    lsb-release
    ca-certificates
    apt-transport-https
)
 
# Preferred local and encoding
LOCALE_NAME="en_US.UTF-8"
LOCALE_CHARSET="UTF-8"
LOCALE_ENTRY="${LOCALE_NAME} ${LOCALE_CHARSET}"

# Adjust as needed for your internal PKI
REQUIRED_CA_CN="Home Lab Root CA"
CA_FETCH_URL="https://sub-ca.pki.home.internal/roots.pem"
CA_DEST_PATH="/usr/local/share/ca-certificates/internal-intermediate.crt"
# Search these directories to see if Root CA installed
CA_TRUST_DIRS=(
    /etc/ssl/certs
    /usr/local/share/ca-certificates
    /usr/share/ca-certificates
)

# Sentinel file to track successful run of script
# Since this script is long-running, we don't want to run at every boot
SENTINEL_FILE="/etc/lxc-provisioned"

# Logger functions 
log()  { echo "[hook][CT ${CTID}][${PHASE}] $*"; }
info() { log "INFO  $*"; }
warn() { log "WARN  $*"; }
fail() { log "ERROR $*"; exit 1; }

# Wrapper around pct exec
ct_exec() {
    pct exec "${CTID}" -- "$@"
}
 
# Upgrade packages on LXC
apt_update_upgrade() {
    info "Running apt update && apt upgrade -y ..."
    ct_exec env DEBIAN_FRONTEND=noninteractive apt-get update -qq
    ct_exec env DEBIAN_FRONTEND=noninteractive apt-get upgrade -y -qq
    info "System packages are up to date."
}
 

# Check for required packages from array above
ensure_packages() {
    info "Checking required packages ..."
 
    local missing=()
 
    for pkg in "${REQUIRED_PACKAGES[@]}"; do
        if ct_exec dpkg-query -W -f='${db:Status-Abbrev}' "${pkg}" 2>/dev/null \
                | grep -q "^ii"; then
            info "  [OK]      ${pkg}"
        else
            warn "  [MISSING] ${pkg} -- will install"
            missing+=("${pkg}")
        fi
    done
 
    if [[ ${#missing[@]} -eq 0 ]]; then
        info "All required packages are already installed."
        return
    fi
 
    info "Installing missing packages: ${missing[*]} ..."
    ct_exec env DEBIAN_FRONTEND=noninteractive \
        apt-get install -y -qq "${missing[@]}" \
        || fail "apt-get install failed for: ${missing[*]}"
 
    local still_missing=()
    for pkg in "${missing[@]}"; do
        if ct_exec dpkg-query -W -f='${db:Status-Abbrev}' "${pkg}" 2>/dev/null \
                | grep -q "^ii"; then
            info "  [INSTALLED] ${pkg}"
        else
            still_missing+=("${pkg}")
        fi
    done
 
    if [[ ${#still_missing[@]} -gt 0 ]]; then
        fail "Package(s) could not be installed: ${still_missing[*]}"
    fi
 
    info "All required packages are now installed."
}
 
# Set system locale and encoding
# Adjust as needed using variable above
ensure_locale() {
    info "Checking locale (${LOCALE_ENTRY}) ..."
 
    # Check if the locale is already generated
    if ct_exec locale -a 2>/dev/null | grep -qi "^${LOCALE_NAME//./\\.}$"; then
        info "Locale ${LOCALE_NAME} is already generated."
    else
        info "Locale ${LOCALE_NAME} not found -- configuring ..."
 
        # Ensure the entry is present and uncommented in /etc/locale.gen
        local locale_gen="/etc/locale.gen"
 
        # Remove any existing (commented or uncommented) entry for this locale
        # then append the correct uncommented entry
        ct_exec bash -c "
            sed -i '/^#\?\s*${LOCALE_ENTRY}/d' ${locale_gen}
            echo '${LOCALE_ENTRY}' >> ${locale_gen}
        " || fail "Failed to update ${locale_gen}"
 
        # Generate the locale
        ct_exec locale-gen \
            || fail "locale-gen failed"
 
        info "Locale ${LOCALE_NAME} generated."
    fi
 
    # Check if LANG is already set correctly system-wide
    local current_lang
    current_lang=$(ct_exec bash -c "
        grep -oP '(?<=^LANG=).+' /etc/default/locale 2>/dev/null || true
    ")
 
    if [[ "${current_lang}" == "${LOCALE_NAME}" ]]; then
        info "System locale already set to ${LOCALE_NAME}."
    else
        info "Setting system locale to ${LOCALE_NAME} ..."
        ct_exec update-locale "LANG=${LOCALE_NAME}" \
            || fail "update-locale failed"
        info "System locale set to ${LOCALE_NAME}."
    fi
}
 
# Checks for Root CA certificate from Step CA
# Update variable at top of script to adjust for your environment
ensure_ca_certificate() {
    info "Checking for CA certificate with CN '${REQUIRED_CA_CN}' ..."
 
    # Collapse the entire search into a single pct exec to eliminate per-cert
    # round-trip overhead. Inside the container, xargs -P fans out openssl
    # calls across multiple workers simultaneously.
    #
    # Output format on match: FOUND:<cert_path>:<subject_line>
    # Output on no match:     (empty)
    local scan_result
    scan_result=$(ct_exec bash -s -- "${REQUIRED_CA_CN}" "${CA_TRUST_DIRS[@]}" <<'SCAN'
        TARGET_CN="${1}"; shift
        SEARCH_DIRS=("$@")
 
        # Exported so the xargs subshell can call it
        _check_cert() {
            local cert="${1}"
            local target_cn="${2}"
            local subject
            subject=$(openssl x509 -noout -subject -in "${cert}" 2>/dev/null) || return
            if echo "${subject}" | grep -qi "CN\s*=\s*${target_cn}"; then
                printf 'FOUND:%s:%s\n' "${cert}" "${subject}"
            fi
        }
        export -f _check_cert
 
        find "${SEARCH_DIRS[@]}" \
            \( -name "*.crt" -o -name "*.pem" \) \
            -type f 2>/dev/null \
        | xargs -P "$(nproc)" -I{} \
            bash -c '_check_cert "$@"' _ {} "${TARGET_CN}" \
        | head -1
SCAN
    ) || true
 
    if [[ "${scan_result}" == FOUND:* ]]; then
        # Strip the leading "FOUND:" tag and split on first colon after the path
        local cert subject
        cert="${scan_result#FOUND:}"
        subject="${cert#*:}"
        cert="${cert%%:*}"
        info "  [FOUND] ${cert}"
        info "          Subject: ${subject}"
        info "CA certificate '${REQUIRED_CA_CN}' is already trusted."
        return
    fi
 
    warn "CA certificate '${REQUIRED_CA_CN}' not found -- fetching from ${CA_FETCH_URL} ..."
 
    ct_exec mkdir -p "$(dirname "${CA_DEST_PATH}")"
 
    # -k used because the system may not yet trust the internal CA
    curl_exit_code=0
    ct_exec curl -fsSk "${CA_FETCH_URL}" -o "${CA_DEST_PATH}" --connect-timeout 3 \
      || curl_exit_code=$?

    case $curl_exit_code in
        0)
             info "Certificate saved to ${CA_DEST_PATH} -- running update-ca-certificates ..."
             ct_exec update-ca-certificates \
                || fail "update-ca-certificates failed"

             info "CA certificate '${REQUIRED_CA_CN}' is now trusted."
             ;;

        6|7|28)
            warn "  Unable to retrieve Root CA certificate due to network error."
            warn "  Error code: ${curl_exit_code}"
            warn "  List of exit codes: https://everything.curl.dev/cmdline/exitcode.html"
            ;;

        *) fail "Failed to download CA certificate from ${CA_FETCH_URL}" ;;
    esac
}
 
# Use the post-start phase to configure the LXC
# Uses 'pct exec' to run commands inside the container after boot
case "${PHASE}" in
    post-start)
        # Invoke helper functions from above
        if ct_exec test -f "${SENTINEL_FILE}" 2>/dev/null; then
            local_ts=$(ct_exec cat "${SENTINEL_FILE}" 2>/dev/null || true)
            info "Already provisioned on ${local_ts} -- skipping."
            exit 0
        fi
 
        info "First boot detected -- running provisioning ..."
        apt_update_upgrade
        ensure_packages
        ensure_locale
        ensure_ca_certificate
 
        # Write the sentinel file with an ISO-8601 timestamp so it is easy to
        # audit when a container was first provisioned.
        local_ts=$(date -u +"%Y-%m-%dT%H:%M:%SZ")
        ct_exec bash -c "echo '${local_ts}' > '${SENTINEL_FILE}'" \
            || warn "Could not write sentinel file -- provisioning will re-run on next boot."
 
        info "Provisioning complete. Sentinel written at ${SENTINEL_FILE} (${local_ts})."
        ;;
    pre-start | pre-stop | post-stop | mount)
        info "Phase '${PHASE}' -- nothing to do; skipping."
        ;;
    *)
        warn "Unknown phase '${PHASE}' -- skipping."
        ;;
esac
 
exit 0