#!/bin/bash
###############################################################################
# CertAuth Setup Script
# Provisions a fresh Ubuntu machine as a Certificate Authority
#
# Prerequisites:
#   - Fresh Ubuntu 24.04+ Server (aarch64/x86_64)
#   - Two YubiKey 5 Nano devices plugged in
#   - Root or sudo access
#   - Network access to Ubuntu repositories
#
# Usage:
#   sudo bash setup-certauth.sh
#
# Configuration: Edit the variables below or pass as environment variables
###############################################################################

set -euo pipefail

# ===================== CONFIGURATION =====================
CA_ORG="${CA_ORG:-Home}"
CA_COUNTRY="${CA_COUNTRY:-US}"
ROOT_CA_CN="${ROOT_CA_CN:-certauth Root CA}"
INT_CA_CN="${INT_CA_CN:-certauth Intermediate CA}"
ROOT_VALID_DAYS="${ROOT_VALID_DAYS:-9125}"    # 25 years
INT_VALID_DAYS="${INT_VALID_DAYS:-5475}"     # 15 years
LEAF_VALID_DAYS="${LEAF_VALID_DAYS:-365}"    # 1 year per cert
NETWORK_CIDR="${NETWORK_CIDR:-192.168.8.0/24}"
ADMIN_USER="${ADMIN_USER:-certauth}"
ADMIN_PASS="${ADMIN_PASS:-CHANGE_ME_ADMIN_PASS}"

# YubiKey assignment (will be detected automatically if not set)
YK1_SERIAL="${YK1_SERIAL:-}"   # Root CA YubiKey serial
YK2_SERIAL="${YK2_SERIAL:-}"   # Intermediate CA YubiKey serial

# PINs (will be generated if not set)
YK1_PIN="${YK1_PIN:-}"
YK1_PUK="${YK1_PUK:-}"
YK2_PIN="${YK2_PIN:-}"
YK2_PUK="${YK2_PUK:-}"

# ===================== COLORS =====================
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

