#!/bin/bash

# Authentication Management Script for docker-compose-nas
# Enables or disables Authelia authentication for specific services
# Created: April 26, 2025

set -e

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

# Files
COMPOSE_FILE="docker-compose.yml"
TIMESTAMP=$(date +"%Y%m%d-%H%M%S")
BACKUP_FILE="docker-compose.${TIMESTAMP}.bak"

# Print section header
print_header() {
    echo -e "\n${CYAN}${BOLD}$1${NC}"
    echo -e "${CYAN}$(printf '=%.0s' $(seq 1 ${#1}))${NC}"
}

# Create a backup of the docker-compose.yml file
create_backup() {
    echo -e "${BLUE}Creating backup of $COMPOSE_FILE as $BACKUP_FILE...${NC}"
    cp "$COMPOSE_FILE" "$BACKUP_FILE"
}

# Get the current auth status for a service
get_auth_status() {
    local service=$1
    # Check if the service has authelia-auth middleware
    if grep -q "traefik.http.routers.$service.middlewares=.*authelia-auth" "$COMPOSE_FILE"; then
        echo "enabled"
    elif grep -q "traefik.http.routers.$service.middlewares=" "$COMPOSE_FILE"; then
        # Service has middlewares, but no authelia-auth
        echo "disabled"
    else
        echo "unknown"
    fi
}

# Enable authentication for a service
enable_auth() {
    local service=$1
    echo -e "${BLUE}Enabling authentication for $service...${NC}"
    
    # Check if the service exists in the compose file
    if ! grep -q "container_name: $service" "$COMPOSE_FILE"; then
        echo -e "${RED}Service $service not found in $COMPOSE_FILE${NC}"
        return 1
    fi
    
    # Update the middlewares configuration
    # This is a bit complex because different services have different middleware configurations
    # We need to find the existing middlewares line and add authelia-auth if it doesn't exist
    
    # First, find the service section
    local service_start=$(grep -n "container_name: $service" "$COMPOSE_FILE" | cut -d':' -f1)
    if [ -z "$service_start" ]; then
        echo -e "${RED}Could not find service $service in $COMPOSE_FILE${NC}"
        return 1
    fi
    
    # Find the next service or EOF
    local next_service=$(tail -n +$((service_start+1)) "$COMPOSE_FILE" | grep -n "container_name:" | head -1 | cut -d':' -f1)
    if [ -n "$next_service" ]; then
        next_service=$((service_start + next_service))
    else
        next_service=$(wc -l < "$COMPOSE_FILE")
    fi
    
    # Extract the service section
    local service_section=$(sed -n "${service_start},${next_service}p" "$COMPOSE_FILE")
    
    # Find if there's a middlewares line
    if echo "$service_section" | grep -q "traefik.http.routers.$service.middlewares="; then
        # Get the middlewares line
        local middlewares_line=$(echo "$service_section" | grep "traefik.http.routers.$service.middlewares=")
        
        # Check if authelia-auth is already in the middlewares
        if echo "$middlewares_line" | grep -q "authelia-auth@docker"; then
            echo -e "${GREEN}Authentication already enabled for $service${NC}"
            return 0
        fi
        
        # Add authelia-auth to the middlewares
        # Different patterns depending on what's in the middlewares line
        if echo "$middlewares_line" | grep -q "middlewares=$"; then
            # Empty middlewares
            sed -i "s|traefik.http.routers.$service.middlewares=|traefik.http.routers.$service.middlewares=authelia-auth@docker|" "$COMPOSE_FILE"
        elif echo "$middlewares_line" | grep -q "middlewares=.*@docker"; then
            # Already has some middleware with @docker
            sed -i "s|traefik.http.routers.$service.middlewares=\(.*\)@docker|traefik.http.routers.$service.middlewares=\1,authelia-auth@docker|" "$COMPOSE_FILE"
        else
            # Has middlewares but no @docker suffix
            sed -i "s|traefik.http.routers.$service.middlewares=\(.*\)|traefik.http.routers.$service.middlewares=\1,authelia-auth@docker|" "$COMPOSE_FILE"
        fi
    else
        # No middlewares line, add one
        local labels_line=$(echo "$service_section" | grep -n "labels:" | cut -d':' -f1)
        if [ -z "$labels_line" ]; then
            echo -e "${RED}Could not find labels section for service $service${NC}"
            return 1
        fi
        
        # Insert the middlewares line after the rule line
        local rule_line=$(echo "$service_section" | grep -n "traefik.http.routers.$service.rule=" | cut -d':' -f1)
        if [ -n "$rule_line" ]; then
            rule_line=$((service_start + rule_line))
            sed -i "${rule_line}a \ \ \ \ \ \ - traefik.http.routers.$service.middlewares=authelia-auth@docker" "$COMPOSE_FILE"
        else
            echo -e "${RED}Could not find rule line for service $service${NC}"
            return 1
        fi
    fi
    
    echo -e "${GREEN}Authentication enabled for $service${NC}"
    return 0
}

