#!/bin/bash
# ==============================================================================
# Script Name: ssl_manager.sh
# Description: Automates SSL certificate issuance and updates LiteSpeed's httpd_config.xml.
#              Ensures robust backups, validation, and error-free updates.
# Version: 2.0.4 (Idempotent, Clean Exit Messages)
# Author: Anthony Garces (tony@mightybox.io)
# Date: 2025-03-26
# ==============================================================================
set -euo pipefail

CONF_FILE="/var/www/conf/httpd_config.xml"
DEFAULT_CONF="/var/www/conf/httpd_config.default.xml"
SERVER_ROOT="/var/www"
LOG_DIR="/var/log/mb-ssl"
CERT_DIR="/etc/letsencrypt/live"
PUBLIC_IP=""
DOMAINS=()
EMAIL=""
VERBOSE=0
VHOST_NAME="Jelastic"  # Default vhost name, can be overridden with --vhost parameter
VHOSTS_DIR="$SERVER_ROOT/conf/vhosts"
DEFAULT_DOCROOT="/var/www/webroot/ROOT"

SCRIPT_LOG="${LOG_DIR}/ssl_manager.log"
ERROR_LOG="${LOG_DIR}/ssl_manager-error.log"
DEBUG_LOG="${LOG_DIR}/ssl_manager-debug.log"
BACKUP_FILE="${LOG_DIR}/httpd_config_backup_$(date +%Y%m%d%H%M%S).xml"
SCRIPT_EXIT_STATUS=0

setup_logging() {
    # Create log directory if it doesn't exist
    sudo mkdir -p "$LOG_DIR" || { echo "❌ ERROR: Cannot create log directory '$LOG_DIR'. Check permissions."; exit 1; }
    
    # Set proper permissions
    sudo chown -R "$(whoami)":"$(id -gn)" "$LOG_DIR"
    sudo chmod 755 "$LOG_DIR"
    
    # Create log files with proper permissions
    touch "$SCRIPT_LOG" "$ERROR_LOG" "$DEBUG_LOG"
    chmod 644 "$SCRIPT_LOG" "$ERROR_LOG" "$DEBUG_LOG"
    
    # Add log rotation if log files are too large (>10MB)
    for log_file in "$SCRIPT_LOG" "$ERROR_LOG" "$DEBUG_LOG"; do
        if [ -f "$log_file" ] && [ "$(stat -f%z "$log_file" 2>/dev/null || stat -c%s "$log_file")" -gt 10485760 ]; then
            mv "$log_file" "${log_file}.$(date +%Y%m%d)"
            touch "$log_file"
            chmod 644 "$log_file"
            gzip "${log_file}.$(date +%Y%m%d)"
        fi
    done
}

log() {
    local level="INFO"
    local message="$1"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    
    # Log to main log file
    echo "[$timestamp] [$level] $message" | sudo tee -a "$SCRIPT_LOG"
    
    # Log errors to error log file
    if [[ "$message" == *"ERROR"* ]] || [[ "$message" == *"❌"* ]]; then
        echo "[$timestamp] [$level] $message" | sudo tee -a "$ERROR_LOG"
    fi
}

log_verbose() {
    if [[ "$VERBOSE" -eq 1 ]]; then
        local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
        echo "[$timestamp] [DEBUG] $1" | sudo tee -a "$DEBUG_LOG"
    fi
}

log_error() {
    local message="$1"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[$timestamp] [ERROR] $message" | sudo tee -a "$ERROR_LOG" "$SCRIPT_LOG"
}

log_success() {
    local message="$1"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[$timestamp] [SUCCESS] $message" | sudo tee -a "$SCRIPT_LOG"
}

on_exit() {
    if [[ $SCRIPT_EXIT_STATUS -ne 0 ]]; then
        log_error "Script ended with error. Exit Code: $SCRIPT_EXIT_STATUS"
        log_error "Backup file retained: $BACKUP_FILE"
    else
        sudo rm -f "$BACKUP_FILE"
        log_success "Script completed successfully. Backup cleaned."
    fi
    exit $SCRIPT_EXIT_STATUS
}
trap 'on_exit' EXIT

check_command() {
    local cmd="$1"
    local pkg="$2"
    if ! command -v "$cmd" &>/dev/null; then
        log "Required command '$cmd' not found. Attempting to install package '$pkg'..."
        if command -v dnf &>/dev/null; then
            if sudo dnf install -y "$pkg"; then
                log_success "Successfully installed '$pkg' via dnf"
                return 0
            fi
        fi
        if command -v yum &>/dev/null; then
            if sudo yum install -y "$pkg"; then
                log_success "Successfully installed '$pkg' via yum"
                return 0
            fi
        fi
        log_error "Failed to install '$pkg'"
        SCRIPT_EXIT_STATUS=1; return 1
    fi
}