info()  { echo -e "${BLUE}[INFO]${NC} $*"; }
success() { echo -e "${GREEN}[OK]${NC} $*"; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $*"; }
error() { echo -e "${RED}[ERROR]${NC} $*"; exit 1; }

# ===================== PRE-FLIGHT CHECKS =====================
check_prerequisites() {
    info "Running pre-flight checks..."

    # Must be root
    if [[ $EUID -ne 0 ]]; then
        error "This script must be run as root (use sudo)"
    fi

    # Check Ubuntu version
    if ! grep -q "Ubuntu" /etc/os-release; then
        error "This script requires Ubuntu"
    fi
    UBUNTU_VERSION=$(grep VERSION_ID /etc/os-release | cut -d'"' -f2)
    info "Ubuntu version: $UBUNTU_VERSION"

    # Check architecture
    ARCH=$(uname -m)
    info "Architecture: $ARCH"

    # Check YubiKeys
    YK_COUNT=$(lsusb 2>/dev/null | grep -c "Yubico" || true)
    if [[ $YK_COUNT -lt 2 ]]; then
        error "Need at least 2 YubiKeys connected (found: $YK_COUNT)"
    fi
    success "Found $YK_COUNT YubiKeys"

    # Check disk space
    AVAIL=$(df -BG / | awk 'NR==2 {print $4}' | tr -d 'G')
    if [[ $AVAIL -lt 5 ]]; then
        error "Need at least 5GB free disk space (have: ${AVAIL}GB)"
    fi
    success "Disk space: ${AVAIL}GB available"
}

# ===================== STEP 1: SYSTEM SETUP =====================
setup_system() {
    info "=== Step 1: System Setup ==="

    # Update system
    info "Updating package lists..."
    apt-get update -qq

    # Install base packages
    info "Installing packages..."
    DEBIAN_FRONTEND=noninteractive apt-get install -y -qq \
        ufw fail2ban opensc pcscd yubikey-manager ykcs11 \
        python3-pip python3-venv caddy sqlite3 libsqlite3-dev \
        haveged apparmor-utils curl wget \
        2>/dev/null || true

    # Install Python packages
    info "Installing Python packages..."
    pip3 install --break-system-packages --ignore-installed \
        typing_extensions 2>/dev/null || true
    pip3 install --break-system-packages \
        fastapi "uvicorn[standard]" jinja2 python-multipart \
        bcrypt "python-jose[cryptography]" ecdsa cryptography \
        2>/dev/null || true

    success "Packages installed"
}

# ===================== STEP 2: CREATE USER =====================
setup_user() {
    info "=== Step 2: Create $ADMIN_USER user ==="

    if ! id "$ADMIN_USER" &>/dev/null; then
        useradd -m -s /bin/bash "$ADMIN_USER"
        usermod -aG sudo,adm,plugdev,pcscd "$ADMIN_USER"
        echo "$ADMIN_USER:$ADMIN_PASS" | chpasswd
        echo "$ADMIN_USER ALL=(ALL) NOPASSWD: ALL" > /etc/sudoers.d/certauth
        chmod 0440 /etc/sudoers.d/certauth
        success "User $ADMIN_USER created"
    else
        warn "User $ADMIN_USER already exists"
    fi
}

# ===================== STEP 3: FIREWALL =====================
setup_firewall() {
    info "=== Step 3: Configure Firewall ==="

    ufw --force reset 2>/dev/null || true
    ufw default deny incoming
    ufw default deny outgoing
    ufw allow in on lo
    ufw allow out on lo
    ufw allow out 53
    ufw allow out 123
    ufw allow out 80/tcp
    ufw allow out 443/tcp
    ufw allow from "$NETWORK_CIDR" port 22 proto tcp
    ufw allow from "$NETWORK_CIDR" port 80 proto tcp
    ufw allow from "$NETWORK_CIDR" port 443 proto tcp
    echo "y" | ufw enable

    # Create swap
    if [[ ! -f /swapfile ]]; then
        fallocate -l 4G /swapfile
        chmod 600 /swapfile
        mkswap /swapfile
        swapon /swapfile
        echo "/swapfile none swap sw 0 0" >> /etc/fstab
        success "Swap created (4GB)"
    fi

    success "Firewall configured"
}

# ===================== STEP 4: DETECT YUBIKEYS =====================
detect_yubikeys() {
    info "=== Step 4: Detect YubiKeys ==="

    # Start pcscd
    systemctl enable pcscd
    systemctl start pcscd
    sleep 2

    # Fix polkit for pcscd access
    mkdir -p /etc/polkit-1/rules.d
    cat > /etc/polkit-1/rules.d/45-access-pcsc.rules << 'POLKIT'
polkit.addRule(function(action, subject) {
    if (action.id == "org.debian.pcsc-lite.access_pcsc") {
        return polkit.Result.YES;
    }
});
POLKIT
    systemctl restart polkit
    systemctl restart pcscd
    sleep 2

    # List YubiKeys
    info "Connected YubiKeys:"
    sudo -u "$ADMIN_USER" ykman list 2>/dev/null || ykman list

    # Auto-detect serials if not set
    SERIALS=$(ykman list 2>/dev/null | grep -oP 'Serial: \K\d+' || true)
    if [[ -z "$YK1_SERIAL" ]]; then
        YK1_SERIAL=$(echo "$SERIALS" | head -1)
    fi
    if [[ -z "$YK2_SERIAL" ]]; then
        YK2_SERIAL=$(echo "$SERIALS" | tail -1)
    fi

    info "YK1 (Root CA):     $YK1_SERIAL"
    info "YK2 (Intermediate): $YK2_SERIAL"

    # Generate PINs if not set
    if [[ -z "$YK1_PIN" ]]; then
        YK1_PIN=$(tr -dc 'A-Za-z0-9' < /dev/urandom | head -c 8)
        YK1_PUK=$(tr -dc 'A-Za-z0-9' < /dev/urandom | head -c 8)
    fi
    if [[ -z "$YK2_PIN" ]]; then
        YK2_PIN=$(tr -dc 'A-Za-z0-9' < /dev/urandom | head -c 8)
        YK2_PUK=$(tr -dc 'A-Za-z0-9' < /dev/urandom | head -c 8)
    fi

    echo ""
    echo "============================================================"
    echo "  YUBIKEY CREDENTIALS - RECORD THESE SECURELY"
    echo "============================================================"
    echo ""
    echo "YK1 (Root CA) - Serial $YK1_SERIAL:"
    echo "  PIN: $YK1_PIN"
    echo "  PUK: $YK1_PUK"
    echo ""
    echo "YK2 (Intermediate) - Serial $YK2_SERIAL:"
    echo "  PIN: $YK2_PIN"
    echo "  PUK: $YK2_PUK"
    echo ""
    echo "============================================================"
    echo ""
}

# ===================== STEP 5: CONFIGURE YUBIKEYS =====================
configure_yubikeys() {
    info "=== Step 5: Configure YubiKeys ==="
    DEFAULT_MGMT="010203040506070801020304050607080102030405060708"

    # --- YK1: Root CA ---
    info "Configuring YK1 (Root CA)..."

    echo "y" | ykman -d "$YK1_SERIAL" piv reset
    ykman -d "$YK1_SERIAL" piv keys generate --algorithm ECCP384 \
        -m "$DEFAULT_MGMT" 9c /tmp/yk1-root-pub.pem
    ykman -d "$YK1_SERIAL" piv access change-pin -P 123456 --new-pin "$YK1_PIN"
    ykman -d "$YK1_SERIAL" piv access change-puk -p 12345678 --new-puk "$YK1_PUK"
    echo "" | ykman -d "$YK1_SERIAL" piv certificates generate \
        -P "$YK1_PIN" -m "$DEFAULT_MGMT" \
        --subject "CN=$ROOT_CA_CN,O=$CA_ORG,C=$CA_COUNTRY" \
        --valid-days "$ROOT_VALID_DAYS" --hash-algorithm SHA384 \
        9c /tmp/yk1-root-pub.pem

    ykman -d "$YK1_SERIAL" piv certificates export 9c /tmp/root-ca-temp.crt
    success "YK1 configured"

    # --- YK2: Intermediate CA ---
    info "Configuring YK2 (Intermediate CA)..."

    echo "y" | ykman -d "$YK2_SERIAL" piv reset
    ykman -d "$YK2_SERIAL" piv keys generate --algorithm ECCP384 \
        -m "$DEFAULT_MGMT" 9c /tmp/yk2-int-pub.pem
    ykman -d "$YK2_SERIAL" piv access change-pin -P 123456 --new-pin "$YK2_PIN"
    ykman -d "$YK2_SERIAL" piv access change-puk -p 12345678 --new-puk "$YK2_PUK"
    echo "" | ykman -d "$YK2_SERIAL" piv certificates request \
        -P "$YK2_PIN" \
        --subject "CN=$INT_CA_CN,O=$CA_ORG,C=$CA_COUNTRY" \
        9c /tmp/yk2-int-pub.pem /tmp/yk2-intermediate.csr

    success "YK2 configured"
}

# ===================== STEP 6: GENERATE CA CERTS =====================
generate_ca_certs() {
    info "=== Step 6: Generate CA Certificates ==="

    # Create directories
    mkdir -p /etc/ssl/ca/{root,intermediate,issued,crl}
    mkdir -p /var/lib/certauth/tmp
    mkdir -p /var/log/certauth
    chown -R "$ADMIN_USER:$ADMIN_USER" /etc/ssl/ca /var/lib/certauth /var/log/certauth
    chmod -R 700 /etc/ssl/ca /var/lib/certauth /var/log/certauth

    # Generate Root CA cert with proper extensions
    info "Generating Root CA certificate..."
    python3 << ROOTPY
import datetime, subprocess, base64
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.x509.oid import NameOID

with open("/tmp/yk1-root-pub.pem", "rb") as f:
    root_pub = serialization.load_pem_public_key(f.read())

subj = x509.Name([
    x509.NameAttribute(NameOID.COUNTRY_NAME, "$CA_COUNTRY"),
    x509.NameAttribute(NameOID.ORGANIZATION_NAME, "$CA_ORG"),
    x509.NameAttribute(NameOID.COMMON_NAME, "$ROOT_CA_CN"),
])

builder = (x509.CertificateBuilder()
    .subject_name(subj).issuer_name(subj)
    .public_key(root_pub).serial_number(x509.random_serial_number())
    .not_valid_before(datetime.datetime.now(datetime.timezone.utc))
    .not_valid_after(datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=$ROOT_VALID_DAYS))
    .add_extension(x509.BasicConstraints(ca=True, path_length=1), critical=True)
    .add_extension(x509.KeyUsage(digital_signature=True, key_cert_sign=True, crl_sign=True,
        key_encipherment=False, content_commitment=False, data_encipherment=False,
        key_agreement=False, encipher_only=False, decipher_only=False), critical=True)
    .add_extension(x509.SubjectKeyIdentifier.from_public_key(root_pub), critical=False)
    .add_extension(x509.AuthorityKeyIdentifier.from_issuer_public_key(root_pub), critical=False))

tmp = ec.generate_private_key(ec.SECP384R1())
temp = builder.sign(tmp, hashes.SHA384())
td = temp.public_bytes(serialization.Encoding.DER)

# Parse TBS
o = 1
if td[o] & 0x80: n = td[o] & 0x7f; o += 1 + n
else: o += 1
tbs_start = o
o += 1
if td[o] & 0x80: n = td[o] & 0x7f; tl = int.from_bytes(td[o+1:o+1+n], "big"); o += 1 + n
else: tl = td[o]; o += 1
tbs_end = o + tl
tbs_full = td[tbs_start:tbs_end]
alg_start = tbs_end
o2 = alg_start + 1
if td[o2] & 0x80: n = td[o2] & 0x7f; al = int.from_bytes(td[o2+1:o2+1+n], "big"); o2 += 1 + n
else: al = td[o2]; o2 += 1
alg_full = td[alg_start:o2+al]

# Sign with YubiKey
with open("/tmp/tbs.der", "wb") as f: f.write(tbs_full)
r = subprocess.run(["pkcs11-tool", "--login", "--pin", "$YK1_PIN",
    "--sign", "--mechanism", "ECDSA-SHA384", "--label", "SIGN key",
    "--input-file", "/tmp/tbs.der", "--output-file", "/tmp/sig.bin"],
    capture_output=True, text=True)
if r.returncode != 0:
    print(f"Sign failed: {r.stderr}")
    exit(1)

with open("/tmp/sig.bin", "rb") as f: raw = f.read()
rb = raw[:48].lstrip(b"\x00") or b"\x00"
sb = raw[48:].lstrip(b"\x00") or b"\x00"
if rb[0] & 0x80: rb = b"\x00" + rb
if sb[0] & 0x80: sb = b"\x00" + sb
r_der = b"\x02" + bytes([len(rb)]) + rb
s_der = b"\x02" + bytes([len(sb)]) + sb
seq = b"\x30" + bytes([len(r_der+s_der)]) + r_der + s_der
bs = b"\x00" + seq
new_sig = b"\x03" + bytes([len(bs)]) + bs

content = tbs_full + alg_full + new_sig
cl = len(content)
final = b"\x30\x82" + bytes([cl>>8, cl&0xff]) + content

with open("/tmp/root.der", "wb") as f: f.write(final)
r = subprocess.run(["openssl", "x509", "-inform", "DER", "-outform", "PEM",
    "-in", "/tmp/root.der", "-out", "/tmp/root.pem"], capture_output=True, text=True)
if r.returncode != 0:
    print(f"DER->PEM failed: {r.stderr}")
    exit(1)

with open("/tmp/root.pem") as f: pem = f.read()
with open("/etc/ssl/ca/root/root-ca.crt", "w") as f: f.write(pem)

# Verify
r = subprocess.run(["openssl", "verify", "-CAfile", "/etc/ssl/ca/root/root-ca.crt",
    "/etc/ssl/ca/root/root-ca.crt"], capture_output=True, text=True)
print(f"Root CA self-verify: {r.stdout.strip() or r.stderr.strip()}")

# Import to YK1
subprocess.run(["echo", "", "|", "ykman", "-d", "$YK1_SERIAL", "piv", "certificates",
    "import", "9c", "/etc/ssl/ca/root/root-ca.crt"], shell=True)
print("Root CA imported to YK1")
ROOTPY

    # Generate Intermediate CA cert signed by Root
    info "Generating Intermediate CA certificate..."
    python3 << INTPY
import datetime, subprocess, base64
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.x509.oid import NameOID

with open("/tmp/yk1-root-pub.pem", "rb") as f:
    root_pub = serialization.load_pem_public_key(f.read())
with open("/tmp/yk2-int-pub.pem", "rb") as f:
    int_pub = serialization.load_pem_public_key(f.read())

root_subj = x509.Name([
    x509.NameAttribute(NameOID.COUNTRY_NAME, "$CA_COUNTRY"),
    x509.NameAttribute(NameOID.ORGANIZATION_NAME, "$CA_ORG"),
    x509.NameAttribute(NameOID.COMMON_NAME, "$ROOT_CA_CN"),
])
int_subj = x509.Name([
    x509.NameAttribute(NameOID.COUNTRY_NAME, "$CA_COUNTRY"),
    x509.NameAttribute(NameOID.ORGANIZATION_NAME, "$CA_ORG"),
    x509.NameAttribute(NameOID.COMMON_NAME, "$INT_CA_CN"),
])

builder = (x509.CertificateBuilder()
    .subject_name(int_subj).issuer_name(root_subj)
    .public_key(int_pub).serial_number(x509.random_serial_number())
    .not_valid_before(datetime.datetime.now(datetime.timezone.utc))
    .not_valid_after(datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=$INT_VALID_DAYS))
    .add_extension(x509.BasicConstraints(ca=True, path_length=0), critical=True)
    .add_extension(x509.KeyUsage(digital_signature=True, key_cert_sign=True, crl_sign=True,
        key_encipherment=False, content_commitment=False, data_encipherment=False,
        key_agreement=False, encipher_only=False, decipher_only=False), critical=True)
    .add_extension(x509.SubjectKeyIdentifier.from_public_key(int_pub), critical=False)
    .add_extension(x509.AuthorityKeyIdentifier.from_issuer_public_key(root_pub), critical=False))