# Disable authentication for a service
disable_auth() {
    local service=$1
    echo -e "${BLUE}Disabling authentication for $service...${NC}"
    
    # Check if the service exists in the compose file
    if ! grep -q "container_name: $service" "$COMPOSE_FILE"; then
        echo -e "${RED}Service $service not found in $COMPOSE_FILE${NC}"
        return 1
    fi
    
    # Update the middlewares configuration by removing authelia-auth
    
    # First, find if there's a middlewares line with authelia-auth
    if grep -q "traefik.http.routers.$service.middlewares=.*authelia-auth@docker" "$COMPOSE_FILE"; then
        # Simple case: only authelia-auth
        if grep -q "traefik.http.routers.$service.middlewares=authelia-auth@docker" "$COMPOSE_FILE"; then
            # Remove authelia-auth@docker, leaving only the label prefix
            sed -i "s|traefik.http.routers.$service.middlewares=authelia-auth@docker|traefik.http.routers.$service.middlewares=|" "$COMPOSE_FILE"
        # authelia-auth is at the beginning with a comma
        elif grep -q "traefik.http.routers.$service.middlewares=authelia-auth@docker," "$COMPOSE_FILE"; then
            # Remove authelia-auth@docker,
            sed -i "s|traefik.http.routers.$service.middlewares=authelia-auth@docker,|traefik.http.routers.$service.middlewares=|" "$COMPOSE_FILE"
        # authelia-auth is at the end with a comma before it
        elif grep -q "traefik.http.routers.$service.middlewares=.*,authelia-auth@docker" "$COMPOSE_FILE"; then
            # Remove ,authelia-auth@docker
            sed -i "s|,authelia-auth@docker||" "$COMPOSE_FILE"
        # authelia-auth is in the middle
        elif grep -q "traefik.http.routers.$service.middlewares=.*,authelia-auth@docker,.*" "$COMPOSE_FILE"; then
            # Remove ,authelia-auth@docker
            sed -i "s|,authelia-auth@docker,|,|" "$COMPOSE_FILE"
        else
            echo -e "${RED}Could not determine how to remove authelia-auth from middlewares for $service${NC}"
            return 1
        fi
        
        echo -e "${GREEN}Authentication disabled for $service${NC}"
    else
        echo -e "${GREEN}Authentication already disabled for $service${NC}"
    fi
    
    return 0
}

