Configure Apache Airflow with SSL certificates and NGINX reverse proxy for secure production deployment

#!/usr/bin/env bash

set -euo pipefail

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Variables
DOMAIN=""
AIRFLOW_USER="${AIRFLOW_USER:-airflow}"
AIRFLOW_HOME="${AIRFLOW_HOME:-/opt/airflow}"

# Usage function
usage() {
    echo "Usage: $0 -d DOMAIN [-u AIRFLOW_USER] [-h AIRFLOW_HOME]"
    echo "  -d DOMAIN         Domain name for SSL certificate (e.g., airflow.example.com)"
    echo "  -u AIRFLOW_USER   Airflow system user (default: airflow)"
    echo "  -h AIRFLOW_HOME   Airflow home directory (default: /opt/airflow)"
    exit 1
}

# Parse arguments
while getopts "d:u:h:" opt; do
    case $opt in
        d) DOMAIN="$OPTARG" ;;
        u) AIRFLOW_USER="$OPTARG" ;;
        h) AIRFLOW_HOME="$OPTARG" ;;
        *) usage ;;
    esac
done

if [[ -z "$DOMAIN" ]]; then
    echo -e "${RED}Error: Domain name is required${NC}"
    usage
fi

# Cleanup function
cleanup() {
    echo -e "${RED}Installation failed. Cleaning up...${NC}"
    systemctl stop nginx 2>/dev/null || true
    rm -f "/etc/nginx/sites-available/airflow" "/etc/nginx/conf.d/airflow.conf" 2>/dev/null || true
}

trap cleanup ERR

# Log function
log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] $1${NC}"
}

warn() {
    echo -e "${YELLOW}[$(date +'%Y-%m-%d %H:%M:%S')] WARNING: $1${NC}"
}

error() {
    echo -e "${RED}[$(date +'%Y-%m-%d %H:%M:%S')] ERROR: $1${NC}"
    exit 1
}

# Check if running as root
if [[ $EUID -ne 0 ]]; then
    error "This script must be run as root or with sudo"
fi

log "[1/8] Detecting operating system..."
if [[ ! -f /etc/os-release ]]; then
    error "/etc/os-release not found. Cannot detect distribution."
fi

. /etc/os-release
case "$ID" in
    ubuntu|debian)
        PKG_MGR="apt"
        PKG_INSTALL="apt install -y"
        PKG_UPDATE="apt update"
        NGINX_SITES_DIR="/etc/nginx/sites-available"
        NGINX_ENABLED_DIR="/etc/nginx/sites-enabled"
        CERTBOT_PKG="python3-certbot-nginx"
        ;;
    almalinux|rocky|centos|rhel|ol)
        PKG_MGR="dnf"
        PKG_INSTALL="dnf install -y"
        PKG_UPDATE="dnf update -y"
        NGINX_SITES_DIR="/etc/nginx/conf.d"
        NGINX_ENABLED_DIR=""
        CERTBOT_PKG="python3-certbot-nginx"
        ;;
    fedora)
        PKG_MGR="dnf"
        PKG_INSTALL="dnf install -y"
        PKG_UPDATE="dnf update -y"
        NGINX_SITES_DIR="/etc/nginx/conf.d"
        NGINX_ENABLED_DIR=""
        CERTBOT_PKG="python3-certbot-nginx"
        ;;
    amzn)
        PKG_MGR="yum"
        PKG_INSTALL="yum install -y"
        PKG_UPDATE="yum update -y"
        NGINX_SITES_DIR="/etc/nginx/conf.d"
        NGINX_ENABLED_DIR=""
        CERTBOT_PKG="python3-certbot-nginx"
        ;;
    *)
        error "Unsupported distribution: $ID"
        ;;
esac

log "Detected: $PRETTY_NAME"

log "[2/8] Installing NGINX and Certbot..."
$PKG_UPDATE
$PKG_INSTALL nginx certbot $CERTBOT_PKG

# Enable and start NGINX
systemctl enable nginx
systemctl start nginx

# Configure firewall for RHEL-based systems
if [[ "$ID" =~ ^(almalinux|rocky|centos|rhel|ol|fedora)$ ]]; then
    if command -v firewall-cmd >/dev/null 2>&1 && systemctl is-active --quiet firewalld; then
        log "Configuring firewall..."
        firewall-cmd --permanent --add-service=http
        firewall-cmd --permanent --add-service=https
        firewall-cmd --reload
    fi
fi

log "[3/8] Checking Airflow installation..."
if [[ ! -d "$AIRFLOW_HOME" ]]; then
    warn "Airflow home directory $AIRFLOW_HOME does not exist"
fi

if ! id "$AIRFLOW_USER" >/dev/null 2>&1; then
    warn "Airflow user $AIRFLOW_USER does not exist"
fi

log "[4/8] Generating SSL certificates..."
echo -e "${YELLOW}Make sure $DOMAIN points to this server's IP address before continuing.${NC}"
read -p "Press Enter to continue with certificate generation..."

certbot certonly --nginx -d "$DOMAIN" --non-interactive --agree-tos --email admin@"$DOMAIN"