tmp = ec.generate_private_key(ec.SECP384R1())
temp = builder.sign(tmp, hashes.SHA384())
td = temp.public_bytes(serialization.Encoding.DER)

o = 1
if td[o] & 0x80: n = td[o] & 0x7f; o += 1 + n
else: o += 1
tbs_start = o
o += 1
if td[o] & 0x80: n = td[o] & 0x7f; tl = int.from_bytes(td[o+1:o+1+n], "big"); o += 1 + n
else: tl = td[o]; o += 1
tbs_end = o + tl
tbs_full = td[tbs_start:tbs_end]
alg_start = tbs_end
o2 = alg_start + 1
if td[o2] & 0x80: n = td[o2] & 0x7f; al = int.from_bytes(td[o2+1:o2+1+n], "big"); o2 += 1 + n
else: al = td[o2]; o2 += 1
alg_full = td[alg_start:o2+al]

with open("/tmp/tbs.der", "wb") as f: f.write(tbs_full)
r = subprocess.run(["pkcs11-tool", "--login", "--pin", "$YK1_PIN",
    "--sign", "--mechanism", "ECDSA-SHA384", "--label", "SIGN key",
    "--input-file", "/tmp/tbs.der", "--output-file", "/tmp/sig.bin"],
    capture_output=True, text=True)
if r.returncode != 0:
    print(f"Sign failed: {r.stderr}")
    exit(1)

with open("/tmp/sig.bin", "rb") as f: raw = f.read()
rb = raw[:48].lstrip(b"\x00") or b"\x00"
sb = raw[48:].lstrip(b"\x00") or b"\x00"
if rb[0] & 0x80: rb = b"\x00" + rb
if sb[0] & 0x80: sb = b"\x00" + sb
r_der = b"\x02" + bytes([len(rb)]) + rb
s_der = b"\x02" + bytes([len(sb)]) + sb
seq = b"\x30" + bytes([len(r_der+s_der)]) + r_der + s_der
bs = b"\x00" + seq
new_sig = b"\x03" + bytes([len(bs)]) + bs

content = tbs_full + alg_full + new_sig
cl = len(content)
final = b"\x30\x82" + bytes([cl>>8, cl&0xff]) + content

with open("/tmp/int.der", "wb") as f: f.write(final)
r = subprocess.run(["openssl", "x509", "-inform", "DER", "-outform", "PEM",
    "-in", "/tmp/int.der", "-out", "/tmp/int.pem"], capture_output=True, text=True)
if r.returncode != 0:
    print(f"DER->PEM failed: {r.stderr}")
    exit(1)

with open("/tmp/int.pem") as f: pem = f.read()
with open("/etc/ssl/ca/intermediate/intermediate-ca.crt", "w") as f: f.write(pem)

