#!/bin/bash
# ==============================================================================
# WALLARM RECONFIGURATION SCRIPT - V1.0
# ==============================================================================
# Purpose: Modify nginx configuration of an existing Wallarm node
# Features:
# - Update set_real_ip_from (trusted proxy IPs/CIDRs)
# - Change wallarm_mode (monitoring/block)
# - Backup current config before changes
# - Interactive prompts with validation
# ==============================================================================

set -euo pipefail

# Color definitions
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[1;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m'

# Defaults
INSTANCE_DIR="/opt"
INSTANCE_NAME=""

# Function to find the Wallarm instance directory
find_wallarm_instance() {
    local dirs=()
    while IFS= read -r dir; do
        if [[ -d "$dir" && -f "$dir/nginx.conf" && -f "$dir/start.sh" ]]; then
            dirs+=("$dir")
        fi
    done < <(find "$INSTANCE_DIR" -maxdepth 1 -type d -name "wallarm-*" 2>/dev/null)
    
    if [ ${#dirs[@]} -eq 0 ]; then
        echo -e "${RED}No Wallarm instance found in $INSTANCE_DIR.${NC}"
        exit 1
    elif [ ${#dirs[@]} -eq 1 ]; then
        INSTANCE_DIR="${dirs[0]}"
        INSTANCE_NAME=$(basename "$INSTANCE_DIR")
        echo -e "${GREEN}Found instance: $INSTANCE_NAME${NC}"
    else
        echo -e "${YELLOW}Multiple Wallarm instances found:${NC}"
        for i in "${!dirs[@]}"; do
            echo "$((i+1)). $(basename "${dirs[$i]}")"
        done
        read -r -p "Select instance number: " choice
        if [[ "$choice" =~ ^[0-9]+$ ]] && [ "$choice" -ge 1 ] && [ "$choice" -le ${#dirs[@]} ]; then
            INSTANCE_DIR="${dirs[$((choice-1))]}"
            INSTANCE_NAME=$(basename "$INSTANCE_DIR")
        else
            echo -e "${RED}Invalid selection.${NC}"
            exit 1
        fi
    fi
}

# Validate IP/CIDR format
validate_proxy() {
    local proxy="$1"
    if [[ "$proxy" =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}(/[0-9]{1,2})?$ ]]; then
        IFS='/' read -r ip cidr <<< "$proxy"
        IFS='.' read -r o1 o2 o3 o4 <<< "$ip"
        if [ "$o1" -le 255 ] && [ "$o2" -le 255 ] && [ "$o3" -le 255 ] && [ "$o4" -le 255 ]; then
            if [ -z "$cidr" ] || ( [ "$cidr" -ge 0 ] && [ "$cidr" -le 32 ] ); then
                return 0
            fi
        fi
    fi
    return 1
}

# Parse current configuration to get existing values
parse_current_config() {
    local config_file="$1"
    # Get current wallarm_mode
    current_mode=$(grep -oP 'wallarm_mode\s+\K\S+' "$config_file" | head -1)
    # Get current set_real_ip_from lines
    current_proxies=$(grep -oP 'set_real_ip_from\s+\K\S+' "$config_file")
}

# Update configuration
update_config() {
    local config_file="$1"
    local backup_file="$config_file.backup.$(date +%Y%m%d_%H%M%S)"
    
    echo -e "${YELLOW}Backing up current config to $backup_file${NC}"
    cp "$config_file" "$backup_file"
    
    # Read new values interactively
    echo -e "\n${CYAN}${BOLD}Current set_real_ip_from entries:${NC}"
    if [ -n "$current_proxies" ]; then
        echo "$current_proxies" | while read -r proxy; do
            echo "  $proxy"
        done
    else
        echo "  (none)"
    fi
    
    echo -e "\n${YELLOW}Do you want to change the trusted proxy IPs/CIDRs? (y/N)${NC}"
    read -r change_proxy
    if [[ "$change_proxy" =~ ^[Yy]$ ]]; then
        echo -e "${YELLOW}Enter new trusted proxy IPs/CIDRs (space-separated, or empty to remove all):${NC}"
        read -r new_proxies_input
        new_proxies=()
        if [[ -n "$new_proxies_input" ]]; then
            IFS=' ' read -ra proxy_array <<< "$new_proxies_input"
            for proxy in "${proxy_array[@]}"; do
                proxy=$(echo "$proxy" | xargs)
                if validate_proxy "$proxy"; then
                    new_proxies+=("$proxy")
                else
                    echo -e "${RED}Invalid format: $proxy. Skipping.${NC}"
                fi
            done
        fi
    else
        # Keep existing
        while read -r proxy; do
            new_proxies+=("$proxy")
        done <<< "$current_proxies"
    fi
    
    echo -e "\n${CYAN}${BOLD}Current wallarm_mode:${NC} ${current_mode:-not set}"
    echo -e "${YELLOW}Do you want to change the wallarm_mode? (y/N)${NC}"
    read -r change_mode
    if [[ "$change_mode" =~ ^[Yy]$ ]]; then
        echo -e "${YELLOW}Select new mode:${NC}"
        echo "1. monitoring"
        echo "2. block"
        read -r mode_choice
        case "$mode_choice" in
            1) new_mode="monitoring" ;;
            2) new_mode="block" ;;
            *) echo -e "${RED}Invalid choice, keeping current mode.${NC}"; new_mode="$current_mode" ;;
        esac
    else
        new_mode="$current_mode"
    fi
    
    # Now rebuild the config file
    # We'll create a temporary file and replace the original
    temp_config=$(mktemp)
    
    # Read original config line by line and modify as needed
    in_server_block=false
    while IFS= read -r line; do
        # Detect start of server block
        if [[ "$line" =~ ^[[:space:]]*server[[:space:]]*{ ]]; then
            in_server_block=true
        fi
        
        # If we are inside server block, we may need to replace lines
        if $in_server_block; then
            # Replace set_real_ip_from lines with new ones
            if [[ "$line" =~ ^[[:space:]]*set_real_ip_from[[:space:]]+ ]]; then
                # Skip original set_real_ip_from lines (will be added later)
                continue
            fi
            # Replace wallarm_mode line
            if [[ "$line" =~ ^[[:space:]]*wallarm_mode[[:space:]]+ ]]; then
                # We'll add new line after processing all lines
                continue
            fi
        fi
        
        # Write line to temp file
        echo "$line" >> "$temp_config"
        
        # After writing the line, if we are at the end of the server block, we may need to insert new directives
        if $in_server_block && [[ "$line" =~ ^[[:space:]]*}$ ]]; then
            in_server_block=false
            # Insert the new set_real_ip_from lines just before the closing brace
            if [ ${#new_proxies[@]} -gt 0 ]; then
                for proxy in "${new_proxies[@]}"; do
                    echo "    set_real_ip_from  $proxy;" >> "$temp_config"
                done
                echo "    real_ip_header    X-Real-IP;" >> "$temp_config"
                echo "    real_ip_recursive on;" >> "$temp_config"
            elif [ -n "$current_proxies" ]; then
                # If we removed all proxies, we should also remove the real_ip_header and real_ip_recursive lines
                # But that's tricky; we'll just not add them, but they might remain in the file if they were separate.
                # Simpler: after rebuild, we need to ensure they are not there. We'll do a final cleanup.
                echo -e "${YELLOW}Removing all set_real_ip_from directives.${NC}"
            fi
            # Insert new wallarm_mode
            if [ -n "$new_mode" ]; then
                echo "    wallarm_mode $new_mode;" >> "$temp_config"
            fi
        fi
    done < "$config_file"
    
    # After building the temp file, we need to ensure any leftover real_ip_header lines are removed if no proxies.
    if [ ${#new_proxies[@]} -eq 0 ]; then
        # Remove lines containing real_ip_header and real_ip_recursive if they exist
        sed -i '/real_ip_header/d' "$temp_config"
        sed -i '/real_ip_recursive/d' "$temp_config"
    fi
    
    # Replace the original config with the new one
    sudo mv "$temp_config" "$config_file"
    sudo chmod 644 "$config_file"
    
    echo -e "${GREEN}Configuration updated.${NC}"
}

restart_container() {
    local container_name="$1"
    echo -e "${YELLOW}Restarting container $container_name to apply changes...${NC}"
    if sudo docker ps --format "{{.Names}}" | grep -q "^$container_name$"; then
        sudo docker restart "$container_name"
        echo -e "${GREEN}Container restarted.${NC}"
    else
        echo -e "${RED}Container $container_name is not running. Starting it...${NC}"
        if [ -f "$INSTANCE_DIR/start.sh" ]; then
            sudo "$INSTANCE_DIR/start.sh"
        else
            echo -e "${RED}No start script found. Please start manually: sudo docker start $container_name${NC}"
            exit 1
        fi
    fi
}

main() {
    echo -e "${BLUE}${BOLD}"
    echo "╔══════════════════════════════════════════════════════════════╗"
    echo "║         WALLARM RECONFIGURATION SCRIPT - V1.0               ║"
    echo "║      Modify nginx.conf (trusted proxies / mode)             ║"
    echo "╚══════════════════════════════════════════════════════════════╝${NC}"
    
    find_wallarm_instance
    
    local config_file="$INSTANCE_DIR/nginx.conf"
    if [ ! -f "$config_file" ]; then
        echo -e "${RED}Configuration file not found: $config_file${NC}"
        exit 1
    fi
    
    parse_current_config "$config_file"
    update_config "$config_file"
    
    echo -e "${YELLOW}Do you want to restart the container now? (Y/n)${NC}"
    read -r restart_choice
    if [[ ! "$restart_choice" =~ ^[Nn]$ ]]; then
        restart_container "$INSTANCE_NAME"
    else
        echo -e "${YELLOW}Changes will take effect after container restart.${NC}"
        echo -e "You can restart later with: sudo docker restart $INSTANCE_NAME"
    fi
    
    echo -e "\n${GREEN}${BOLD}Reconfiguration completed.${NC}"
}

main "$@"