#!/bin/bash
#
# Fast Comprehensive Discovery Script
# Optimized version that avoids filesystem bottlenecks
#

set -euo pipefail

# --- Configuration ---
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
HOSTNAME=$(hostname -f)
OUTPUT_BASE_DIR="/tmp/system_audit_${HOSTNAME}_${TIMESTAMP}"
DISCOVERY_DIR="${OUTPUT_BASE_DIR}/discovery"
mkdir -p "$DISCOVERY_DIR"
LOG_FILE="${OUTPUT_BASE_DIR}/discovery.log"

# --- Logging ---
exec > >(tee -a "$LOG_FILE") 2>&1
echo "Starting Fast Comprehensive Discovery on ${HOSTNAME} at $(date)"
echo "Output will be saved in ${OUTPUT_BASE_DIR}"
echo "-----------------------------------------------------"

# --- Helper Functions ---
print_header() {
    echo ""
    echo "====================================================="
    echo ">= $1"
    echo "====================================================="
}

run_command() {
    local title="$1"
    local command="$2"
    local output_file="$3"
    
    print_header "$title"
    echo "Running command: $command"
    echo "Outputting to: $output_file"
    
    if eval "$command" > "$output_file" 2>/dev/null; then
        echo "Successfully captured $title."
    else
        echo "Warning: Command for '$title' failed or produced no output." > "$output_file"
    fi
}

# --- 1. Infrastructure Discovery ---
infra_discovery() {
    local out_dir="${DISCOVERY_DIR}/1_infrastructure"
    mkdir -p "$out_dir"
    
    run_command "OS Information" "cat /etc/os-release" "${out_dir}/os_info.txt"
    run_command "CPU Information" "lscpu" "${out_dir}/cpu_info.txt"
    run_command "Memory Information" "free -h" "${out_dir}/memory_info.txt"
    run_command "Disk Usage" "df -h" "${out_dir}/disk_usage.txt"
    run_command "Storage Layout" "lsblk -a" "${out_dir}/storage_layout.txt"
    run_command "Network Interfaces" "ip addr show" "${out_dir}/network_interfaces.txt"
    run_command "Routing Table" "ip route" "${out_dir}/routing_table.txt"
    run_command "DNS Configuration" "cat /etc/resolv.conf" "${out_dir}/dns_config.txt"
    run_command "PCI Devices" "lspci" "${out_dir}/pci_devices.txt"
    run_command "USB Devices" "lsusb" "${out_dir}/usb_devices.txt"
    run_command "RAID Status" "cat /proc/mdstat || echo 'No software RAID'" "${out_dir}/raid_status.txt"
    run_command "Firewall (iptables)" "sudo iptables -L -n" "${out_dir}/firewall_iptables.txt"
    run_command "Firewall (UFW)" "sudo ufw status verbose" "${out_dir}/firewall_ufw.txt"
}

# --- 2. Services & Container Discovery ---
services_inventory() {
    local out_dir="${DISCOVERY_DIR}/2_services"
    mkdir -p "$out_dir"
    
    # Docker Discovery (optimized)
    if command -v docker >/dev/null 2>&1; then
        run_command "Docker Information" "docker info" "${out_dir}/docker_info.txt"
        run_command "Docker Images" "docker images" "${out_dir}/docker_images.txt"
        run_command "Docker Containers (All)" "docker ps -a" "${out_dir}/docker_ps.txt"
        run_command "Docker Networks" "docker network ls" "${out_dir}/docker_networks.txt"
        run_command "Docker Volumes" "docker volume ls" "${out_dir}/docker_volumes.txt"
        
        print_header "Docker Container Details"
        for id in $(docker ps -q); do
            local name=$(docker inspect --format '{{.Name}}' "$id" | sed 's,^/,,')
            echo "Inspecting container: $name"
            docker inspect "$id" > "${out_dir}/container_${name}.json"
        done
        
        print_header "Finding Docker Compose files (optimized search)"
        # Search only common Docker Compose locations instead of entire filesystem
        {
            find /home -name "docker-compose.yml" -o -name "docker-compose.yaml" -o -name "compose.yml" 2>/dev/null || true
            find /opt -name "docker-compose.yml" -o -name "docker-compose.yaml" -o -name "compose.yml" 2>/dev/null || true
            find /srv -name "docker-compose.yml" -o -name "docker-compose.yaml" -o -name "compose.yml" 2>/dev/null || true
            find /etc -name "docker-compose.yml" -o -name "docker-compose.yaml" -o -name "compose.yml" 2>/dev/null || true
            find /root -name "docker-compose.yml" -o -name "docker-compose.yaml" -o -name "compose.yml" 2>/dev/null || true
        } > "${out_dir}/docker_compose_locations.txt"
        
        while IFS= read -r file; do
            if [[ -n "$file" ]]; then
                sudo cp "$file" "${out_dir}/compose_file_$(basename "$(dirname "$file")").yml" 2>/dev/null || true
            fi
        done < "${out_dir}/docker_compose_locations.txt"
    else
        echo "Docker not found." > "${out_dir}/docker_status.txt"
    fi
    
    # Systemd Services
    run_command "Systemd Services (Enabled)" "systemctl list-unit-files --state=enabled" "${out_dir}/systemd_enabled_services.txt"
    run_command "Systemd Services (Running)" "systemctl list-units --type=service --state=running" "${out_dir}/systemd_running_services.txt"
}