# Verify chain
r = subprocess.run(["openssl", "verify", "-CAfile", "/etc/ssl/ca/root/root-ca.crt",
    "/etc/ssl/ca/intermediate/intermediate-ca.crt"], capture_output=True, text=True)
print(f"Chain verify: {r.stdout.strip() or r.stderr.strip()}")

# Import to YK2
subprocess.run(["echo", "", "|", "ykman", "-d", "$YK2_SERIAL", "piv", "certificates",
    "import", "9c", "/etc/ssl/ca/intermediate/intermediate-ca.crt"], shell=True)
print("Intermediate CA imported to YK2")
INTPY

    # Create CA chain file
    cat /etc/ssl/ca/intermediate/intermediate-ca.crt /etc/ssl/ca/root/root-ca.crt \
        > /etc/ssl/ca/ca-chain.crt

    success "CA certificates generated and verified"
}

# ===================== STEP 7: INSTALL API =====================
install_api() {
    info "=== Step 7: Install Certificate API ==="

    # Create API directory
    mkdir -p /opt/certauth/{api/templates,api/static,logs}
    chown -R "$ADMIN_USER:$ADMIN_USER" /opt/certauth

    # Deploy API from the repository (single source of truth: the api/
    # directory next to this script). Never inline a second copy here.
    SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
    for f in config.py models.py auth.py signing.py main.py; do
        if [ ! -f "$SCRIPT_DIR/api/$f" ]; then
            error "api/$f not found next to setup-certauth.sh (run from the repository checkout)"
        fi
    done
    cp "$SCRIPT_DIR/api/config.py" "$SCRIPT_DIR/api/models.py" "$SCRIPT_DIR/api/auth.py" \
        "$SCRIPT_DIR/api/signing.py" "$SCRIPT_DIR/api/main.py" /opt/certauth/api/
    cp -r "$SCRIPT_DIR/api/templates/." /opt/certauth/api/templates/
    mkdir -p "$SCRIPT_DIR/api/static"
    cp -r "$SCRIPT_DIR/api/static/." /opt/certauth/api/static/

    # Runtime configuration. The API fails closed on startup if any variable
    # is missing, so this file must be complete and private.
    JWT_SECRET_GENERATED="$(python3 -c 'import secrets; print(secrets.token_hex(32))')"
    install -d -m 750 -o "$ADMIN_USER" -g "$ADMIN_USER" /etc/certauth
    cat > /etc/certauth/certauth.env << ENV_EOF
YK_ROOT_SERIAL=$YK1_SERIAL
YK_INT_SERIAL=$YK2_SERIAL
YK_ROOT_PIN=$YK1_PIN
YK_INT_PIN=$YK2_PIN
JWT_SECRET=$JWT_SECRET_GENERATED
ADMIN_USERNAME=$ADMIN_USER
ADMIN_PASSWORD=$ADMIN_PASS
ENV_EOF
    chown "$ADMIN_USER:$ADMIN_USER" /etc/certauth/certauth.env
    chmod 600 /etc/certauth/certauth.env

    # Initialize database (loads the private env file as the API user)
    sudo -u "$ADMIN_USER" bash -c 'set -a; . /etc/certauth/certauth.env; set +a; cd /opt/certauth/api; python3 -c "from models import init_db; init_db()"'

    success "API installed"
}

