#!/bin/bash
#
# Automated kernel patch workflow with embedded metadata
# 
# Usage:
#   ./script.sh <kernel_version> <patch_version> [--debug]
#   ./script.sh 4.18.0-513.9.1.el8_9 v1
#   ./script.sh 4.18.0-513.9.1.el8_9 v1 --debug
#

set -e  # Exit on any error
set -u  # Exit on undefined variables
set -o pipefail  # Catch errors in pipes

# Trap errors and show where they occurred
trap 'error "Script failed at line $LINENO. Exit code: $?"' ERR

# ============================================================================
# CONFIGURATION
# ============================================================================

# Check if patch version is provided
if [[ $# -lt 2 ]]; then
    echo "Usage: $0 <kernel_version> <patch_version> [--debug]"
    echo ""
    echo "Example:"
    echo "  $0 4.18.0-513.9.1.el8_9 v1"
    echo "  $0 4.18.0-513.9.1.el8_9 v2"
    echo "  $0 4.18.0-513.9.1.el8_9 v1 --debug"
    echo ""
    exit 1
fi

VERSION="$1"
PATCH_VERSION="$2"
DEBUG_MODE=false

# Check for debug flag
if [[ $# -ge 3 ]] && [[ "$3" == "--debug" ]]; then
    DEBUG_MODE=true
    set -x  # Enable command tracing
fi

# Validate patch version format (v1, v2, v3, etc.)
if [[ ! "${PATCH_VERSION}" =~ ^v[0-9]+$ ]]; then
    echo "ERROR: Patch version must be in format: v1, v2, v3, etc."
    echo "You provided: ${PATCH_VERSION}"
    exit 1
fi

# Base directories
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
WORK_DIR="${SCRIPT_DIR}/build_${VERSION}_${PATCH_VERSION}"
PATCH_DIR="${SCRIPT_DIR}/selected_patches_${VERSION}"
TARBALL="${SCRIPT_DIR}/linux-${VERSION}.tar.xz"

# Work directory structure
SOURCE_DIR="${WORK_DIR}/source"
TARGET_DIR="${WORK_DIR}/target"
BASE_DIR="/root/rebootless"
CUMULATIVE_PATCH="${WORK_DIR}/cumulative-${VERSION}-${PATCH_VERSION}.patch"
MODULE_OUTPUT="${WORK_DIR}"

# ============================================================================
# UTILITY FUNCTIONS
# ============================================================================

log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"
}

error() {
    echo "[ERROR] $*" >&2
    exit 1
}

check_file() {
    local file="$1"
    [[ -f "$file" ]] || error "Required file not found: $file"
}

check_dir() {
    local dir="$1"
    [[ -d "$dir" ]] || error "Required directory not found: $dir"
}

check_command() {
    local cmd="$1"
    command -v "$cmd" >/dev/null 2>&1 || error "Required command not found: $cmd"
}

# ============================================================================
# METADATA GENERATION
# ============================================================================

generate_module_metadata_source() {
    local version="$1"
    local patch_version="$2"
    local patch_count="$3"
    local cve_list="$4"
    local metadata_file="${TARGET_DIR}/kpatch_metadata.c"
    
    log "Generating module metadata source file..."
    
    cat > "${metadata_file}" <<EOF
#include <linux/module.h>

MODULE_INFO(kernel_version, "${version}");
MODULE_INFO(patch_version, "${patch_version}");
MODULE_INFO(build_date, "$(date +%Y-%m-%d)");
MODULE_INFO(build_timestamp, "$(date -Iseconds)");
MODULE_INFO(patch_count, "${patch_count}");
MODULE_INFO(builder, "${USER}@$(hostname)");

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Rebootless Patching Team");
MODULE_DESCRIPTION("Cumulative kernel security livepatch for ${version} ${patch_version}");
MODULE_VERSION("${patch_version}");
EOF
    
    log "Created metadata source: ${metadata_file}"
}

create_metadata_companion_file() {
    local ko_file="$1"
    local patch_version="$2"
    local patch_count="$3"
    local applied="$4"
    local failed="$5"
    local line_count="$6"
    
    local metadata_txt="${ko_file%.ko}.metadata.txt"
    
    log "Creating companion metadata file..."
    
    cat > "${metadata_txt}" <<EOF
╔══════════════════════════════════════════════════════════════════════════╗
║                    Livepatch Module Metadata                             ║
╚══════════════════════════════════════════════════════════════════════════╝

MODULE INFORMATION:
  Module File:        $(basename ${ko_file})
  Kernel Version:     ${VERSION}
  Patch Version:      ${patch_version}
  Module Name:        $(basename ${ko_file} .ko)

BUILD INFORMATION:
  Build Date:         $(date)
  Build Timestamp:    $(date -Iseconds)
  Build Host:         $(hostname)
  Build User:         ${USER}
  Working Directory:  ${WORK_DIR}

PATCH STATISTICS:
  Total Patches:      ${patch_count}
  Successfully Applied: ${applied}
  Failed:             ${failed}
  Cumulative Patch:   $(basename ${CUMULATIVE_PATCH})
  Patch Size:         ${line_count} lines

SOURCE INFORMATION:
  Source Directory:   ${SOURCE_DIR}
  Original Tarball:   $(basename ${TARBALL})

PATCH LIST:
$(find "${PATCH_DIR}" -name "*.patch" -printf "  - %f\n" 2>/dev/null | sort)

USAGE:
  # Load the livepatch
  sudo insmod $(basename ${ko_file})
  
  # Check if loaded
  sudo kpatch list
  
  # View module info
  modinfo $(basename ${ko_file})
  
  # Unload the livepatch
  sudo kpatch unload $(basename ${ko_file} .ko)

Generated by: $(basename $0)
Working directory preserved at: ${WORK_DIR}
EOF
    
    log "Metadata file created: ${metadata_txt}"
}

# ============================================================================
# MAIN WORKFLOW
# ============================================================================

main() {
    log "╔══════════════════════════════════════════════════════════════════════════╗"
    log "║           Kernel Patch Workflow - Rebootless Patching                    ║"
    log "╚══════════════════════════════════════════════════════════════════════════╝"
    log ""
    log "Kernel Version: ${VERSION}"
    log "Patch Version:  ${PATCH_VERSION}"
    log "Work Directory: ${WORK_DIR}"
    log ""
    
    # Step 0: Check prerequisites
    log "Checking prerequisites..."
    check_file "${TARBALL}"
    check_dir "${PATCH_DIR}"
    check_command "tar"
    check_command "patch"
    check_command "diff"
    check_command "make"
    check_command "kpatch-build"
    
    # Count patches
    local patch_count=$(find "${PATCH_DIR}" -name "*.patch" 2>/dev/null | wc -l)
    log "Found ${patch_count} patches in ${PATCH_DIR}"
    [[ ${patch_count} -gt 0 ]] || error "No patches found in ${PATCH_DIR}"
    
    # Step 1: Create work directory
    log "Creating work directory..."
    if [[ -d "${WORK_DIR}" ]]; then
        log "WARNING: ${WORK_DIR} already exists"
        read -p "Remove existing directory? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            rm -rf "${WORK_DIR}"
        else
            error "Work directory exists. Please remove it or use a different version."
        fi
    fi
    mkdir -p "${WORK_DIR}"
    cd "${WORK_DIR}" || error "Cannot cd to ${WORK_DIR}"
    
    # Step 2: Extract tarball to create base kernel source
    log "Extracting kernel source from ${TARBALL}..."
    tar -xf "${TARBALL}" || error "Failed to extract tarball"
    
    # Find the extracted directory (it might have a different name)
    local extracted_dir=$(find . -maxdepth 1 -type d -name "linux-*" | head -1)
    [[ -n "${extracted_dir}" ]] || error "Could not find extracted kernel directory"
    
    log "Found extracted directory: ${extracted_dir}"
    
    # Step 3: Create source and target directories
    log "Creating source and target directories..."
    mv "${extracted_dir}" source || error "Failed to create source directory"
    cp -a source target || error "Failed to create target directory"
    
    check_dir "${SOURCE_DIR}"
    check_dir "${TARGET_DIR}"
    
    # Step 4: Apply patches to target directory
    log "Applying patches to ${TARGET_DIR}..."
    cd "${TARGET_DIR}" || error "Cannot cd to ${TARGET_DIR}"
    
    local applied=0
    local failed=0
    local skipped=0
    local failed_patches=()
    local total_patches=${patch_count}
    local current=0
    
    # Temporarily disable exit on error for patch application
    set +e
    
    for patch_file in "${PATCH_DIR}"/*.patch; do
        [[ -f "${patch_file}" ]] || continue
        
        ((current++))
        local patch_name=$(basename "${patch_file}")
        
        # Show progress every 10 patches or for important patches
        if [[ $((current % 10)) -eq 0 ]] || [[ $current -eq 1 ]] || [[ $current -eq $total_patches ]]; then
            log "  Progress: ${current}/${total_patches} - Applying: ${patch_name}"
        fi
        
        # Try to apply patch
        local patch_output=$(patch -p1 --forward < "${patch_file}" 2>&1)
        local patch_status=$?
        
        if [[ ${patch_status} -eq 0 ]]; then
            ((applied++))
        else
            # Check if patch was already applied (reverse check)
            if patch -p1 --reverse --dry-run < "${patch_file}" >/dev/null 2>&1; then
                ((skipped++))
            else
                # Check if it's a fuzz/offset issue (partial success)
                if echo "${patch_output}" | grep -qi "succeeded\|applied"; then
                    ((applied++))
                    log "  INFO: ${patch_name} applied with fuzz/offset"

                else
                    ((failed++))
                    failed_patches+=("${patch_name}")
                    log "  WARNING: Failed to apply ${patch_name}"
                fi
            fi
        fi
    done
    
    # Re-enable exit on error
    set -e
    
    log "Patch application complete:"
    log "  Applied: ${applied}"
    log "  Skipped: ${skipped} (already applied)"
    log "  Failed:  ${failed}"
    
    # Save detailed patch log
    local patch_log="${WORK_DIR}/patch_application-${PATCH_VERSION}.log"
    {
        echo "Patch Application Log"
        echo "===================="
        echo "Kernel Version: ${VERSION}"
        echo "Patch Version: ${PATCH_VERSION}"
        echo "Date: $(date)"
        echo ""
        echo "Summary:"
        echo "  Total patches: ${total_patches}"
        echo "  Applied: ${applied}"
        echo "  Skipped: ${skipped}"
        echo "  Failed: ${failed}"
        echo ""
        if [[ ${#failed_patches[@]} -gt 0 ]]; then
            echo "Failed patches:"
            for fp in "${failed_patches[@]}"; do
                echo "  - ${fp}"
            done
        fi
    } > "${patch_log}"
    
    log "Patch application log saved to: ${patch_log}"
    
    if [[ ${failed} -gt 0 ]]; then
        log "WARNING: Failed patches:"
        for fp in "${failed_patches[@]}"; do
            log "  - ${fp}"
        done
    fi
    
    # Step 5: Check for .orig and .rej files
    log "Checking for .orig and .rej files..."
    local orig_count=$(find "${TARGET_DIR}" -name "*.orig" 2>/dev/null | wc -l)
    local rej_count=$(find "${TARGET_DIR}" -name "*.rej" 2>/dev/null | wc -l)
    
    log "  Found ${orig_count} .orig files"
    log "  Found ${rej_count} .rej files"
    
    if [[ ${rej_count} -gt 0 ]]; then
        log "WARNING: Patch rejections found:"
        find "${TARGET_DIR}" -name "*.rej" | while read f; do log "  - $f"; done
    fi
    
    # Step 6: Clean up .orig files
    if [[ ${orig_count} -gt 0 ]]; then
        log "Removing .orig files..."
        find "${TARGET_DIR}" -name "*.orig" -delete
    fi
    
    # Step 7: Add module metadata to target
    cd "${WORK_DIR}" || error "Cannot cd to ${WORK_DIR}"
    
    generate_module_metadata_source "${VERSION}" "${PATCH_VERSION}" "${applied}" "none"
    
    # Step 8: Generate cumulative patch
    log "Generating cumulative patch..."
    diff -urN source/ target/ > "${CUMULATIVE_PATCH}" || true
    
    local line_count=$(wc -l < "${CUMULATIVE_PATCH}")
    log "Cumulative patch created: ${line_count} lines"
    
    if [[ ${line_count} -eq 0 ]]; then
        error "Cumulative patch is empty! No changes detected."
    fi
    
    # Step 9: Validate cumulative patch
    log "Validating cumulative patch..."
    
    if grep -qE "\.orig|\.rej" "${CUMULATIVE_PATCH}"; then
        log "WARNING: Found .orig or .rej references in cumulative patch!"
        grep -E "\.orig|\.rej" "${CUMULATIVE_PATCH}" | head -5
    fi
    
    local init_count=$(grep -c "__init" "${CUMULATIVE_PATCH}" 2>/dev/null || true)
    local struct_count=$(grep -cE "^[+-]struct .*\{" "${CUMULATIVE_PATCH}" 2>/dev/null || true)
    
    log "  __init occurrences: ${init_count}"
    log "  struct definitions: ${struct_count}"
    
    if [[ ${init_count} -gt 0 ]] || [[ ${struct_count} -gt 0 ]]; then
        log "WARNING: Patch contains __init or struct definitions"
        log "         This may not be suitable for kpatch!"
    fi
    
    # Step 9b: Test cumulative patch applies cleanly
    if [[ "${TEST_CUMULATIVE_BEFORE_BUILD:-1}" == "1" ]]; then
        log "Testing cumulative patch applies cleanly..."
        
        local test_dir="${WORK_DIR}/test_apply"
        rm -rf "$test_dir"
        mkdir -p "$test_dir"
        
        cp -a source/* "$test_dir/"
        
        if patch -p1 --forward --dry-run < "${CUMULATIVE_PATCH}" > "$test_dir/patch_test.log" 2>&1; then
            log "  ✓ Cumulative patch applies cleanly (dry-run)"
        else
            local failed_hunks=$(grep -c "FAILED" "$test_dir/patch_test.log" 2>/dev/null || echo "0")
            if [[ "$failed_hunks" -gt 0 ]]; then
                log "WARNING: Cumulative patch has $failed_hunks failed hunks!"
                log "  See: $test_dir/patch_test.log"
                log "Continuing anyway (kpatch may handle this)..."
            else
                log "  ✓ Cumulative patch applies with offsets/fuzz (dry-run)"
            fi
        fi
        
        rm -rf "$test_dir"
    fi
    
    # Step 10: Build vmlinux in source directory
    log "Building vmlinux in source directory..."
    cd "${SOURCE_DIR}" || error "Cannot cd to ${SOURCE_DIR}"
    
    # Check if .config exists
    if [[ ! -f .config ]]; then
        log "Copying running kernel config..."
        if [[ -f /boot/config-$(uname -r) ]]; then
            cp /boot/config-$(uname -r) .config
        elif [[ -f /proc/config.gz ]]; then
            zcat /proc/config.gz > .config
        else
            error "No kernel config found. Please provide .config file."
        fi
    fi
    
    log "Preparing kernel build (this may take a while)..."
    sudo openssl req -new -x509 -nodes -days 3650 -subj "/CN=kernel-build-rhel/" -out certs/rhel.pem
    make olddefconfig >/dev/null 2>&1 || error "make olddefconfig failed"
    
    log "Building vmlinux (this will take several minutes)..."
    make vmlinux -j$(nproc) || error "vmlinux build failed"
    
    check_file "${SOURCE_DIR}/vmlinux"
    log "vmlinux built successfully: ${SOURCE_DIR}/vmlinux"
    
    # Step 11: Build kpatch module
    cd "${WORK_DIR}" || error "Cannot cd to ${WORK_DIR}"
    
    log "Building kpatch module..."
    
    local build_date=$(date +%Y%m%d)
    # Replace dots with underscores for modinfo name
    local version_safe=$(echo "${VERSION}" | tr '.' '_')
    local module_name="RL-${version_safe}-${build_date}-${PATCH_VERSION}"
    
    log "Module name: RL-${VERSION}-${build_date}-${PATCH_VERSION}"
    log "Safe name:   ${module_name}"
    
    kpatch-build \
        --skip-compiler-check \
        --non-replace \
        --name "${module_name}" \
        --sourcedir "${SOURCE_DIR}" \
        --vmlinux "${SOURCE_DIR}/vmlinux" \
        "${CUMULATIVE_PATCH}" || error "kpatch-build failed!"
    
    log "kpatch module built successfully!"
    
    # Step 12: Find and rename the module
    local ko_file=$(find "${WORK_DIR}" -name "${module_name}.ko" -type f 2>/dev/null | head -1)
    
    if [[ -z "${ko_file}" ]]; then
        # Fallback: find any .ko file
        ko_file=$(find "${WORK_DIR}" -name "*.ko" -type f 2>/dev/null | head -1)
    fi
    
    if [[ -n "${ko_file}" ]]; then
        # Rename the module 
        local new_name="RL-${VERSION}-${build_date}-${PATCH_VERSION}.ko"
        local new_ko_file="${WORK_DIR}/${new_name}"
        
        if [[ "${ko_file}" != "${new_ko_file}" ]]; then
            log "Renaming module to use dots in version..."
            mv "${ko_file}" "${new_ko_file}"
            ko_file="${new_ko_file}"
        fi
        
        # Rollback: backup previous livepatch modules
        local rollback_dir="${BASE_DIR}/rollback"
        mkdir -p "${rollback_dir}"
        
        # Find and backup previous modules for this kernel
        local prev_modules=($(find "${rollback_dir}" -name "RL-${VERSION}-*.ko" 2>/dev/null || true))
        if [[ ${#prev_modules[@]} -gt 0 ]]; then
            log "Found ${#prev_modules[@]} previous module(s) in rollback directory"
        fi
        
        # Copy current module to rollback dir
        local backup_name="RL-${VERSION}-${build_date}-${PATCH_VERSION}.ko"
        cp "${ko_file}" "${rollback_dir}/${backup_name}"
        log "Saved module to rollback: ${rollback_dir}/${backup_name}"
        
        # Keep only last 5 versions
        local kept=0
        for old_mod in $(ls -t "${rollback_dir}"/RL-"${VERSION}"-*.ko 2>/dev/null | tail -n +6); do
            rm -f "$old_mod" 2>/dev/null || true
            ((kept++)) || true
        done
        if [[ $kept -gt 0 ]]; then
            log "Cleaned up $kept old rollback module(s)"
        fi
        
        log "═══════════════════════════════════════════════════════════"
        log "Generated kpatch module: ${ko_file}"
        ls -lh "${ko_file}"
        
        log "═══════════════════════════════════════════════════════════"
        log "Module information (modinfo):"
        modinfo "${ko_file}" 2>/dev/null || log "  (modinfo not available)"
        
        # Create companion metadata file
        create_metadata_companion_file "${ko_file}" "${PATCH_VERSION}" "${patch_count}" "${applied}" "${failed}" "${line_count}"
        
        log "═══════════════════════════════════════════════════════════"
    else
        log "WARNING: Could not locate generated .ko file"
    fi

    # Remove previous source and target directories
    rm -rf "${SOURCE_DIR}"
    rm -rf "${TARGET_DIR}"

    # Copy CVE index for the current version/patch
    cp "${BASE_DIR}/rebootlessdata/${VERSION}_cve_index.json" \
        "${WORK_DIR}/${VERSION}_cve_index-${PATCH_VERSION}.json"
    
    log "Workflow complete!"
    
    # Summary
    log "═══════════════════════════════════════════════════════════"
    log "SUMMARY:"
    log "  Kernel Version:    ${VERSION}"
    log "  Patch Version:     ${PATCH_VERSION}"
    log "  Patches Applied:   ${applied}/${patch_count}"
    log "  Module Name:       RL-${VERSION}-${build_date}-${PATCH_VERSION}"
    log "  Module File:       ${ko_file}"
    log "  Cumulative Patch:  ${CUMULATIVE_PATCH}"
    log "  Work Directory:    ${WORK_DIR}"
    log ""
    log "Files preserved in: ${WORK_DIR}"
    log "  - source/          : Original kernel source + built vmlinux"
    log "  - target/          : Patched kernel source"
    log "  - cumulative patch : ${CUMULATIVE_PATCH}"
    log "  - module           : ${ko_file}"
    log "═══════════════════════════════════════════════════════════"
}

# ============================================================================
# SCRIPT ENTRY POINT
# ============================================================================

if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi