#!/bin/bash
set -euo pipefail

# ===============================================================================
# Linux System Audit Deployment Script
# This script helps deploy and run the audit across multiple devices
# Version: 2.0 - Enhanced with better error handling and security features
# ===============================================================================

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
INVENTORY_FILE="$SCRIPT_DIR/inventory.ini"
PLAYBOOK_FILE="$SCRIPT_DIR/linux_audit_playbook.yml"
AUDIT_SCRIPT="$SCRIPT_DIR/linux_system_audit.sh"
RESULTS_DIR="$SCRIPT_DIR/audit_results"

# Enhanced error handling
error_handler() {
    local exit_code=$1
    local line_no=$2
    local last_command=$3
    
    echo -e "${RED}Error occurred in: $last_command${NC}"
    echo -e "${RED}Exit code: $exit_code${NC}"
    echo -e "${RED}Line number: $line_no${NC}"
    exit $exit_code
}

# Set up error trap
trap 'error_handler $? $LINENO "$BASH_COMMAND"' ERR

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

print_status() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

check_dependencies() {
    print_status "Checking dependencies..."

    local missing_deps=()
    local optional_deps=()

    # Required dependencies
    if ! command -v ansible >/dev/null 2>&1; then
        missing_deps+=("ansible")
    fi

    if ! command -v ansible-playbook >/dev/null 2>&1; then
        missing_deps+=("ansible-playbook")
    fi

    # Optional but recommended dependencies
    if ! command -v jq >/dev/null 2>&1; then
        optional_deps+=("jq")
    fi

    if ! command -v tar >/dev/null 2>&1; then
        optional_deps+=("tar")
    fi

    if [ ${#missing_deps[@]} -ne 0 ]; then
        print_error "Missing required dependencies: ${missing_deps[*]}"
        echo "Please install Ansible:"
        echo "  Ubuntu/Debian: sudo apt update && sudo apt install ansible"
        echo "  Fedora: sudo dnf install ansible"
        echo "  Or via pip: pip3 install ansible"
        exit 1
    fi

    if [ ${#optional_deps[@]} -ne 0 ]; then
        print_warning "Missing optional dependencies: ${optional_deps[*]}"
        echo "These are recommended but not required:"
        echo "  jq: For enhanced JSON processing"
        echo "  tar: For result compression"
    fi

    print_success "All required dependencies found"
    
    # Check Ansible version
    local ansible_version=$(ansible --version | head -1 | awk '{print $2}')
    print_status "Ansible version: $ansible_version"
}

check_files() {
    print_status "Checking required files..."

    local missing_files=()

    if [ ! -f "$INVENTORY_FILE" ]; then
        missing_files+=("$INVENTORY_FILE")
    fi

    if [ ! -f "$PLAYBOOK_FILE" ]; then
        missing_files+=("$PLAYBOOK_FILE")
    fi

    if [ ! -f "$AUDIT_SCRIPT" ]; then
        missing_files+=("$AUDIT_SCRIPT")
    fi

    if [ ${#missing_files[@]} -ne 0 ]; then
        print_error "Missing files: ${missing_files[*]}"
        exit 1
    fi

    print_success "All required files found"
}

test_connectivity() {
    print_status "Testing connectivity to hosts..."

    # Test connectivity with detailed output
    local connectivity_test=$(ansible all -i "$INVENTORY_FILE" -m ping --one-line 2>&1)
    local success_count=$(echo "$connectivity_test" | grep -c "SUCCESS" || echo "0")
    local total_hosts=$(ansible all -i "$INVENTORY_FILE" --list-hosts 2>/dev/null | grep -v "hosts" | wc -l)
    
    print_status "Connectivity test results: $success_count/$total_hosts hosts reachable"
    
    if [ "$success_count" -eq "$total_hosts" ]; then
        print_success "Successfully connected to all hosts"
    elif [ "$success_count" -gt 0 ]; then
        print_warning "Connected to $success_count/$total_hosts hosts"
        echo "Some hosts may not be reachable. Check your inventory file and SSH keys."
        echo "Failed connections:"
        echo "$connectivity_test" | grep -v "SUCCESS" || true
        
        read -p "Continue with available hosts? (y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            exit 1
        fi
    else
        print_error "No hosts are reachable. Please check:"
        echo "  1. SSH key authentication is set up"
        echo "  2. Inventory file is correct"
        echo "  3. Hosts are online and accessible"
        echo "  4. Firewall rules allow SSH connections"
        exit 1
    fi
}

run_audit() {
    print_status "Starting audit across all hosts..."

    # Create results directory
    mkdir -p "$RESULTS_DIR"

    # Run the playbook
    if ansible-playbook -i "$INVENTORY_FILE" "$PLAYBOOK_FILE" -v; then
        print_success "Audit completed successfully"
    else
        print_error "Audit failed. Check the output above for details."
        exit 1
    fi
}

generate_summary_report() {
    print_status "Generating summary report..."

    local summary_file="$RESULTS_DIR/MASTER_SUMMARY_$(date +%Y%m%d_%H%M%S).txt"

    cat > "$summary_file" << EOF
===============================================================================
COMPREHENSIVE LINUX SYSTEM AUDIT SUMMARY REPORT
Generated: $(date)
Script Version: 1.0
===============================================================================

OVERVIEW:
---------
EOF

    # Count hosts
    local total_hosts=$(find "$RESULTS_DIR" -maxdepth 1 -type d | grep -v "^$RESULTS_DIR$" | wc -l)
    echo "Total hosts audited: $total_hosts" >> "$summary_file"
    echo "" >> "$summary_file"

    # Process each host
    for host_dir in "$RESULTS_DIR"/*; do
        if [ -d "$host_dir" ] && [ "$(basename "$host_dir")" != "audit_results" ]; then
            local hostname=$(basename "$host_dir")
            echo "=== $hostname ===" >> "$summary_file"

            if [ -f "$host_dir/SUMMARY.txt" ]; then
                cat "$host_dir/SUMMARY.txt" >> "$summary_file"
            else
                echo "No summary file found" >> "$summary_file"
            fi
            echo "" >> "$summary_file"
        fi
    done

    print_success "Summary report generated: $summary_file"
}

create_dashboard() {
    print_status "Creating dynamic HTML dashboard..."

    local dashboard_file="$RESULTS_DIR/dashboard.html"
    local all_results_file="$RESULTS_DIR/all_results.json"

    # Aggregate all JSON results into a single file
    echo "[" > "$all_results_file"
    first=true
    for host_dir in "$RESULTS_DIR"/*; do
        if [ -d "$host_dir" ]; then
            local json_file="$host_dir/results.json"
            if [ -f "$json_file" ]; then
                if [ "$first" = false ]; then
                    echo "," >> "$all_results_file"
                fi
                cat "$json_file" >> "$all_results_file"
                first=false
            fi
        fi
    done
    echo "]" >> "$all_results_file"

    cat > "$dashboard_file" << 'EOF'
<!DOCTYPE html>
<html>
<head>
    <title>Linux System Audit Dashboard</title>
    <style>
        body { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif; margin: 0; background-color: #f0f2f5; }
        .header { background-color: #fff; padding: 20px; border-bottom: 1px solid #ddd; text-align: center; }
        .header h1 { margin: 0; }
        .container { display: flex; flex-wrap: wrap; padding: 20px; justify-content: center; }
        .card { background-color: #fff; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1); margin: 10px; padding: 20px; width: 300px; }
        .card h2 { margin-top: 0; color: #333; border-bottom: 1px solid #eee; padding-bottom: 10px; }
        .card p { color: #666; }
        .card .label { font-weight: bold; color: #333; }
        .status-ok { color: #28a745; }
        .status-warning { color: #ffc107; }
        .status-error { color: #dc3545; }
    </style>
</head>
<body>
    <div class="header">
        <h1>Linux System Audit Dashboard</h1>
        <p id="generation-date"></p>
    </div>

    <div id="hosts-container" class="container">
        <!-- Host information will be inserted here -->
    </div>

    <script>
        document.addEventListener('DOMContentLoaded', function() {
            document.getElementById('generation-date').textContent = `Generated: ${new Date().toLocaleString()}`;

            fetch('all_results.json')
                .then(response => response.json())
                .then(data => {
                    const container = document.getElementById('hosts-container');
                    data.forEach(host => {
                        const card = document.createElement('div');
                        card.className = 'card';

                        let ufwStatusClass = 'status-ok';
                        if (host.security.ufw_status !== 'active') {
                            ufwStatusClass = 'status-warning';
                        }

                        let rootLoginClass = 'status-ok';
                        if (host.security.ssh_root_login.toLowerCase() !== 'no') {
                            rootLoginClass = 'status-error';
                        }

                        card.innerHTML = `
                            <h2>${host.system.hostname}</h2>
                            <p><span class="label">OS:</span> ${host.system.os}</p>
                            <p><span class="label">Kernel:</span> ${host.system.kernel}</p>
                            <p><span class="label">IP Address:</span> ${host.system.ip_addresses}</p>
                            <p><span class="label">Uptime:</span> ${host.system.uptime}</p>
                            <p><span class="label">Running Containers:</span> ${host.containers.running_containers}</p>
                            <p><span class="label">UFW Status:</span> <span class="${ufwStatusClass}">${host.security.ufw_status}</span></p>
                            <p><span class="label">SSH Root Login:</span> <span class="${rootLoginClass}">${host.security.ssh_root_login}</span></p>
                            <p><span class="label">Open Ports:</span> ${host.security.open_ports.join(', ')}</p>
                        `;
                        container.appendChild(card);
                    });
                })
                .catch(error => {
                    console.error('Error loading audit data:', error);
                    const container = document.getElementById('hosts-container');
                    container.innerHTML = '<p style="color: red;">Could not load audit data. Make sure all_results.json is available.</p>';
                });
        });
    </script>
</body>
</html>
EOF

    print_success "HTML dashboard created: $dashboard_file"
}

cleanup_old_results() {
    if [ -d "$RESULTS_DIR" ]; then
        print_status "Cleaning up old results..."
        rm -rf "$RESULTS_DIR"/*
        print_success "Old results cleaned"
    fi
}

show_help() {
    cat << EOF
Linux System Audit Deployment Script

Usage: $0 [OPTIONS]

OPTIONS:
    -h, --help           Show this help message
    -c, --check          Check dependencies and connectivity only
    -n, --no-cleanup     Don't cleanup old results
    -q, --quick          Skip connectivity test
    --inventory FILE     Use custom inventory file
    --results-dir DIR    Use custom results directory
    --setup-sudo         Set up passwordless sudo on all hosts in inventory

EXAMPLES:
    $0                   Run full audit
    $0 -c                Check setup only
    $0 --setup-sudo      Set up passwordless sudo
    $0 --inventory /path/to/custom/inventory.ini

Before running:
1. Edit inventory.ini with your server details
2. Ensure SSH key authentication is set up
3. Verify sudo access on target hosts
EOF
}

setup_passwordless_sudo() {
    print_status "Setting up passwordless sudo for all hosts in inventory..."

    if [ ! -f "$INVENTORY_FILE" ]; then
        print_error "Inventory file not found at $INVENTORY_FILE"
        exit 1
    fi

    # Extract user and host from inventory, ignore comments and empty lines
    mapfile -t hosts < <(grep -vE '^\s*(#|$|\[)' "$INVENTORY_FILE" | grep 'ansible_host' | awk '{print $1, $2}' | sed 's/ansible_host=//')

    if [ ${#hosts[@]} -eq 0 ]; then
        print_error "No hosts found in inventory file."
        exit 1
    fi

    for host_info in "${hosts[@]}"; do
        read -r alias host <<< "$host_info"
        user_host=$(awk -v alias="$alias" '$1 == alias {print $3}' "$INVENTORY_FILE" | sed 's/ansible_user=//')
        
        if [[ -z "$user_host" ]]; then
            # Fallback for different inventory formats
            user_host=$(grep "$host" "$INVENTORY_FILE" | grep "ansible_user" | sed 's/.*ansible_user=\([^ ]*\).*/\1@'${host}'/')
            user_host=${user_host%@*}
        fi

        if [[ -z "$user_host" ]]; then
            print_warning "Could not determine user for host $alias. Skipping."
            continue
        fi

        user=$(echo "$user_host" | cut -d'@' -f1)
        
        print_status "Configuring $alias ($user@$host)..."

        if [ "$user" = "root" ]; then
            print_success "Skipping $alias - already root user"
            continue
        fi

        # Create sudoers file for the user
        if ssh "$user@$host" "echo '$user ALL=(ALL) NOPASSWD: ALL' | sudo tee /etc/sudoers.d/audit-$user && sudo chmod 440 /etc/sudoers.d/audit-$user"; then
            print_success "Successfully configured passwordless sudo for $alias"
        else
            print_error "Failed to configure passwordless sudo for $alias"
        fi
    done

    print_success "Passwordless sudo setup complete!"
}

main() {
    local cleanup=true
    local check_only=false
    local quick=false
    local setup_sudo=false

    # Parse command line arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_help
                exit 0
                ;;
            -c|--check)
                check_only=true
                shift
                ;;
            -n|--no-cleanup)
                cleanup=false
                shift
                ;;
            -q|--quick)
                quick=true
                shift
                ;;
            --inventory)
                INVENTORY_FILE="$2"
                shift 2
                ;;
            --results-dir)
                RESULTS_DIR="$2"
                shift 2
                ;;
            --setup-sudo)
                setup_sudo=true
                shift
                ;;
            *)
                print_error "Unknown option: $1"
                show_help
                exit 1
                ;;
        esac
    done

    if [ "$setup_sudo" = true ]; then
        setup_passwordless_sudo
        exit 0
    fi

    print_status "Starting Linux System Audit Deployment"

    check_dependencies
    check_files

    if [ "$check_only" = true ]; then
        test_connectivity
        print_success "All checks passed!"
        exit 0
    fi

    if [ "$cleanup" = true ]; then
        cleanup_old_results
    fi

    if [ "$quick" = false ]; then
        test_connectivity
    fi

    run_audit
    generate_summary_report
    create_dashboard

    print_success "Audit deployment completed!"
    echo ""
    echo "Results available in: $RESULTS_DIR"
    echo "Open dashboard.html in your browser for a visual overview"
}

# Run main function with all arguments
main "$@"