# --- 3. Data & Storage Discovery ---
data_discovery() {
    local out_dir="${DISCOVERY_DIR}/3_data_storage"
    mkdir -p "$out_dir"
    
    run_command "NFS Exports" "showmount -e localhost || echo 'No NFS exports'" "${out_dir}/nfs_exports.txt"
    run_command "Mounted File Systems" "mount" "${out_dir}/mounts.txt"
    
    print_header "Searching for critical data directories (optimized)"
    # Target specific directories instead of full filesystem search
    {
        find /etc/postgresql -name "postgresql.conf" 2>/dev/null || true
        find /var/lib/postgresql -name "postgresql.conf" 2>/dev/null || true
    } > "${out_dir}/postgres_locations.txt"
    
    {
        find /etc/mysql -name "my.cnf" 2>/dev/null || true
        find /etc -name "my.cnf" 2>/dev/null || true
    } > "${out_dir}/mysql_locations.txt"
    
    find /var/lib -name "*.db" 2>/dev/null | head -20 > "${out_dir}/sqlite_locations.txt" || true
    
    # Common data directories (limit depth to avoid long searches)
    {
        ls -la /srv 2>/dev/null || true
        ls -la /mnt 2>/dev/null || true
        ls -la /opt 2>/dev/null || true
    } > "${out_dir}/common_data_dirs.txt"
}

# --- 4. Security & Access Discovery ---
security_discovery() {
    local out_dir="${DISCOVERY_DIR}/4_security"
    mkdir -p "$out_dir"
    
    run_command "User Accounts" "cat /etc/passwd" "${out_dir}/users.txt"
    run_command "Sudoers Configuration" "sudo cat /etc/sudoers" "${out_dir}/sudoers.txt"
    run_command "SSH Daemon Configuration" "sudo cat /etc/ssh/sshd_config" "${out_dir}/sshd_config.txt"
    run_command "Last Logins" "last -a | head -50" "${out_dir}/last_logins.txt"
    run_command "Open Ports" "sudo ss -tuln" "${out_dir}/open_ports.txt"
    run_command "Cron Jobs (System)" "sudo cat /etc/crontab || echo 'No system crontab'" "${out_dir}/crontab_system.txt"
    
    print_header "User Cron Jobs"
    {
        for user in $(cut -f1 -d: /etc/passwd); do 
            crontab -u "$user" -l 2>/dev/null | sed "s/^/[$user] /" || true
        done
    } > "${out_dir}/crontab_users.txt"
}

# --- 5. Performance & Usage (30-second sample) ---
performance_discovery() {
    local out_dir="${DISCOVERY_DIR}/5_performance"
    mkdir -p "$out_dir"
    
    run_command "Current Processes" "ps aux" "${out_dir}/processes.txt"
    run_command "Uptime & Load" "uptime" "${out_dir}/uptime.txt"
    run_command "Network Stats" "netstat -s || ss -s" "${out_dir}/netstat.txt"
    
    print_header "Collecting 30-second performance sample"
    if command -v iostat >/dev/null 2>&1; then
        iostat -x 1 30 > "${out_dir}/iostat.txt" &
        IOSTAT_PID=$!
    fi
    
    # Collect system stats every 5 seconds for 30 seconds
    {
        for i in {1..6}; do
            echo "=== Sample $i/6 at $(date) ==="
            echo "Load: $(uptime)"
            echo "Memory: $(free -m | grep '^Mem:')"
            echo "CPU: $(grep 'cpu ' /proc/stat)"
            sleep 5
        done
    } > "${out_dir}/system_samples.txt"
    
    # Wait for iostat to complete
    if [[ -n "${IOSTAT_PID:-}" ]]; then
        wait $IOSTAT_PID 2>/dev/null || true
    fi
}

# --- Main Execution ---
main() {
    infra_discovery
    services_inventory
    data_discovery
    security_discovery
    performance_discovery
    
    print_header "Packaging Results"
    cd "$(dirname "$OUTPUT_BASE_DIR")"
    tar -czf "${OUTPUT_BASE_DIR}.tar.gz" "$(basename "$OUTPUT_BASE_DIR")"
    
    echo "-----------------------------------------------------"
    echo "Fast comprehensive discovery complete."
    echo "Results packaged in ${OUTPUT_BASE_DIR}.tar.gz"
    echo "Log file: $LOG_FILE"
}

main