#!/bin/bash
#
# Ubuntu Kernel CVE Finder - Like yum updateinfo for RHEL
# Finds CVEs affecting Ubuntu kernels (kernel 5.4.x series)
#

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
DATA_DIR="${DATA_DIR:-rebootlessdata}"
CACHE_DIR="${DATA_DIR}/cve_cache"
PATCH_DIR="${DATA_DIR}/kernel_cve_patches"

mkdir -p "$DATA_DIR" "$CACHE_DIR" "$PATCH_DIR"

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

log_info() { echo -e "${BLUE}[INFO]${NC} $*"; }
log_success() { echo -e "${GREEN}[SUCCESS]${NC} $*"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $*"; }
log_error() { echo -e "${RED}[ERROR]${NC} $*" >&2; }

KERNEL_VERSION="${1:-$(uname -r)}"
KERNEL_BASE=$(echo "$KERNEL_VERSION" | cut -d'-' -f1)

main() {
    echo ""
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║       Ubuntu Kernel CVE Finder (yum updateinfo equiv) ║"
    echo "╚══════════════════════════════════════════════════════════╝"
    echo ""
    
    log_info "Kernel Version: $KERNEL_VERSION"
    log_info "Kernel Base: $KERNEL_BASE"
    echo ""
    
    # Get full kernel version from dpkg
    local full_version=$(dpkg -l linux-image-${KERNEL_VERSION} 2>/dev/null | awk '/^ii/ {print $3}')
    
    if [[ -z "$full_version" ]]; then
        full_version="${KERNEL_VERSION}.33"
    fi
    
    log_info "Full version: $full_version"
    
    local all_cves=""
    
    # Method 1: Kernel changelog CVEs (what's fixed in this kernel)
    # This is the PRIMARY source - tells us exactly which CVEs affect this kernel version
    log_info "[1/3] Parsing kernel changelog (primary source)..."
    local changelog_url="https://changelogs.ubuntu.com/changelogs/pool/main/l/linux/linux_${full_version}/changelog"
    
    local changelog_cves=$(curl -sL --connect-timeout 15 --max-time 60 "$changelog_url" 2>/dev/null | \
                  grep -oE 'CVE-[0-9]{4}-[0-9]+' | sort -u || echo "")
    
    local changelog_count=$(echo "$changelog_cves" | grep -c 'CVE' || echo "0")
    changelog_count=${changelog_count//[^0-9]/}
    log_success "Found $changelog_count CVEs from kernel changelog"
    
    # Add changelog CVEs as primary source
    all_cves+=$'\n'"$changelog_cves"
    
    # Method 2: Ubuntu Security Notices (recent CVEs)
    log_info "[2/3] Fetching Ubuntu Security Notices..."
    local usn_cves=$(curl -sL --connect-timeout 30 --max-time 60 "https://ubuntu.com/security/notices/rss.xml" 2>/dev/null | \
                     grep -oE 'CVE-[0-9]{4}-[0-9]+' | sort -u || echo "")
    local usn_count=$(echo "$usn_cves" | grep -c 'CVE' || echo "0")
    usn_count=${usn_count//[^0-9]/}
    log_success "Found $usn_count CVEs from Ubuntu Security Notices"
    all_cves+=$'\n'"$usn_cves"
    
    # Method 3: Check for patches directory CVEs (from previous downloads)
    log_info "[4/4] Fetching KernelCare CVEs..."
    local patch_cves=$(find "$PATCH_DIR" -maxdepth 1 -type d -name 'CVE-*' 2>/dev/null | \
                       xargs -I{} basename {} 2>/dev/null | sort -u || echo "")
    local patch_count
    patch_count=$(echo "$patch_cves" | grep -c 'CVE' || echo "0")
    patch_count=${patch_count//[^0-9]/}
    if [[ -n "$patch_count" && "$patch_count" -gt 0 ]]; then
        log_success "Found $patch_count CVEs from existing patches"
        all_cves+=$'\n'"$patch_cves"
    fi
    
    # Method 4: KernelCare CVEs
    log_info "[4/4] Fetching KernelCare CVEs..."
    local kc_cves=""
    local kc_count=0
    
    # Map Ubuntu version to KernelCare distro
    local ubuntu_version=$(lsb_release -r -s 2>/dev/null || echo "")
    local kc_distro=""
    case "$ubuntu_version" in
        24.04) kc_distro="ubuntu-noble" ;;
        22.04) kc_distro="ubuntu-jammy" ;;
        20.04) kc_distro="ubuntu-focal" ;;
        18.04) kc_distro="ubuntu-bionic" ;;
        *) ;;
    esac
    
    if [[ -n "$kc_distro" ]]; then
        log_info "KernelCare distro: $kc_distro"
        
        # Fetch index data
        local index_json=$(curl -sL --connect-timeout 15 --max-time 60 "https://patches.kernelcare.com/indexdata.json?v=1" 2>/dev/null || echo "")
        
        if [[ -n "$index_json" && "${#index_json}" -gt 100 ]]; then
            # Find hash for our kernel
            local kernel_search="$KERNEL_BASE"
            local kc_hash=$(echo "$index_json" | python3 -c "
import sys, json
data = json.load(sys.stdin)
for v in data.get('versions', []):
    if v.get('product') != 'kernel': continue
    boot = v.get('boot_version', '')
    if '$kc_distro' == v.get('distro', '') and '$kernel_search' in boot:
        print(v.get('id', ''))
        break
" 2>/dev/null || echo "")
            
            if [[ -n "$kc_hash" ]]; then
                log_info "Found KernelCare hash: $kc_hash"
                
                # Get latest level
                local kc_level=$(echo "$index_json" | python3 -c "
import sys, json
data = json.load(sys.stdin)
for v in data.get('versions', []):
    if v.get('id') == '$kc_hash':
        print(v.get('latest', 1))
        break
" 2>/dev/null || echo "1")
                
                # Fetch CVEs from kpatch.html
                local kc_html=$(curl -sL --connect-timeout 15 --max-time 60 "https://patches.kernelcare.com/${kc_hash}/${kc_level}/kpatch.html" 2>/dev/null || echo "")
                
                if [[ -n "$kc_html" && "${#kc_html}" -gt 500 ]]; then
                    kc_cves=$(echo "$kc_html" | grep -oE 'CVE-[0-9]{4}-[0-9]+' | sort -u || echo "")
                    kc_count=$(echo "$kc_cves" | grep -c 'CVE' || echo "0")
                    kc_count=${kc_count//[^0-9]/}
                    log_success "Found $kc_count CVEs from KernelCare"
                    all_cves+=$'\n'"$kc_cves"
                fi
            fi
        fi
    else
        log_warn "Could not determine Ubuntu version for KernelCare"
    fi
    
    # Deduplicate and save
    echo ""
    log_info "Merging and deduplicating CVEs..."
    
    local all_cves_sorted=$(echo "$all_cves" | grep -E '^CVE-' | sort -u)
    local total_count
    total_count=$(echo "$all_cves_sorted" | grep -c 'CVE' || echo "0")
    total_count=${total_count//[^0-9]/}
    
    # Save files
    echo "$all_cves_sorted" > "$DATA_DIR/cves_all_${KERNEL_VERSION}.txt"
    echo "$changelog_cves" > "$CACHE_DIR/fixed_cves_${KERNEL_VERSION}.txt"
    echo "$all_cves_sorted" > "$DATA_DIR/cves_${KERNEL_VERSION}.txt"
    echo "$all_cves_sorted" > "$DATA_DIR/cves_needed_for_livepatch.txt"
    
    echo ""
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║                    Results                               ║"
    echo "╠══════════════════════════════════════════════════════════╣"
    printf "║  %-50s ║\n" "CVEs from kernel changelog: $changelog_count"
    printf "║  %-50s ║\n" "CVEs from Ubuntu USN: $usn_count"
    printf "║  %-50s ║\n" "CVEs from KernelCare: $kc_count"
    printf "║  %-50s ║\n" "Total unique CVEs: $total_count"
    printf "║  %-50s ║\n" "CVEs needing patches: $total_count"
    echo "╠══════════════════════════════════════════════════════════╣"
    printf "║  %-50s ║\n" "Output: rebootlessdata/cves_${KERNEL_VERSION}.txt"
    echo "╚══════════════════════════════════════════════════════════╝"
    
    return 0
}

main "$@"