# ===================== STEP 8: SERVICES =====================
setup_services() {
    info "=== Step 8: Configure Services ==="

    # Caddy
    cat > /etc/caddy/Caddyfile << 'CADDYEOF'
:80 {
    encode gzip
    reverse_proxy 127.0.0.1:8000 {
        header_up Host {host}
        header_up X-Real-IP {remote}
    }
}
CADDYEOF

    # Systemd service for API
    cat > /etc/systemd/system/certauth-api.service << SVCCEOF
[Unit]
Description=CertAuth Certificate Management API
After=network.target pcscd.service

[Service]
Type=simple
User=$ADMIN_USER
Group=$ADMIN_USER
WorkingDirectory=/opt/certauth/api
Environment=PYTHONUNBUFFERED=1
EnvironmentFile=/etc/certauth/certauth.env
ExecStart=/usr/bin/python3 -m uvicorn main:app --host 127.0.0.1 --port 8000
Restart=on-failure
RestartSec=5

[Install]
WantedBy=multi-user.target
SVCCEOF

    systemctl daemon-reload
    systemctl enable caddy certauth-api pcscd
    systemctl start caddy certauth-api

    success "Services configured and started"
}

# ===================== STEP 9: VERIFY =====================
verify_setup() {
    info "=== Step 9: Verification ==="

    sleep 3

    # Check services
    for svc in certauth-api caddy pcscd; do
        STATUS=$(systemctl is-active "$svc")
        if [[ "$STATUS" == "active" ]]; then
            success "$svc: active"
        else
            warn "$svc: $STATUS"
        fi
    done

    # Test API
    HEALTH=$(curl -s http://localhost:8000/api/health 2>/dev/null || echo "FAILED")
    if echo "$HEALTH" | grep -q "ok"; then
        success "API health check: OK"
    else
        warn "API health check: $HEALTH"
    fi

    # Verify CA chain
    openssl verify -CAfile /etc/ssl/ca/root/root-ca.crt \
        /etc/ssl/ca/intermediate/intermediate-ca.crt 2>&1 | while read line; do
        if echo "$line" | grep -q "OK"; then
            success "CA chain verification: OK"
        else
            warn "CA chain: $line"
        fi
    done

    echo ""
    echo "============================================================"
    echo "  SETUP COMPLETE"
    echo "============================================================"
    echo ""
    echo "  Access the Key Vault UI at:"
    echo "    http://$(hostname -I | awk '{print $1}')/"
    echo ""
    echo "  Admin login:"
    echo "    Username: $ADMIN_USER"
    echo "    Password: $ADMIN_PASS"
    echo ""
    echo "  API endpoints:"
    echo "    POST /api/token          - Login (returns JWT)"
    echo "    GET  /api/domains        - List domains"
    echo "    POST /api/domains        - Register domain"
    echo "    GET  /api/certs          - List certificates"
    echo "    POST /api/certs/request  - Request certificate"
    echo "    POST /api/certs/{id}/sign - Sign certificate (requires YK2)"
    echo "    GET  /api/certs/{id}/download - Download cert"
    echo "    GET  /api/certs/{id}/key   - Download private key"
    echo "    GET  /api/ca-chain       - Download CA chain"
    echo ""
    echo "  IMPORTANT: Change the admin password after first login!"
    echo "============================================================"
}

# ===================== MAIN =====================
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

echo ""
echo "============================================================"
echo "  CertAuth Setup"
echo "============================================================"
echo ""

check_prerequisites
setup_system
setup_user
setup_firewall
detect_yubikeys
configure_yubikeys
generate_ca_certs
install_api
setup_services
verify_setup

echo ""
success "All done!"