validate_domain() { [[ "$1" =~ ^([a-zA-Z0-9-]+\.)+[a-zA-Z]{2,}$ ]]; }
validate_ip() { [[ "$1" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]]; }
validate_email() { [[ "$1" =~ ^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}$ ]]; }

create_default_backup() {
    if [[ ! -f "$DEFAULT_CONF" ]]; then
        log "Creating initial backup of '$CONF_FILE' as '$DEFAULT_CONF'..."
        sudo cp -a "$CONF_FILE" "$DEFAULT_CONF" || {
            log_error "FATAL: Failed to create initial backup '$DEFAULT_CONF'"
            SCRIPT_EXIT_STATUS=2; return 1
        }
        log_success "Created initial backup successfully"
    else
        log "Default backup '$DEFAULT_CONF' already exists. Skipping creation."
    fi
}

validate_xml() {
    log_verbose "Validating XML structure of '$1'..."
    sudo xmllint --noout "$1" 2>/dev/null || {
        log_error "Invalid XML structure in '$1'"
        SCRIPT_EXIT_STATUS=1
        return 1
    }
    log_verbose "XML validation passed for '$1'"
}

validate_dns() {
    local resolved_ip
    resolved_ip=$(dig +short "$1" | grep -E '^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$' | head -n1)
    [[ "$resolved_ip" == "$2" ]]
}