log "[5/8] Configuring Airflow for reverse proxy..."
AIRFLOW_CONFIG="$AIRFLOW_HOME/airflow.cfg"
if [[ -f "$AIRFLOW_CONFIG" ]]; then
    # Backup original config
    cp "$AIRFLOW_CONFIG" "$AIRFLOW_CONFIG.backup.$(date +%s)"
    
    # Update configuration
    sed -i "s/^web_server_host = .*/web_server_host = 127.0.0.1/" "$AIRFLOW_CONFIG"
    sed -i "s/^web_server_port = .*/web_server_port = 8080/" "$AIRFLOW_CONFIG"
    sed -i "s|^base_url = .*|base_url = https://$DOMAIN|" "$AIRFLOW_CONFIG"
    sed -i "s/^enable_proxy_fix = .*/enable_proxy_fix = True/" "$AIRFLOW_CONFIG"
    
    chown "$AIRFLOW_USER":"$AIRFLOW_USER" "$AIRFLOW_CONFIG"
    chmod 640 "$AIRFLOW_CONFIG"
else
    warn "Airflow configuration file not found at $AIRFLOW_CONFIG"
fi

log "[6/8] Creating NGINX configuration..."
NGINX_CONFIG_FILE="$NGINX_SITES_DIR/airflow.conf"

cat > "$NGINX_CONFIG_FILE" << EOF
upstream airflow_webserver {
    server 127.0.0.1:8080 fail_timeout=0;
}

server {
    listen 80;
    server_name $DOMAIN;
    return 301 https://\$server_name\$request_uri;
}

server {
    listen 443 ssl http2;
    server_name $DOMAIN;
    
    ssl_certificate /etc/letsencrypt/live/$DOMAIN/fullchain.pem;
    ssl_certificate_key /etc/letsencrypt/live/$DOMAIN/privkey.pem;
    ssl_protocols TLSv1.2 TLSv1.3;
    ssl_ciphers ECDHE-RSA-AES256-GCM-SHA512:DHE-RSA-AES256-GCM-SHA512:ECDHE-RSA-AES256-GCM-SHA384:DHE-RSA-AES256-GCM-SHA384:ECDHE-RSA-AES256-SHA384;
    ssl_prefer_server_ciphers off;
    ssl_session_cache shared:SSL:10m;
    ssl_session_timeout 10m;
    
    add_header Strict-Transport-Security "max-age=63072000; includeSubDomains; preload";
    add_header X-Content-Type-Options nosniff;
    add_header X-Frame-Options DENY;
    add_header X-XSS-Protection "1; mode=block";
    add_header Referrer-Policy "strict-origin-when-cross-origin";
    add_header Content-Security-Policy "default-src 'self'; script-src 'self' 'unsafe-inline' 'unsafe-eval'; style-src 'self' 'unsafe-inline'; img-src 'self' data:; font-src 'self'; connect-src 'self'";
    
    client_max_body_size 100M;
    
    location / {
        proxy_pass http://airflow_webserver;
        proxy_set_header Host \$http_host;
        proxy_set_header X-Real-IP \$remote_addr;
        proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for;
        proxy_set_header X-Forwarded-Proto \$scheme;
        proxy_redirect off;
        
        proxy_http_version 1.1;
        proxy_set_header Upgrade \$http_upgrade;
        proxy_set_header Connection "upgrade";
        
        proxy_connect_timeout 60s;
        proxy_send_timeout 60s;
        proxy_read_timeout 60s;
    }
    
    location /health {
        proxy_pass http://airflow_webserver/health;
        proxy_set_header Host \$http_host;
        access_log off;
    }
    
    location ~* \.(css|js|png|jpg|jpeg|gif|ico|svg)\$ {
        proxy_pass http://airflow_webserver;
        proxy_set_header Host \$http_host;
        expires 1M;
        add_header Cache-Control "public, immutable";
    }
}
EOF

chmod 644 "$NGINX_CONFIG_FILE"

log "[7/8] Enabling NGINX configuration..."
if [[ -n "$NGINX_ENABLED_DIR" ]]; then
    # Debian/Ubuntu style
    ln -sf "$NGINX_CONFIG_FILE" "$NGINX_ENABLED_DIR/airflow.conf"
    rm -f "$NGINX_ENABLED_DIR/default"
fi

# Test and reload NGINX
nginx -t
systemctl reload nginx

log "[8/8] Setting up SSL certificate renewal..."
# Add renewal cron job
(crontab -l 2>/dev/null || echo "") | grep -v "certbot renew" | \
{ cat; echo "0 12 * * * /usr/bin/certbot renew --quiet --deploy-hook 'systemctl reload nginx'"; } | crontab -

log "Installation completed successfully!"
echo
echo -e "${GREEN}Next steps:${NC}"
echo "1. Start your Airflow webserver: sudo -u $AIRFLOW_USER airflow webserver"
echo "2. Access your Airflow instance at: https://$DOMAIN"
echo "3. SSL certificates will auto-renew via cron"
echo
echo -e "${YELLOW}Note: Ensure Airflow is running on 127.0.0.1:8080${NC}"