# List services and their auth status
list_services() {
    print_header "Services Authentication Status"
    
    echo -e "${BLUE}Checking services in $COMPOSE_FILE...${NC}"
    echo -e "${CYAN}SERVICE\t\tAUTH STATUS${NC}"
    echo -e "${CYAN}-------\t\t-----------${NC}"
    
    # Find all container names
    local services=$(grep "container_name:" "$COMPOSE_FILE" | awk '{print $3}')
    
    for service in $services; do
        # Skip infrastructure containers
        if [[ "$service" == "redis" || "$service" == "authelia" || "$service" == "traefik" || "$service" == "tailscale" || "$service" == "watchtower" || "$service" == "autoheal" || "$service" == "middlewares" ]]; then
            continue
        fi
        
        local status=$(get_auth_status "$service")
        
        # Pad service name for better formatting
        printf "${BOLD}%-20s${NC}" "$service"
        
        case "$status" in
            "enabled")
                echo -e "${GREEN}Enabled${NC}"
                ;;
            "disabled")
                echo -e "${YELLOW}Disabled${NC}"
                ;;
            *)
                echo -e "${RED}Unknown${NC}"
                ;;
        esac
    done
}

# Display help
show_help() {
    print_header "Authentication Management Script"
    echo -e "Usage: $0 [options]"
    echo -e ""
    echo -e "Options:"
    echo -e "  ${CYAN}list${NC}                       List all services and their auth status"
    echo -e "  ${CYAN}enable <service>${NC}           Enable authentication for a service"
    echo -e "  ${CYAN}disable <service>${NC}          Disable authentication for a service"
    echo -e "  ${CYAN}enable-all${NC}                 Enable authentication for all services"
    echo -e "  ${CYAN}disable-all${NC}                Disable authentication for all services"
    echo -e "  ${CYAN}help${NC}                       Show this help message"
    echo -e ""
    echo -e "Examples:"
    echo -e "  $0 list"
    echo -e "  $0 enable jellyfin"
    echo -e "  $0 disable homepage"
    echo -e "  $0 enable-all"
    echo -e ""
}

# Main function
main() {
    if [ $# -eq 0 ]; then
        show_help
        exit 0
    fi
    
    # Create backup of compose file
    create_backup
    
    case "$1" in
        "list")
            list_services
            ;;
        "enable")
            if [ -z "$2" ]; then
                echo -e "${RED}Error: No service specified${NC}"
                echo -e "Usage: $0 enable <service>"
                exit 1
            fi
            enable_auth "$2"
            ;;
        "disable")
            if [ -z "$2" ]; then
                echo -e "${RED}Error: No service specified${NC}"
                echo -e "Usage: $0 disable <service>"
                exit 1
            fi
            disable_auth "$2"
            ;;
        "enable-all")
            print_header "Enabling Authentication for All Services"
            
            local services=$(grep "container_name:" "$COMPOSE_FILE" | awk '{print $3}')
            for service in $services; do
                # Skip infrastructure containers
                if [[ "$service" == "redis" || "$service" == "authelia" || "$service" == "traefik" || "$service" == "tailscale" || "$service" == "watchtower" || "$service" == "autoheal" || "$service" == "middlewares" ]]; then
                    continue
                fi
                
                enable_auth "$service"
            done
            ;;
        "disable-all")
            print_header "Disabling Authentication for All Services"
            
            local services=$(grep "container_name:" "$COMPOSE_FILE" | awk '{print $3}')
            for service in $services; do
                # Skip infrastructure containers
                if [[ "$service" == "redis" || "$service" == "authelia" || "$service" == "traefik" || "$service" == "tailscale" || "$service" == "watchtower" || "$service" == "autoheal" || "$service" == "middlewares" ]]; then
                    continue
                fi
                
                disable_auth "$service"
            done
            ;;
        "help")
            show_help
            ;;
        *)
            echo -e "${RED}Error: Unknown command: $1${NC}"
            show_help
            exit 1
            ;;
    esac
    
    echo -e "\n${GREEN}${BOLD}Done!${NC}"
    echo -e "${BLUE}Original configuration backed up to: $BACKUP_FILE${NC}"
    echo -e "${YELLOW}Remember to restart the stack for changes to take effect:${NC}"
    echo -e "${CYAN}  docker compose down${NC}"
    echo -e "${CYAN}  docker compose up -d${NC}"
}

main "$@"