#!/bin/bash
# Enable Kubernetes Audit Logging in K3s
# Citadel requires comprehensive audit logging for compliance
# This script configures K3s to use our audit policy

set -euo pipefail

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

log() { echo "[$(date +'%Y-%m-%d %H:%M:%S')] $*"; }
error() { echo "[ERROR] $*" >&2; }

log "Enabling Kubernetes audit logging..."

# Check if running as root (needed to modify K3s config)
if [[ $EUID -ne 0 ]]; then
    error "This script must be run as root to modify K3s configuration"
    exit 1
fi

# Create audit log directory
AUDIT_LOG_DIR="/var/log/kubernetes/audit"
mkdir -p "${AUDIT_LOG_DIR}"
chmod 700 "${AUDIT_LOG_DIR}"
log "Created audit log directory: ${AUDIT_LOG_DIR}"

# Copy audit policy to K3s config location
AUDIT_POLICY_DEST="/etc/rancher/k3s/audit-policy.yaml"
cp "${REPO_DIR}/config/audit-policy.yaml" "${AUDIT_POLICY_DEST}"
chmod 600 "${AUDIT_POLICY_DEST}"
log "Copied audit policy to: ${AUDIT_POLICY_DEST}"

# Create K3s configuration file with audit settings
K3S_CONFIG="/etc/rancher/k3s/config.yaml"

# Backup existing config if it exists
if [[ -f "${K3S_CONFIG}" ]]; then
    cp "${K3S_CONFIG}" "${K3S_CONFIG}.backup.$(date +%Y%m%d%H%M%S)"
    log "Backed up existing K3s config"
fi

# Check if audit settings already exist
if grep -q "kube-apiserver-arg:" "${K3S_CONFIG}" 2>/dev/null; then
    log "K3s config already has kube-apiserver-arg. Adding audit settings..."

    # Check if audit-policy-file is already configured
    if grep -q "audit-policy-file" "${K3S_CONFIG}"; then
        log "Audit logging already configured. Skipping..."
    else
        # Append audit settings to existing kube-apiserver-arg section
        cat >> "${K3S_CONFIG}" << EOF
  - "audit-policy-file=${AUDIT_POLICY_DEST}"
  - "audit-log-path=${AUDIT_LOG_DIR}/audit.log"
  - "audit-log-maxage=30"
  - "audit-log-maxbackup=10"
  - "audit-log-maxsize=100"
EOF
    fi
else
    # Create or append new config with audit settings
    cat >> "${K3S_CONFIG}" << EOF

# Audit logging configuration for Citadel compliance
kube-apiserver-arg:
  - "audit-policy-file=${AUDIT_POLICY_DEST}"
  - "audit-log-path=${AUDIT_LOG_DIR}/audit.log"
  - "audit-log-maxage=30"
  - "audit-log-maxbackup=10"
  - "audit-log-maxsize=100"
EOF
fi

log "Updated K3s configuration with audit settings"

# Create log rotation configuration
cat > /etc/logrotate.d/kubernetes-audit << EOF
${AUDIT_LOG_DIR}/audit.log {
    daily
    rotate 30
    compress
    delaycompress
    missingok
    notifempty
    create 0600 root root
    sharedscripts
    postrotate
        /bin/kill -HUP \$(cat /var/run/k3s.pid 2>/dev/null) 2>/dev/null || true
    endscript
}
EOF
log "Created log rotation configuration"

# Create Fluent Bit config for log forwarding (optional)
cat > "${REPO_DIR}/config/fluent-bit-audit.yaml" << 'EOF'
# Fluent Bit configuration for forwarding audit logs
apiVersion: v1
kind: ConfigMap
metadata:
  name: fluent-bit-audit-config
  namespace: kube-system
data:
  fluent-bit.conf: |
    [SERVICE]
        Flush         5
        Daemon        Off
        Log_Level     info
        Parsers_File  parsers.conf

    [INPUT]
        Name          tail
        Path          /var/log/kubernetes/audit/audit.log
        Parser        json
        Tag           kube.audit
        Refresh_Interval  5
        Mem_Buf_Limit     50MB
        Skip_Long_Lines   On

    [FILTER]
        Name          modify
        Match         kube.audit
        Add           cluster citadel-local
        Add           log_type audit

    [OUTPUT]
        Name          stdout
        Match         *
        Format        json

    # Uncomment to forward to external SIEM
    # [OUTPUT]
    #     Name          http
    #     Match         *
    #     Host          siem.example.com
    #     Port          8080
    #     URI           /api/logs
    #     Format        json

  parsers.conf: |
    [PARSER]
        Name        json
        Format      json
        Time_Key    requestReceivedTimestamp
        Time_Format %Y-%m-%dT%H:%M:%S.%L%z
EOF
log "Created Fluent Bit configuration for log forwarding"

log ""
log "=========================================="
log "Audit Logging Configuration Complete"
log "=========================================="
log ""
log "IMPORTANT: K3s must be restarted for changes to take effect!"
log ""
log "To restart K3s:"
log "  sudo systemctl restart k3s"
log ""
log "After restart, verify audit logging:"
log "  tail -f ${AUDIT_LOG_DIR}/audit.log"
log ""
log "Audit log location: ${AUDIT_LOG_DIR}/audit.log"
log "Audit policy: ${AUDIT_POLICY_DEST}"
log ""
log "What is being logged:"
log "  - Secret/ConfigMap modifications"
log "  - Pod exec/attach/port-forward"
log "  - RBAC changes (roles, bindings)"
log "  - Network policy changes"
log "  - Namespace changes"
log "  - Node lifecycle events"
log "  - Failed requests"
log ""
log "Log rotation: Daily, 30 days retention"
log ""
log "To deploy Fluent Bit for log forwarding:"
log "  kubectl apply -f ${REPO_DIR}/config/fluent-bit-audit.yaml"
log "=========================================="

# Optional: Restart K3s now
read -p "Restart K3s now to apply changes? [y/N] " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
    log "Restarting K3s..."
    systemctl restart k3s
    log "Waiting for K3s to be ready..."
    sleep 10
    kubectl wait --for=condition=Ready nodes --all --timeout=120s
    log "K3s restarted successfully!"

    # Verify audit logging is working
    sleep 5
    if [[ -f "${AUDIT_LOG_DIR}/audit.log" ]]; then
        log "✓ Audit log file created"
        log "Sample entries:"
        tail -3 "${AUDIT_LOG_DIR}/audit.log" | jq -r '.verb + " " + .objectRef.resource + "/" + .objectRef.name' 2>/dev/null || tail -3 "${AUDIT_LOG_DIR}/audit.log"
    else
        error "Audit log file not found. Check K3s logs for errors."
    fi
else
    log "Skipping restart. Remember to restart K3s manually."
fi
