#!/bin/bash
set -euo pipefail

# ===============================================================================
# Comprehensive Linux System Enumeration and Security Audit Script
# For Ubuntu, Debian, and Fedora systems
# Created for automated device inventory and security assessment
# Version: 2.0 - Enhanced with better error handling and security features
# ===============================================================================

# Configuration
OUTPUT_DIR="/tmp/system_audit_$(hostname)_$(date +%Y%m%d_%H%M%S)"
LOG_FILE="${OUTPUT_DIR}/audit.log"
RESULTS_FILE="${OUTPUT_DIR}/results.json"
START_TIME=$(date +%s)

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

# Enhanced error handling
error_handler() {
    local exit_code=$1
    local line_no=$2
    local bash_lineno=$3
    local last_command=$4
    local func_stack=$5
    
    echo -e "${RED}Error occurred in: $last_command${NC}" | tee -a "$LOG_FILE"
    echo -e "${RED}Exit code: $exit_code${NC}" | tee -a "$LOG_FILE"
    echo -e "${RED}Line number: $line_no${NC}" | tee -a "$LOG_FILE"
    echo -e "${RED}Function stack: $func_stack${NC}" | tee -a "$LOG_FILE"
    
    # Cleanup on error
    if [ -d "$OUTPUT_DIR" ]; then
        rm -rf "$OUTPUT_DIR"
    fi
    
    exit $exit_code
}

# Set up error trap
trap 'error_handler $? $LINENO $BASH_LINENO "$BASH_COMMAND" $(printf "::%s" ${FUNCNAME[@]:-})' ERR

# Create output directory
mkdir -p "$OUTPUT_DIR"

# Enhanced logging function with levels
log() {
    local level=$1
    shift
    local message="$*"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[$timestamp] [$level] $message" | tee -a "$LOG_FILE"
}

log_info() { log "INFO" "$*"; }
log_warn() { log "WARN" "$*"; }
log_error() { log "ERROR" "$*"; }
log_debug() { log "DEBUG" "$*"; }

print_section() {
    echo -e "\n${BLUE}==== $1 ====${NC}" | tee -a "$LOG_FILE"
}

print_subsection() {
    echo -e "\n${GREEN}--- $1 ---${NC}" | tee -a "$LOG_FILE"
}

print_warning() {
    echo -e "${YELLOW}WARNING: $1${NC}" | tee -a "$LOG_FILE"
}

print_error() {
    echo -e "${RED}ERROR: $1${NC}" | tee -a "$LOG_FILE"
}

# Input validation and environment checking
validate_environment() {
    log_info "Validating environment and dependencies..."
    
    local required_tools=("hostname" "uname" "lsblk" "ip" "cat" "grep" "awk")
    local missing_tools=()
    
    for tool in "${required_tools[@]}"; do
        if ! command -v "$tool" >/dev/null 2>&1; then
            missing_tools+=("$tool")
        fi
    done
    
    if [ ${#missing_tools[@]} -gt 0 ]; then
        log_error "Missing required tools: ${missing_tools[*]}"
        exit 1
    fi
    
    # Check for optional but recommended tools
    local optional_tools=("jq" "docker" "podman" "nmap" "vnstat" "ethtool")
    for tool in "${optional_tools[@]}"; do
        if ! command -v "$tool" >/dev/null 2>&1; then
            log_warn "Optional tool not found: $tool"
        fi
    done
    
    log_info "Environment validation completed"
}

# Check if running as root with enhanced security
check_privileges() {
    if [[ $EUID -ne 0 ]]; then
        print_warning "Script not running as root. Some checks may be limited."
        SUDO_CMD="sudo"
        # Verify sudo access
        if ! sudo -n true 2>/dev/null; then
            log_warn "Passwordless sudo not available. Some operations may fail."
        fi
    else
        SUDO_CMD=""
        log_info "Running with root privileges"
    fi
}

# System Information Collection
collect_system_info() {
    print_section "SYSTEM INFORMATION"

    print_subsection "Basic System Details"
    {
        echo "Hostname: $(hostname)"
        echo "FQDN: $(hostname -f 2>/dev/null || echo 'N/A')"
        echo "IP Addresses: $(hostname -I 2>/dev/null || echo 'N/A')"
        echo "Date/Time: $(date)"
        echo "Uptime: $(uptime)"
        echo "Load Average: $(cat /proc/loadavg)"
        echo "Architecture: $(uname -m)"
        echo "Kernel: $(uname -r)"
        echo "Distribution: $(lsb_release -d 2>/dev/null | cut -f2 || cat /etc/os-release | grep PRETTY_NAME | cut -d'=' -f2 | tr -d '"')"
        echo "Kernel Version: $(uname -v)"
    } | tee -a "$LOG_FILE"

    print_subsection "Hardware Information"
    {
        echo "CPU Info:"
        lscpu | tee -a "$LOG_FILE"
        echo -e "\nMemory Info:"
        free -h | tee -a "$LOG_FILE"
        echo -e "\nDisk Usage:"
        df -h | tee -a "$LOG_FILE"
        echo -e "\nBlock Devices:"
        lsblk | tee -a "$LOG_FILE"
        echo -e "\nPCI Devices:"
        lspci | tee -a "$LOG_FILE"
        echo -e "\nUSB Devices:"
        lsusb 2>/dev/null | tee -a "$LOG_FILE"
    }
}

# Network Information
collect_network_info() {
    print_section "NETWORK INFORMATION"

    print_subsection "Network Interfaces"
    {
        echo "Network Interfaces (ip addr):"
        if command -v ip >/dev/null 2>&1; then
            ip addr show | tee -a "$LOG_FILE"
        else
            log_warn "ip command not available, using ifconfig fallback"
            ifconfig 2>/dev/null | tee -a "$LOG_FILE" || echo "No network interface information available"
        fi
        
        echo -e "\nRouting Table:"
        if command -v ip >/dev/null 2>&1; then
            ip route show | tee -a "$LOG_FILE"
        else
            route -n 2>/dev/null | tee -a "$LOG_FILE" || echo "No routing information available"
        fi
        
        echo -e "\nDNS Configuration:"
        if [ -f /etc/resolv.conf ]; then
            cat /etc/resolv.conf | tee -a "$LOG_FILE"
        else
            echo "DNS configuration file not found" | tee -a "$LOG_FILE"
        fi
        
        echo -e "\nNetwork Connections:"
        # Try multiple methods to get listening ports
        if command -v ss >/dev/null 2>&1; then
            $SUDO_CMD ss -tuln | tee -a "$LOG_FILE"
        elif command -v netstat >/dev/null 2>&1; then
            $SUDO_CMD netstat -tuln 2>/dev/null | tee -a "$LOG_FILE"
        else
            echo "No network connection tools available (ss/netstat)" | tee -a "$LOG_FILE"
        fi
        
        echo -e "\nActive Network Connections:"
        if command -v ss >/dev/null 2>&1; then
            $SUDO_CMD ss -tupln | tee -a "$LOG_FILE"
        elif command -v netstat >/dev/null 2>&1; then
            $SUDO_CMD netstat -tupln 2>/dev/null | tee -a "$LOG_FILE"
        else
            echo "No active connection tools available" | tee -a "$LOG_FILE"
        fi
        
        echo -e "\nNetwork Statistics:"
        if [ -f /proc/net/dev ]; then
            cat /proc/net/dev | tee -a "$LOG_FILE"
        else
            echo "Network statistics not available" | tee -a "$LOG_FILE"
        fi
        
        echo -e "\nNetwork Interface Details:"
        if command -v ip >/dev/null 2>&1; then
            for iface in $(ip link show | grep -E '^[0-9]+:' | cut -d: -f2 | tr -d ' '); do
                if [ "$iface" != "lo" ]; then
                    echo "Interface: $iface" | tee -a "$LOG_FILE"
                    if command -v ethtool >/dev/null 2>&1; then
                        ethtool "$iface" 2>/dev/null | grep -E '(Speed|Duplex|Link detected)' | tee -a "$LOG_FILE" || echo "  ethtool info not available"
                    else
                        echo "  ethtool not available" | tee -a "$LOG_FILE"
                    fi
                fi
            done
        fi
        
        echo -e "\nBandwidth Usage (if available):"
        if command -v vnstat >/dev/null 2>&1; then
            # Try multiple interfaces
            for iface in eth0 eth1 wlan0; do
                if ip link show "$iface" >/dev/null 2>&1; then
                    echo "Interface $iface:" | tee -a "$LOG_FILE"
                    vnstat -i "$iface" 2>/dev/null | tee -a "$LOG_FILE" || echo "No vnstat data available for $iface"
                    break
                fi
            done
        else
            echo "vnstat not installed" | tee -a "$LOG_FILE"
        fi
    }

    print_subsection "Firewall Status"
    {
        if command -v ufw >/dev/null 2>&1; then
            echo "UFW Status:"
            $SUDO_CMD ufw status verbose | tee -a "$LOG_FILE"
        fi

        if command -v iptables >/dev/null 2>&1; then
            echo -e "\nIPTables Rules:"
            $SUDO_CMD iptables -L -n | tee -a "$LOG_FILE"
        fi

        if command -v firewall-cmd >/dev/null 2>&1; then
            echo -e "\nFirewalld Status:"
            $SUDO_CMD firewall-cmd --list-all | tee -a "$LOG_FILE"
        fi
    }
}

# Docker and Container Information
collect_container_info() {
    print_section "CONTAINER INFORMATION"

    if command -v docker >/dev/null 2>&1; then
        print_subsection "Docker Information"
        {
            echo "Docker Version:"
            docker --version | tee -a "$LOG_FILE"
            echo -e "\nDocker System Info:"
            $SUDO_CMD docker system info 2>/dev/null | tee -a "$LOG_FILE"
            echo -e "\nRunning Containers:"
            $SUDO_CMD docker ps -a | tee -a "$LOG_FILE"
            echo -e "\nDocker Images:"
            $SUDO_CMD docker images | tee -a "$LOG_FILE"
            echo -e "\nDocker Networks:"
            $SUDO_CMD docker network ls | tee -a "$LOG_FILE"
            echo -e "\nDocker Volumes:"
            $SUDO_CMD docker volume ls | tee -a "$LOG_FILE"
            
            echo -e "\nDocker Compose Services:"
            if command -v docker-compose >/dev/null 2>&1; then
                find /opt /home -name "docker-compose.yml" -o -name "docker-compose.yaml" 2>/dev/null | head -10 | tee -a "$LOG_FILE"
            fi
            
            echo -e "\nContainer Management Tools:"
            $SUDO_CMD docker ps --format "table {{.Names}}	{{.Image}}	{{.Ports}}" | grep -E "(portainer|watchtower|traefik|nginx-proxy|heimdall|dashboard)" | tee -a "$LOG_FILE" || echo "No common management tools detected"
            
            echo -e "\nContainer Resource Usage:"
            $SUDO_CMD docker stats --no-stream --format "table {{.Container}}	{{.CPUPerc}}	{{.MemUsage}}	{{.NetIO}}" 2>/dev/null | head -20 | tee -a "$LOG_FILE"
        }

        # Check for docker socket permissions
        if [ -S /var/run/docker.sock ]; then
            echo -e "\nDocker Socket Permissions:" | tee -a "$LOG_FILE"
            ls -la /var/run/docker.sock | tee -a "$LOG_FILE"
        fi
    else
        echo "Docker not installed or not in PATH" | tee -a "$LOG_FILE"
    fi

    if command -v podman >/dev/null 2>&1; then
        print_subsection "Podman Information"
        {
            echo "Podman Version:"
            podman --version | tee -a "$LOG_FILE"
            echo -e "\nPodman Containers:"
            podman ps -a | tee -a "$LOG_FILE"
            echo -e "\nPodman Images:"
            podman images | tee -a "$LOG_FILE"
        }
    fi

    # Check for container runtime files
    if [ -f "/.dockerenv" ]; then
        print_warning "This system appears to be running inside a Docker container"
    fi
}

# Software and Package Information
collect_software_info() {
    print_section "SOFTWARE INFORMATION"

    print_subsection "Installed Packages"

    # Determine package manager and list packages
    if command -v dpkg >/dev/null 2>&1; then
        echo "Installed Debian/Ubuntu packages:" | tee -a "$LOG_FILE"
        dpkg -l > "${OUTPUT_DIR}/packages_dpkg.txt"
        echo "Package list saved to packages_dpkg.txt ($(wc -l < "${OUTPUT_DIR}/packages_dpkg.txt") packages)" | tee -a "$LOG_FILE"

        # Check for security updates
        if command -v apt >/dev/null 2>&1; then
            echo -e "\nAvailable Security Updates:" | tee -a "$LOG_FILE"
            apt list --upgradable 2>/dev/null | grep -i security | tee -a "$LOG_FILE"
        fi
    fi

    if command -v rpm >/dev/null 2>&1; then
        echo "Installed RPM packages:" | tee -a "$LOG_FILE"
        rpm -qa > "${OUTPUT_DIR}/packages_rpm.txt"
        echo "Package list saved to packages_rpm.txt ($(wc -l < "${OUTPUT_DIR}/packages_rpm.txt") packages)" | tee -a "$LOG_FILE"

        # Check for security updates (Fedora/RHEL)
        if command -v dnf >/dev/null 2>&1; then
            echo -e "\nAvailable Security Updates:" | tee -a "$LOG_FILE"
            dnf check-update --security 2>/dev/null | tee -a "$LOG_FILE"
        fi
    fi

    print_subsection "Running Services"
    {
        echo "Systemd Services:"
        systemctl list-units --type=service --state=running | tee -a "$LOG_FILE"
        echo -e "\nEnabled Services:"
        systemctl list-unit-files --type=service --state=enabled | tee -a "$LOG_FILE"
    }

    print_subsection "Running Processes"
    {
        echo "Process List (top 20 by CPU):"
        ps aux --sort=-%cpu | head -20 | tee -a "$LOG_FILE"
        echo -e "\nProcess Tree:"
        pstree | tee -a "$LOG_FILE"
    }
}

# Security Assessment
collect_security_info() {
    print_section "SECURITY ASSESSMENT"

    print_subsection "User Accounts"
    {
        echo "User accounts with shell access:"
        set +o pipefail
        cat /etc/passwd | grep -E '/bin/(bash|sh|zsh|fish)' | tee -a "$LOG_FILE"
        set -o pipefail
        echo -e "\nUsers with UID 0 (root privileges):"
        awk -F: '$3 == 0 {print $1}' /etc/passwd | tee -a "$LOG_FILE"
        echo -e "\nSudo group members:"
        if [ -f /etc/group ]; then
            grep -E '^(sudo|wheel):' /etc/group | tee -a "$LOG_FILE"
        fi
        echo -e "\nCurrently logged in users:"
        who | tee -a "$LOG_FILE"
        echo -e "\nLast logins:"
        last -10 | tee -a "$LOG_FILE"
    }

    print_subsection "SSH Configuration"
    {
        if [ -f /etc/ssh/sshd_config ]; then
            echo "SSH Configuration highlights:"
            grep -E '^(Port|PermitRootLogin|PasswordAuthentication|PubkeyAuthentication|Protocol)' /etc/ssh/sshd_config | tee -a "$LOG_FILE"
        fi

        echo -e "\nSSH Failed Login Attempts (last 50):"
        $SUDO_CMD grep "Failed password" /var/log/auth.log 2>/dev/null | tail -50 | tee -a "$LOG_FILE" || echo "Auth log not accessible"
    }

    print_subsection "File Permissions and SUID"
    {
        echo "World-writable files (excluding /proc, /sys, /dev):"
        # Use parallel processing for better performance
        find / -type f -perm -002 -not -path "/proc/*" -not -path "/sys/*" -not -path "/dev/*" -not -path "/tmp/*" -not -path "/var/tmp/*" 2>/dev/null | head -20 | tee -a "$LOG_FILE"

        echo -e "\nSUID/SGID files:"
        find / -type f \( -perm -4000 -o -perm -2000 \) 2>/dev/null | head -30 | tee -a "$LOG_FILE"
        
        echo -e "\nSUID/SGID Risk Assessment:"
        # Check for dangerous SUID binaries
        local dangerous_suid=("/bin/su" "/usr/bin/sudo" "/usr/bin/passwd" "/usr/bin/chfn" "/usr/bin/chsh" "/usr/bin/gpasswd" "/usr/bin/newgrp" "/usr/bin/mount" "/usr/bin/umount" "/usr/bin/ping" "/usr/bin/ping6")
        for binary in "${dangerous_suid[@]}"; do
            if [ -f "$binary" ] && [ -u "$binary" ]; then
                echo "  WARNING: Potentially dangerous SUID binary found: $binary" | tee -a "$LOG_FILE"
            fi
        done
        
        echo -e "\nWorld-writable directories (excluding /proc, /sys, /dev, /tmp):"
        find / -type d -perm -002 -not -path "/proc/*" -not -path "/sys/*" -not -path "/dev/*" -not -path "/tmp/*" -not -path "/var/tmp/*" 2>/dev/null | head -10 | tee -a "$LOG_FILE"
    }

    print_subsection "Cron Jobs"
    {
        echo "Root crontab:"
        $SUDO_CMD crontab -l 2>/dev/null | tee -a "$LOG_FILE" || echo "No root crontab"

        echo -e "\nSystem cron jobs:"
        if [ -d /etc/cron.d ]; then
            ls -la /etc/cron.d/ | tee -a "$LOG_FILE"
        fi

        if [ -f /etc/crontab ]; then
            echo -e "\n/etc/crontab contents:"
            cat /etc/crontab | tee -a "$LOG_FILE"
        fi
    }

    print_subsection "Shell History"
    {
        echo "Checking for sensitive information in shell history..."
        local sensitive_patterns=(
            "password"
            "passwd"
            "secret"
            "token"
            "key"
            "api_key"
            "private_key"
            "ssh_key"
            "aws_access"
            "aws_secret"
            "database_url"
            "connection_string"
            "credential"
            "auth"
            "login"
        )
        
        for histfile in /home/*/.bash_history /root/.bash_history /home/*/.zsh_history /home/*/.fish_history; do
            if [ -f "$histfile" ] && [ -r "$histfile" ]; then
                echo "Analyzing: $histfile" | tee -a "$LOG_FILE"
                local found_sensitive=false
                
                for pattern in "${sensitive_patterns[@]}"; do
                    if grep -q -i "$pattern" "$histfile" 2>/dev/null; then
                        echo "  WARNING: Pattern '$pattern' found in $histfile" | tee -a "$LOG_FILE"
                        found_sensitive=true
                    fi
                done
                
                if [ "$found_sensitive" = false ]; then
                    echo "  No obvious sensitive patterns found" | tee -a "$LOG_FILE"
                fi
            fi
        done
        
        # Check for command history in memory
        if [ -f /proc/$$/environ ]; then
            echo -e "\nChecking for sensitive environment variables..."
            if grep -q -i "password\|secret\|token\|key" /proc/$$/environ 2>/dev/null; then
                print_warning "Sensitive environment variables detected in current process"
            fi
        fi
    }

    print_subsection "Tailscale Configuration"
    {
        if command -v tailscale >/dev/null 2>&1; then
            echo "Tailscale Status:"
            tailscale status 2>/dev/null | tee -a "$LOG_FILE" || echo "Tailscale not running"
            echo -e "\nTailscale IP:"
            tailscale ip -4 2>/dev/null | tee -a "$LOG_FILE" || echo "No Tailscale IP"
        else
            echo "Tailscale not installed" | tee -a "$LOG_FILE"
        fi
    }
}

# Vulnerability Assessment
run_vulnerability_scan() {
    print_section "VULNERABILITY ASSESSMENT"

    print_subsection "Kernel Vulnerabilities"
    {
        echo "Kernel version and potential CVEs:"
        uname -r | tee -a "$LOG_FILE"

        # Enhanced kernel vulnerability assessment
        kernel_version=$(uname -r)
        kernel_major=$(echo "$kernel_version" | cut -d. -f1)
        kernel_minor=$(echo "$kernel_version" | cut -d. -f2)
        
        echo "Current kernel: $kernel_version" | tee -a "$LOG_FILE"
        echo "Kernel major version: $kernel_major" | tee -a "$LOG_FILE"
        echo "Kernel minor version: $kernel_minor" | tee -a "$LOG_FILE"

        # Enhanced kernel version checking
        local risk_level="LOW"
        local risk_message=""
        
        if [ "$kernel_major" -lt 4 ]; then
            risk_level="CRITICAL"
            risk_message="Kernel version is very outdated and likely has multiple critical vulnerabilities"
        elif [ "$kernel_major" -eq 4 ] && [ "$kernel_minor" -lt 19 ]; then
            risk_level="HIGH"
            risk_message="Kernel version is outdated and may have security vulnerabilities"
        elif [ "$kernel_major" -eq 5 ] && [ "$kernel_minor" -lt 4 ]; then
            risk_level="MEDIUM"
            risk_message="Kernel version may be outdated. Consider updating for security patches"
        elif [ "$kernel_major" -eq 5 ] && [ "$kernel_minor" -lt 10 ]; then
            risk_level="LOW"
            risk_message="Kernel version is relatively recent but may have some vulnerabilities"
        else
            risk_level="LOW"
            risk_message="Kernel version is recent and likely secure"
        fi
        
        echo "Risk Level: $risk_level" | tee -a "$LOG_FILE"
        echo "Assessment: $risk_message" | tee -a "$LOG_FILE"
        
        # Check for specific known vulnerable kernel versions
        local known_vulnerable_versions=(
            "4.9.0" "4.9.1" "4.9.2" "4.9.3" "4.9.4" "4.9.5" "4.9.6" "4.9.7" "4.9.8" "4.9.9"
            "4.14.0" "4.14.1" "4.14.2" "4.14.3" "4.14.4" "4.14.5" "4.14.6" "4.14.7" "4.14.8" "4.14.9"
            "4.19.0" "4.19.1" "4.19.2" "4.19.3" "4.19.4" "4.19.5" "4.19.6" "4.19.7" "4.19.8" "4.19.9"
        )
        
        for vulnerable_version in "${known_vulnerable_versions[@]}"; do
            if [[ "$kernel_version" == "$vulnerable_version"* ]]; then
                print_warning "Kernel version $kernel_version matches known vulnerable pattern: $vulnerable_version"
                break
            fi
        done
        
        # Check kernel configuration for security features
        echo -e "\nKernel Security Features:" | tee -a "$LOG_FILE"
        if [ -f /proc/sys/kernel/randomize_va_space ]; then
            local aslr=$(cat /proc/sys/kernel/randomize_va_space)
            if [ "$aslr" -eq 2 ]; then
                echo "  ASLR (Address Space Layout Randomization): ENABLED" | tee -a "$LOG_FILE"
            else
                print_warning "ASLR is not fully enabled (value: $aslr)"
            fi
        fi
        
        if [ -f /proc/sys/kernel/dmesg_restrict ]; then
            local dmesg_restrict=$(cat /proc/sys/kernel/dmesg_restrict)
            if [ "$dmesg_restrict" -eq 1 ]; then
                echo "  Dmesg restriction: ENABLED" | tee -a "$LOG_FILE"
            else
                print_warning "Dmesg restriction is disabled"
            fi
        fi
    }

    print_subsection "Open Ports Security Check"
    {
        echo "Potentially risky open ports:"
        $SUDO_CMD netstat -tuln 2>/dev/null | awk 'NR>2 {print $4}' | sed 's/.*://' | sort -n | uniq | while read port; do
            case $port in
                21) echo "Port $port (FTP) - Consider secure alternatives" ;; 
                23) echo "Port $port (Telnet) - Insecure, use SSH instead" ;; 
                53) echo "Port $port (DNS) - Ensure properly configured" ;; 
                80) echo "Port $port (HTTP) - Consider HTTPS" ;; 
                135|139|445) echo "Port $port (SMB/NetBIOS) - Potentially risky" ;; 
                3389) echo "Port $port (RDP) - Ensure secure configuration" ;; 
            esac
        done | tee -a "$LOG_FILE"
    }
}

# Environment Variables and Configuration
collect_env_info() {
    print_section "ENVIRONMENT AND CONFIGURATION"

    print_subsection "Environment Variables"
    {
        echo "Key environment variables:"
        env | grep -E '^(PATH|HOME|USER|SHELL|LANG)=' | tee -a "$LOG_FILE"
    }

    print_subsection "Mount Points"
    {
        echo "Current mounts:"
        mount | tee -a "$LOG_FILE"
        echo -e "\nDisk usage by mount point:"
        df -h | tee -a "$LOG_FILE"
    }

    print_subsection "System Limits"
    {
        echo "System limits:"
        ulimit -a | tee -a "$LOG_FILE"
    }
}

# Generate JSON summary
generate_json_summary() {
    print_section "GENERATING SUMMARY"
    log "Generating JSON summary..."

    # Calculate execution time
    END_TIME=$(date +%s)
    EXECUTION_TIME=$((END_TIME - START_TIME))

    # Gather data safely, with fallbacks
    hostname_val=$(hostname)
    fqdn_val=$(hostname -f 2>/dev/null || echo 'unknown')
    ip_addresses_val=$(hostname -I 2>/dev/null | tr ' ' ',' || echo 'unknown')
    os_val=$(lsb_release -d 2>/dev/null | cut -f2 | tr -d '"' || cat /etc/os-release | grep PRETTY_NAME | cut -d'=' -f2 | tr -d '"')
    kernel_val=$(uname -r)
    arch_val=$(uname -m)
    uptime_val=$(uptime -p 2>/dev/null || uptime)
    docker_installed_val=$(command -v docker >/dev/null 2>&1 && echo "true" || echo "false")
    podman_installed_val=$(command -v podman >/dev/null 2>&1 && echo "true" || echo "false")
    running_containers_val=$(if command -v docker >/dev/null 2>&1; then $SUDO_CMD docker ps -q 2>/dev/null | wc -l; else echo "0"; fi)
    ssh_root_login_val=$(grep -E '^PermitRootLogin' /etc/ssh/sshd_config 2>/dev/null | awk '{print $2}' || echo 'unknown')
    ufw_status_val=$(command -v ufw >/dev/null 2>&1 && $SUDO_CMD ufw status | head -1 | awk '{print $2}' || echo 'not_installed')
    failed_ssh_attempts_val=$(grep "Failed password" /var/log/auth.log 2>/dev/null | wc -l || echo "0")
    
    # Collect open ports with better fallback
    if command -v ss >/dev/null 2>&1; then
        open_ports=$($SUDO_CMD ss -tuln 2>/dev/null | awk 'NR>2 {print $5}' | sed 's/.*://' | sort -n | uniq | tr '\n' ',' | sed 's/,$//')
    elif command -v netstat >/dev/null 2>&1; then
        open_ports=$($SUDO_CMD netstat -tuln 2>/dev/null | awk 'NR>2 {print $4}' | sed 's/.*://' | sort -n | uniq | tr '\n' ',' | sed 's/,$//')
    else
        open_ports="unknown"
    fi

    # Check for jq and use it if available, otherwise generate basic JSON
    if command -v jq >/dev/null 2>&1; then
        # Build JSON with jq
        jq -n \
          --arg timestamp "$(date -Iseconds)" \
          --arg hostname "$hostname_val" \
          --arg scan_duration "${EXECUTION_TIME}s" \
          --arg fqdn "$fqdn_val" \
          --arg ip_addresses "$ip_addresses_val" \
          --arg os "$os_val" \
          --arg kernel "$kernel_val" \
          --arg architecture "$arch_val" \
          --arg uptime "$uptime_val" \
          --argjson docker_installed "$docker_installed_val" \
          --argjson podman_installed "$podman_installed_val" \
          --arg running_containers "$running_containers_val" \
          --arg ssh_root_login "$ssh_root_login_val" \
          --arg ufw_status "$ufw_status_val" \
          --arg failed_ssh_attempts "$failed_ssh_attempts_val" \
          --arg open_ports "$open_ports" \
          '{ 
            "scan_info": {
              "timestamp": $timestamp,
              "hostname": $hostname,
              "scanner_version": "2.0",
              "scan_duration": $scan_duration
            },
            "system": {
              "hostname": $hostname,
              "fqdn": $fqdn,
              "ip_addresses": $ip_addresses,
              "os": $os,
              "kernel": $kernel,
              "architecture": $architecture,
              "uptime": $uptime
            },
            "containers": {
              "docker_installed": $docker_installed,
              "podman_installed": $podman_installed,
              "running_containers": ($running_containers | tonumber)
            },
            "security": {
              "ssh_root_login": $ssh_root_login,
              "ufw_status": $ufw_status,
              "failed_ssh_attempts": ($failed_ssh_attempts | tonumber),
              "open_ports": ($open_ports | split(","))
            }
          }' > "$RESULTS_FILE"
    else
        # Generate basic JSON without jq
        log_warn "jq not available, generating basic JSON summary"
        cat > "$RESULTS_FILE" << EOF
{
  "scan_info": {
    "timestamp": "$(date -Iseconds)",
    "hostname": "$hostname_val",
    "scanner_version": "2.0",
    "scan_duration": "${EXECUTION_TIME}s"
  },
  "system": {
    "hostname": "$hostname_val",
    "fqdn": "$fqdn_val",
    "ip_addresses": "$ip_addresses_val",
    "os": "$os_val",
    "kernel": "$kernel_val",
    "architecture": "$arch_val",
    "uptime": "$uptime_val"
  },
  "containers": {
    "docker_installed": $docker_installed_val,
    "podman_installed": $podman_installed_val,
    "running_containers": $running_containers_val
  },
  "security": {
    "ssh_root_login": "$ssh_root_login_val",
    "ufw_status": "$ufw_status_val",
    "failed_ssh_attempts": $failed_ssh_attempts_val,
    "open_ports": ["$open_ports"]
  }
}
EOF
    fi

    if [ $? -eq 0 ]; then
        log_info "JSON summary generated successfully: $RESULTS_FILE"
    else
        print_error "Failed to generate JSON summary."
        return 1
    fi
}

# Network scan function (optional - can be run from central location)
network_discovery() {
    if [ "$1" = "--network-scan" ]; then
        print_section "NETWORK DISCOVERY"

        if command -v nmap >/dev/null 2>&1; then
            local_subnet=$(ip route | grep "src $(hostname -I | awk '{print $1}')" | awk '{print $1}' | head -1)
            if [ ! -z "$local_subnet" ]; then
                echo "Scanning local subnet: $local_subnet" | tee -a "$LOG_FILE"
                nmap -sn "$local_subnet" > "${OUTPUT_DIR}/network_scan.txt" 2>&1
                echo "Network scan results saved to network_scan.txt" | tee -a "$LOG_FILE"
            fi
        else
            echo "nmap not available for network scanning" | tee -a "$LOG_FILE"
        fi
    fi
}

# Main execution
main() {
    log_info "Starting comprehensive system audit on $(hostname)"
    log_info "Output directory: $OUTPUT_DIR"
    log_info "Script version: 2.0"

    # Validate environment first
    validate_environment
    
    # Check privileges
    check_privileges
    
    # Run audit modules with error handling
    local modules=(
        "collect_system_info"
        "collect_network_info"
        "collect_container_info"
        "collect_software_info"
        "collect_security_info"
        "run_vulnerability_scan"
        "collect_env_info"
    )
    
    for module in "${modules[@]}"; do
        log_info "Running module: $module"
        if ! $module; then
            log_warn "Module $module encountered issues, continuing..."
        fi
    done
    
    # Run network discovery if requested
    if [ $# -gt 0 ]; then
        log_info "Running network discovery with args: $*"
        network_discovery "$@"
    fi
    
    # Generate JSON summary
    log_info "Generating JSON summary"
    if ! generate_json_summary; then
        log_warn "JSON summary generation failed, but continuing..."
    fi

    print_section "AUDIT COMPLETE"
    END_TIME=$(date +%s)
    EXECUTION_TIME=$((END_TIME - START_TIME))
    log_info "Audit completed successfully in ${EXECUTION_TIME} seconds"
    log_info "Results available in: $OUTPUT_DIR"

    # Create an enhanced summary file
    create_enhanced_summary

    # Display final information
    echo -e "\n${GREEN}Audit Complete!${NC}"
    echo -e "${BLUE}Results directory: $OUTPUT_DIR${NC}"
    echo -e "${BLUE}Main log file: $LOG_FILE${NC}"
    echo -e "${BLUE}JSON summary: $RESULTS_FILE${NC}"

    # Compress results with better error handling
    compress_results
}

# Create enhanced summary file
create_enhanced_summary() {
    local summary_file="${OUTPUT_DIR}/SUMMARY.txt"
    
    cat > "$summary_file" << EOF
=== COMPREHENSIVE AUDIT SUMMARY ===
Generated: $(date)
Script Version: 2.0
Hostname: $(hostname)
FQDN: $(hostname -f 2>/dev/null || echo 'N/A')
IP Addresses: $(hostname -I 2>/dev/null || echo 'N/A')

=== SYSTEM INFORMATION ===
OS: $(lsb_release -d 2>/dev/null | cut -f2 || cat /etc/os-release | grep PRETTY_NAME | cut -d'=' -f2 | tr -d '"')
Kernel: $(uname -r)
Architecture: $(uname -m)
Uptime: $(uptime -p 2>/dev/null || uptime)

=== SECURITY STATUS ===
SSH Root Login: $(grep -E '^PermitRootLogin' /etc/ssh/sshd_config 2>/dev/null | awk '{print $2}' || echo 'unknown')
UFW Status: $(command -v ufw >/dev/null 2>&1 && $SUDO_CMD ufw status | head -1 | awk '{print $2}' || echo 'not_installed')
Failed SSH Attempts: $(grep "Failed password" /var/log/auth.log 2>/dev/null | wc -l || echo "0")

=== CONTAINER STATUS ===
Docker: $(command -v docker >/dev/null 2>&1 && echo "Installed" || echo "Not installed")
Podman: $(command -v podman >/dev/null 2>&1 && echo "Installed" || echo "Not installed")
Running Containers: $(if command -v docker >/dev/null 2>&1; then $SUDO_CMD docker ps -q 2>/dev/null | wc -l; else echo "0"; fi)

=== FILES GENERATED ===
EOF

    ls -la "$OUTPUT_DIR" >> "$summary_file"
    
    log_info "Enhanced summary created: $summary_file"
}

# Compress results with better error handling
compress_results() {
    if ! command -v tar >/dev/null 2>&1; then
        log_warn "tar not available, skipping compression"
        return 0
    fi
    
    log_info "Compressing audit results..."
    
    local parent_dir=$(dirname "$OUTPUT_DIR")
    local archive_name=$(basename "$OUTPUT_DIR").tar.gz
    
    if cd "$parent_dir" && tar -czf "$archive_name" "$(basename "$OUTPUT_DIR")"; then
        log_info "Compressed results: $parent_dir/$archive_name"
        
        # Verify archive integrity
        if tar -tzf "$archive_name" >/dev/null 2>&1; then
            log_info "Archive verified successfully"
            
            # Calculate archive size
            local archive_size=$(du -h "$archive_name" | cut -f1)
            log_info "Archive size: $archive_size"
        else
            log_error "Archive verification failed"
            return 1
        fi
    else
        log_error "Compression failed"
        return 1
    fi
}

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