validate_http_access() {
    local token
    token=$(openssl rand -hex 16)
    local acme_dir="/var/www/webroot/ROOT/.well-known/acme-challenge"
    sudo mkdir -p "$acme_dir"
    echo "$token" | sudo tee "$acme_dir/test-token" >/dev/null
    [[ "$(curl -s "http://$1/.well-known/acme-challenge/test-token")" == "$token" ]]
}

issue_certificate() {
	if [[ -f "$CERT_DIR/$1/fullchain.pem" ]]; then
		log_success "Certificate already exists for '$1'. Skipping issuance."
		return
	fi
	log "Issuing SSL certificate for domain '$1' with email '$2'..."
	sudo certbot certonly --webroot -w "/var/www/webroot/ROOT" -d "$1" --non-interactive --agree-tos --email "$2" || {
		log_error "Failed to issue certificate for '$1'"
		SCRIPT_EXIT_STATUS=1; return 1
	}
	log_success "Certificate successfully issued for '$1'"
}

issue_certificate_san() {
	local email="$1"; shift
	local primary="$1"; shift
	local others=("$@")

	local args=(certonly --webroot -w "$DEFAULT_DOCROOT" --non-interactive --agree-tos --email "$email")
	for d in "${others[@]}"; do
		args+=( -d "$d" )
	done
	log "Requesting SAN certificate for: $primary ${others[*]}"
	sudo certbot "${args[@]}" || { log_error "SAN issuance failed"; return 1; }
	return 0
}

update_httpd_config() {
	local domain="$1"
	local ip="$2"
	local vhost_name="$domain"   # vhost named after the domain
	local vhost_dir="$VHOSTS_DIR/$domain"
	local vhconf_file="$vhost_dir/vhconf.xml"

	log "Configuring SNI for domain '$domain' on existing HTTPS listener (443)"
	sudo cp -a "$CONF_FILE" "$BACKUP_FILE"

	# 1) Ensure virtualHostList entry exists for this domain
	local vh_exists
	vh_exists=$(xmlstarlet sel -t -v "/httpServerConfig/virtualHostList/virtualHost[name='$vhost_name']/name" "$CONF_FILE" 2>/dev/null || true)
	if [[ -z "$vh_exists" ]]; then
		log "Adding virtualHost '$vhost_name' to httpd_config.xml"
		# Create vhost directory
		sudo mkdir -p "$vhost_dir"
		# Write minimal vhconf.xml with SSL at vhost level (SNI)
		cat <<EOF | sudo tee "$vhconf_file" >/dev/null
<?xml version="1.0" encoding="UTF-8"?>
<virtualHostConfig>
  <docRoot>$VH_ROOT/ROOT/</docRoot>
  <enableGzip>1</enableGzip>
  <vhssl>
    <keyFile>/etc/letsencrypt/live/$domain/privkey.pem</keyFile>
    <certFile>/etc/letsencrypt/live/$domain/fullchain.pem</certFile>
    <certChain>1</certChain>
  </vhssl>
</virtualHostConfig>
EOF
		# Add virtualHost entry
		sudo xmlstarlet ed -L \
			-s "/httpServerConfig/virtualHostList" -t elem -n "virtualHost" \
			"$CONF_FILE" || { log_error "Failed to create virtualHost node"; return 1; }

		sudo xmlstarlet ed -L \
			-s "/httpServerConfig/virtualHostList/virtualHost[last()]" -t elem -n "name" -v "$vhost_name" \
			"$CONF_FILE" || { log_error "Failed to set virtualHost name"; return 1; }

		sudo xmlstarlet ed -L \
			-s "/httpServerConfig/virtualHostList/virtualHost[last()]" -t elem -n "vhRoot" -v "$SERVER_ROOT/webroot/" \
			"$CONF_FILE" || { log_error "Failed to set vhRoot"; return 1; }

		sudo xmlstarlet ed -L \
			-s "/httpServerConfig/virtualHostList/virtualHost[last()]" -t elem -n "configFile" -v "$vhost_dir/vhconf.xml" \
			"$CONF_FILE" || { log_error "Failed to set configFile"; return 1; }

		sudo xmlstarlet ed -L \
			-s "/httpServerConfig/virtualHostList/virtualHost[last()]" -t elem -n "allowSymbolLink" -v "1" \
			"$CONF_FILE" || { log_error "Failed to set allowSymbolLink"; return 1; }

		sudo xmlstarlet ed -L \
			-s "/httpServerConfig/virtualHostList/virtualHost[last()]" -t elem -n "enableScript" -v "1" \
			"$CONF_FILE" || { log_error "Failed to set enableScript"; return 1; }

		sudo xmlstarlet ed -L \
			-s "/httpServerConfig/virtualHostList/virtualHost[last()]" -t elem -n "restrained" -v "1" \
			"$CONF_FILE" || { log_error "Failed to set restrained"; return 1; }
	else
		log "virtualHost '$vhost_name' already exists"
	fi

	# 2) Map domain to this vhost in existing HTTPS listener
	local https_listener_exists
	https_listener_exists=$(xmlstarlet sel -t -v "/httpServerConfig/listenerList/listener[name='HTTPS']/name" "$CONF_FILE" 2>/dev/null || true)
	local https6_listener_exists
	https6_listener_exists=$(xmlstarlet sel -t -v "/httpServerConfig/listenerList/listener[name='HTTPS-ipv6']/name" "$CONF_FILE" 2>/dev/null || true)
	if [[ -z "$https_listener_exists" ]]; then
		log_error "HTTPS listener not found in $CONF_FILE. Cannot proceed with SNI mapping."
		SCRIPT_EXIT_STATUS=1; return 1
	fi

	for ln in HTTPS ${https6_listener_exists:+HTTPS-ipv6}; do
		local mapping_exists
		mapping_exists=$(xmlstarlet sel -t -v "/httpServerConfig/listenerList/listener[name='$ln']/vhostMapList/vhostMap[domain='$domain']/domain" "$CONF_FILE" 2>/dev/null || true)
		if [[ -z "$mapping_exists" ]]; then
			log "Adding vhostMap for domain '$domain' to $ln listener"
			# Ensure vhostMapList exists
			local vhost_maplist_exists
			vhost_maplist_exists=$(xmlstarlet sel -t -v "/httpServerConfig/listenerList/listener[name='$ln']/vhostMapList/vhostMap/domain" "$CONF_FILE" 2>/dev/null || true)
			if [[ -z "$vhost_maplist_exists" ]]; then
				sudo xmlstarlet ed -L -s "/httpServerConfig/listenerList/listener[name='$ln']" -t elem -n "vhostMapList" "$CONF_FILE" || { log_error "Failed to create vhostMapList for $ln"; return 1; }
			fi

			# Add vhostMap
			sudo xmlstarlet ed -L -s "/httpServerConfig/listenerList/listener[name='$ln']/vhostMapList" -t elem -n "vhostMap" "$CONF_FILE" || { log_error "Failed to create vhostMap for $ln"; return 1; }
			sudo xmlstarlet ed -L -s "/httpServerConfig/listenerList/listener[name='$ln']/vhostMapList/vhostMap[last()]" -t elem -n "vhost" -v "$vhost_name" "$CONF_FILE" || { log_error "Failed to set vhost in map for $ln"; return 1; }
			sudo xmlstarlet ed -L -s "/httpServerConfig/listenerList/listener[name='$ln']/vhostMapList/vhostMap[last()]" -t elem -n "domain" -v "$domain" "$CONF_FILE" || { log_error "Failed to set domain in map for $ln"; return 1; }
		else
			log "vhostMap for '$domain' already exists on $ln listener"
		fi
	done

	# 3) Validate final config
	if ! xmllint --noout "$CONF_FILE" 2>/dev/null; then
		log_error "Invalid XML structure after SNI configuration. Restoring backup..."
		sudo cp -a "$BACKUP_FILE" "$CONF_FILE"
		SCRIPT_EXIT_STATUS=1; return 1
	fi

	log_success "SNI configured for '$domain' on port 443 with vhost '$vhost_name'"
}

cleanup_xml() {
    local domain="$1"
    log "Cleaning up XML structure..."
    sudo cp -a "$CONF_FILE" "$BACKUP_FILE"
    
    # Remove any incorrectly placed vhostMapList elements
    sudo xmlstarlet ed -L \
        -d "//listenerList/vhostMapList[not(parent::listener)]" \
        "$CONF_FILE" || {
            log_error "Failed to clean up XML structure"
            SCRIPT_EXIT_STATUS=1
            return 1
        }
    
    # Validate the cleaned config
    if ! xmllint --noout "$CONF_FILE" 2>/dev/null; then
        log_error "Invalid XML structure after cleanup. Restoring backup..."
        sudo cp -a "$BACKUP_FILE" "$CONF_FILE"
        SCRIPT_EXIT_STATUS=1
        return 1
    fi
    log_success "XML structure cleaned up successfully"
}

restart_litespeed() {
    log "Restarting LiteSpeed server..."
    sudo systemctl restart lsws || {
        log_error "Failed to restart LiteSpeed. Restoring backup..."
        sudo cp -a "$BACKUP_FILE" "$CONF_FILE"
        SCRIPT_EXIT_STATUS=3
        return 1
    }
    log_success "LiteSpeed server restarted successfully"
}

# === Main Script Logic ===
main() {
    # Setup logging first
    setup_logging
    log "Starting SSL Manager V2.0.4"

    # Parse parameters
    for arg in "$@"; do
        case $arg in
            --public-ip=*) PUBLIC_IP="${arg#*=}"; log_verbose "Set public IP: $PUBLIC_IP";;
            --domain=*) PRIMARY_DOMAIN="${arg#*=}"; DOMAINS=("$PRIMARY_DOMAIN"); log_verbose "Set primary domain: $PRIMARY_DOMAIN";;
            --domains=*) IFS=',' read -ra DOMAINS <<< "${arg#*=}"; PRIMARY_DOMAIN="${DOMAINS[0]}"; log_verbose "Set domains: ${DOMAINS[*]}";;
            --email=*) EMAIL="${arg#*=}"; log_verbose "Set email: $EMAIL";;
            --vhost=*) VHOST_NAME="${arg#*=}"; log_verbose "Set vhost name: $VHOST_NAME";;
            --verbose) VERBOSE=1; log "Verbose mode enabled";;
            *) log_error "Invalid argument: $arg"; SCRIPT_EXIT_STATUS=1; exit 1;;
        esac
    done

    [[ -z "$PRIMARY_DOMAIN" || -z "$PUBLIC_IP" || -z "$EMAIL" ]] && {
        log_error "Missing required parameters. Provide --domains, --public-ip, and --email."
        SCRIPT_EXIT_STATUS=1; exit 1
    }

    validate_domain "$PRIMARY_DOMAIN" || { log_error "Invalid domain '$PRIMARY_DOMAIN'"; SCRIPT_EXIT_STATUS=1; exit 1; }
    validate_ip "$PUBLIC_IP" || { log_error "Invalid IP '$PUBLIC_IP'"; SCRIPT_EXIT_STATUS=1; exit 1; }
    validate_email "$EMAIL" || { log_error "Invalid email '$EMAIL'"; SCRIPT_EXIT_STATUS=1; exit 1; }

    validate_dns "$PRIMARY_DOMAIN" "$PUBLIC_IP" || { log_error "DNS validation failed"; SCRIPT_EXIT_STATUS=1; exit 1; }
    validate_http_access "$PRIMARY_DOMAIN" || { log_error "HTTP access validation failed"; SCRIPT_EXIT_STATUS=1; exit 1; }

    log_verbose "Checking dependencies..."
    check_command xmllint libxml2
    check_command xmlstarlet xmlstarlet
    check_command certbot certbot
    check_command dig bind-utils
    check_command curl curl
    check_command openssl openssl

    create_default_backup
    for domain in "${DOMAINS[@]}"; do
        issue_certificate "$domain" "$EMAIL"
        update_httpd_config "$domain" "$PUBLIC_IP"
        cleanup_xml "$domain"
    done

    restart_litespeed

    log_success "SSL Manager completed successfully"
    SCRIPT_EXIT_STATUS=0
}

# === Entry Point ===
main "$@"