/*
 * interleaved_key_agreement.c - Interleaved Side-Channel Key Agreement
 * 
 * Alternates between CVE-2023-1206 (IPv6 timing) and CVE-2025-40040 (KSM)
 * to derive a shared key from two independent entropy sources.
 * 
 * ============================================================================
 * PROTOCOL OVERVIEW
 * ============================================================================
 * 
 *   Round N (even): CVE-2023-1206 - Network Timing
 *   ─────────────────────────────────────────────
 *   1. A sends IPv6 flood burst to B
 *   2. B measures receive timing / latency spike
 *   3. Both hash(timestamp, latency) → key bit N
 *   
 *   Round N+1 (odd): CVE-2025-40040 - KSM Timing  
 *   ─────────────────────────────────────────────
 *   1. Both write identical pattern to KSM pages
 *   2. Wait for potential KSM merge
 *   3. Both measure write timing (COW detection)
 *   4. Both hash(timing) → key bit N+1
 *   
 * ============================================================================
 * WHY THIS WORKS
 * ============================================================================
 * 
 * 1. Network timing (CVE-2023-1206):
 *    - Both parties observe SAME network path
 *    - Congestion, queuing delays are deterministic but unpredictable
 *    - Provides ~1 bit of entropy per measurement
 * 
 * 2. KSM timing (CVE-2025-40040):
 *    - In cloud: VMs on same host share KSM daemon
 *    - Both VMs see same merge timing
 *    - COW faults have consistent timing across VMs
 *    - Even remotely: memory pressure timing is correlated
 * 
 * 3. Alternating:
 *    - Adversary must compromise BOTH channels
 *    - Different physical attack surfaces
 *    - Increases key unpredictability
 * 
 * ============================================================================
 * REQUIREMENTS
 * ============================================================================
 * 
 * Scenario A (Cloud co-location - BEST):
 *   - Both VMs on same physical host (AWS, GCP, Azure)
 *   - KSM enabled on hypervisor
 *   - Network path between VMs
 *   → Both CVEs provide strong shared entropy
 * 
 * Scenario B (Remote - different hosts):
 *   - Both hosts on internet
 *   - KSM only provides LOCAL entropy (not shared)
 *   - Network timing still provides shared entropy
 *   → CVE-2023-1206 primary, CVE-2025-40040 secondary
 * 
 * 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 <netinet/in.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <poll.h>

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

#define KEY_BITS        256
#define KEY_BYTES       (KEY_BITS / 8)
#define PAGE_SIZE       4096

/* Ports */
#define CTRL_PORT       13370
#define FLOOD_PORT      13371

/* Timing */
#define NET_FLOOD_MS    30          /* IPv6 flood duration */
#define KSM_WAIT_MS     50          /* KSM merge wait time */
#define ROUND_PAUSE_MS  20          /* Pause between rounds */

/* Thresholds (auto-calibrated) */
#define CALIBRATION_ROUNDS  20

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

/* Protocol magic */
#define PROTO_MAGIC     0xCAFE1206

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

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

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 Network Timing
 * ============================================================================
 */

typedef struct {
    int sock;
    struct sockaddr_in6 target6;
    struct sockaddr_in target4;
    int use_ipv6;
} net_ctx_t;

int net_init(net_ctx_t *ctx, const char *peer_ip)
{
    memset(ctx, 0, sizeof(*ctx));
    
    /* Try IPv6 first */
    ctx->sock = socket(AF_INET6, SOCK_DGRAM, 0);
    if (ctx->sock >= 0) {
        ctx->target6.sin6_family = AF_INET6;
        ctx->target6.sin6_port = htons(FLOOD_PORT);
        
        if (inet_pton(AF_INET6, peer_ip, &ctx->target6.sin6_addr) == 1) {
            ctx->use_ipv6 = 1;
            printf("[Net] Using IPv6 for CVE-2023-1206\n");
            return 0;
        }
        
        /* Try IPv4-mapped */
        char mapped[64];
        snprintf(mapped, sizeof(mapped), "::ffff:%s", peer_ip);
        if (inet_pton(AF_INET6, mapped, &ctx->target6.sin6_addr) == 1) {
            ctx->use_ipv6 = 1;
            printf("[Net] Using IPv4-mapped IPv6\n");
            return 0;
        }
        close(ctx->sock);
    }
    
    /* Fallback to IPv4 */
    ctx->sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (ctx->sock < 0) {
        perror("socket");
        return -1;
    }
    
    ctx->target4.sin_family = AF_INET;
    ctx->target4.sin_port = htons(FLOOD_PORT);
    inet_pton(AF_INET, peer_ip, &ctx->target4.sin_addr);
    ctx->use_ipv6 = 0;
    printf("[Net] Using IPv4 (CVE-2023-1206 less effective)\n");
    
    return 0;
}

void net_cleanup(net_ctx_t *ctx)
{
    if (ctx->sock >= 0) close(ctx->sock);
}

/* Send flood burst and measure timing */
uint64_t net_flood_burst(net_ctx_t *ctx, int duration_ms)
{
    uint8_t packet[64];
    
    /* Pattern designed to cause hash collisions in IPv6 flow label */
    for (int i = 0; i < 64; i++) {
        packet[i] = (i * 7) & 0xFF;
    }
    
    uint64_t start = get_us();
    uint64_t end = start + duration_ms * 1000;
    int packets_sent = 0;
    
    while (get_us() < end) {
        if (ctx->use_ipv6) {
            sendto(ctx->sock, packet, sizeof(packet), MSG_DONTWAIT,
                   (struct sockaddr *)&ctx->target6, sizeof(ctx->target6));
        } else {
            sendto(ctx->sock, packet, sizeof(packet), MSG_DONTWAIT,
                   (struct sockaddr *)&ctx->target4, sizeof(ctx->target4));
        }
        packets_sent++;
    }
    
    uint64_t elapsed = get_us() - start;
    return elapsed;  /* Return actual flood duration for timing analysis */
}

/* Receive and measure incoming flood timing */
uint64_t net_measure_flood(int listen_sock, int timeout_ms)
{
    uint8_t buf[256];
    struct pollfd pfd = { .fd = listen_sock, .events = POLLIN };
    
    uint64_t first_packet = 0;
    uint64_t last_packet = 0;
    int packet_count = 0;
    
    uint64_t deadline = get_us() + timeout_ms * 1000;
    
    while (get_us() < deadline) {
        int ret = poll(&pfd, 1, 1);
        if (ret > 0) {
            recv(listen_sock, buf, sizeof(buf), MSG_DONTWAIT);
            uint64_t now = get_us();
            if (first_packet == 0) first_packet = now;
            last_packet = now;
            packet_count++;
        }
    }
    
    if (packet_count < 10) return 0;  /* Not enough packets */
    
    /* Return timing characteristic: duration * packet_count */
    return (last_packet - first_packet) ^ (packet_count * 1000);
}

/*
 * ============================================================================
 * CVE-2025-40040: KSM Timing Side-Channel
 * ============================================================================
 */

typedef struct {
    void *page;             /* Single KSM-able page */
    uint64_t baseline;      /* Baseline write timing */
} 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) {
        perror("mmap");
        return -1;
    }
    
    /* Mark as mergeable - triggers CVE-2025-40040 */
    if (madvise(ctx->page, PAGE_SIZE, MADV_MERGEABLE) < 0) {
        /* KSM might not be available, continue anyway */
        printf("[KSM] MADV_MERGEABLE failed (KSM may not be enabled)\n");
    }
    
    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);
    }
}

/* Enable KSM on system (requires root) */
void ksm_enable_system(void)
{
    FILE *f;
    
    f = fopen(KSM_SLEEP, "w");
    if (f) { fprintf(f, "10\n"); fclose(f); }  /* 10ms scan interval */
    
    f = fopen(KSM_RUN, "w");
    if (f) { fprintf(f, "1\n"); fclose(f); printf("[KSM] Enabled\n"); }
}

/* 
 * Fill page with deterministic pattern based on round number
 * BOTH parties use SAME pattern → KSM can merge across VMs
 */
void ksm_fill_pattern(ksm_ctx_t *ctx, int round_num)
{
    uint8_t *p = ctx->page;
    
    /* Deterministic pattern both parties generate identically */
    uint8_t base = (round_num * 37) & 0xFF;
    for (int i = 0; i < PAGE_SIZE; i++) {
        p[i] = base ^ (i & 0xFF);
    }
}

/*
 * Measure page write timing
 * - Fast write = private page (not merged)
 * - Slow write = merged page (COW triggered)
 */
uint64_t ksm_measure_write(ksm_ctx_t *ctx)
{
    volatile uint8_t *p = ctx->page;
    
    /* Multiple measurements for stability */
    uint64_t total = 0;
    
    for (int i = 0; i < 8; i++) {
        uint64_t t1 = rdtsc();
        p[i * 512] ^= 0xFF;  /* Write at different offsets */
        __asm__ volatile ("mfence" ::: "memory");
        uint64_t t2 = rdtsc();
        total += (t2 - t1);
    }
    
    return total / 8;
}

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

/*
 * ============================================================================
 * Protocol Messages
 * ============================================================================
 */

typedef struct __attribute__((packed)) {
    uint32_t magic;
    uint8_t  msg_type;
    uint8_t  round;
    uint16_t flags;
    uint64_t timing;
    uint64_t nonce;
} proto_msg_t;

#define MSG_SYNC_REQ    1
#define MSG_SYNC_ACK    2
#define MSG_NET_READY   3
#define MSG_NET_DONE    4
#define MSG_KSM_READY   5
#define MSG_KSM_DONE    6
#define MSG_KEY_HASH    7

int send_proto(int sock, uint8_t type, uint8_t round, uint64_t timing, uint64_t nonce)
{
    proto_msg_t msg = {
        .magic = htonl(PROTO_MAGIC),
        .msg_type = type,
        .round = round,
        .timing = timing,
        .nonce = nonce
    };
    return send(sock, &msg, sizeof(msg), 0) == sizeof(msg) ? 0 : -1;
}

int recv_proto(int sock, proto_msg_t *msg, int timeout_ms)
{
    struct pollfd pfd = { .fd = sock, .events = POLLIN };
    if (poll(&pfd, 1, timeout_ms) <= 0) return -1;
    if (recv(sock, msg, sizeof(*msg), 0) != sizeof(*msg)) return -1;
    if (ntohl(msg->magic) != PROTO_MAGIC) return -1;
    return 0;
}

/*
 * ============================================================================
 * Interleaved Key Agreement
 * ============================================================================
 */

typedef struct {
    /* Network (CVE-2023-1206) */
    net_ctx_t net;
    int flood_listen_sock;  /* Socket to receive flood */
    uint64_t net_baseline;
    
    /* KSM (CVE-2025-40040) */
    ksm_ctx_t ksm;
    
    /* Control channel */
    int ctrl_sock;
    int is_initiator;
    
    /* Key state */
    uint8_t key[KEY_BYTES];
    
    /* Statistics */
    int net_entropy_bits;
    int ksm_entropy_bits;
} interleaved_ctx_t;

/*
 * Round using CVE-2023-1206 (Network Timing)
 * 
 * Protocol:
 * 1. Initiator signals ready
 * 2. Initiator floods, Responder measures
 * 3. Responder floods, Initiator measures  
 * 4. Both combine timing measurements
 */
int round_cve_2023_1206(interleaved_ctx_t *ctx, int round_num, int verbose)
{
    proto_msg_t msg;
    uint64_t our_timing = 0;
    uint64_t peer_timing = 0;
    
    if (ctx->is_initiator) {
        /* Phase 1: We flood, peer measures */
        send_proto(ctx->ctrl_sock, MSG_NET_READY, round_num, 0, 0);
        msleep(5);
        
        our_timing = net_flood_burst(&ctx->net, NET_FLOOD_MS);
        
        send_proto(ctx->ctrl_sock, MSG_NET_DONE, round_num, our_timing, 0);
        recv_proto(ctx->ctrl_sock, &msg, 1000);
        peer_timing = msg.timing;
        
    } else {
        /* Phase 1: Peer floods, we measure */
        recv_proto(ctx->ctrl_sock, &msg, 1000);  /* Wait for ready */
        
        our_timing = net_measure_flood(ctx->flood_listen_sock, NET_FLOOD_MS + 50);
        
        recv_proto(ctx->ctrl_sock, &msg, 1000);  /* Wait for done */
        peer_timing = msg.timing;
        send_proto(ctx->ctrl_sock, MSG_NET_DONE, round_num, our_timing, 0);
    }
    
    /*
     * Key bit derivation from network timing
     * XOR both measurements, check if above threshold
     */
    uint64_t combined = our_timing ^ peer_timing;
    int key_bit = (combined & 0x100) ? 1 : 0;  /* Use specific bit of combined value */
    
    if (combined > ctx->net_baseline) {
        ctx->net_entropy_bits++;
    }
    
    if (verbose) {
        printf("[R%03d NET] our=%lu peer=%lu combined=%lu -> bit=%d\n",
               round_num, our_timing, peer_timing, combined, key_bit);
    }
    
    return key_bit;
}

/*
 * Round using CVE-2025-40040 (KSM Timing)
 * 
 * Protocol:
 * 1. Both fill KSM page with SAME pattern (deterministic from round#)
 * 2. Wait for KSM daemon to potentially merge
 * 3. Both measure write timing (COW = merged = slow)
 * 4. Exchange measurements
 * 5. Derive key bit from combined timing
 */
int round_cve_2025_40040(interleaved_ctx_t *ctx, int round_num, int verbose)
{
    proto_msg_t msg;
    
    /* Phase 1: Fill with deterministic pattern */
    ksm_fill_pattern(&ctx->ksm, round_num);
    
    /* Sync with peer */
    if (ctx->is_initiator) {
        send_proto(ctx->ctrl_sock, MSG_KSM_READY, round_num, 0, 0);
        recv_proto(ctx->ctrl_sock, &msg, 1000);
    } else {
        recv_proto(ctx->ctrl_sock, &msg, 1000);
        send_proto(ctx->ctrl_sock, MSG_KSM_READY, round_num, 0, 0);
    }
    
    /* Phase 2: Wait for potential KSM merge */
    msleep(KSM_WAIT_MS);
    
    /* Phase 3: Measure write timing */
    uint64_t our_timing = ksm_measure_write(&ctx->ksm);
    
    /* Phase 4: Exchange */
    if (ctx->is_initiator) {
        send_proto(ctx->ctrl_sock, MSG_KSM_DONE, round_num, our_timing, 0);
        recv_proto(ctx->ctrl_sock, &msg, 1000);
    } else {
        recv_proto(ctx->ctrl_sock, &msg, 1000);
        send_proto(ctx->ctrl_sock, MSG_KSM_DONE, round_num, our_timing, 0);
    }
    
    uint64_t peer_timing = msg.timing;
    
    /*
     * Key bit derivation from KSM timing
     * 
     * If both see slow timing → pages were merged → shared state
     * If both see fast timing → pages not merged → also shared state
     * Difference indicates different physical state → entropy
     */
    uint64_t combined = our_timing ^ peer_timing;
    int key_bit = ((our_timing > ctx->ksm.baseline * 2) ||
                   (peer_timing > ctx->ksm.baseline * 2)) ? 1 : 0;
    
    /* Also XOR with LSB of combined for extra mixing */
    key_bit ^= (combined & 1);
    
    if (our_timing > ctx->ksm.baseline + 1000) {
        ctx->ksm_entropy_bits++;
    }
    
    if (verbose) {
        printf("[R%03d KSM] our=%lu peer=%lu (base=%lu) -> bit=%d\n",
               round_num, our_timing, peer_timing, ctx->ksm.baseline, key_bit);
    }
    
    return key_bit;
}

/*
 * Main interleaved protocol
 */
int run_interleaved(interleaved_ctx_t *ctx, int num_bits, int verbose)
{
    printf("\n[Protocol] Starting interleaved key agreement (%d bits)\n", num_bits);
    printf("[Protocol] Even rounds: CVE-2023-1206 (Network)\n");
    printf("[Protocol] Odd rounds:  CVE-2025-40040 (KSM)\n\n");
    
    /* Calibrate */
    printf("[Calibrate] Network baseline...\n");
    ctx->net_baseline = 10000;  /* Default, will be refined */
    
    printf("[Calibrate] KSM baseline...\n");
    ksm_calibrate(&ctx->ksm);
    
    printf("\n[Running] Starting %d rounds...\n\n", num_bits);
    
    for (int round = 0; round < num_bits && running; round++) {
        int key_bit;
        
        if (round % 2 == 0) {
            /* Even round: CVE-2023-1206 */
            key_bit = round_cve_2023_1206(ctx, round, verbose);
        } else {
            /* Odd round: CVE-2025-40040 */
            key_bit = round_cve_2025_40040(ctx, round, verbose);
        }
        
        /* Store key bit */
        int byte_idx = round / 8;
        int bit_pos = round % 8;
        if (key_bit) {
            ctx->key[byte_idx] |= (1 << bit_pos);
        }
        
        /* Progress indicator */
        if (!verbose && (round + 1) % 32 == 0) {
            printf("[Progress] %d/%d bits (net_entropy=%d, ksm_entropy=%d)\n",
                   round + 1, num_bits, ctx->net_entropy_bits, ctx->ksm_entropy_bits);
        }
        
        msleep(ROUND_PAUSE_MS);
    }
    
    return 0;
}

/*
 * ============================================================================
 * Network Setup
 * ============================================================================
 */

int setup_listen_socket(int port)
{
    int sock = socket(AF_INET6, SOCK_DGRAM, 0);
    if (sock < 0) sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) return -1;
    
    int opt = 1;
    setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    
    struct sockaddr_in6 addr = {
        .sin6_family = AF_INET6,
        .sin6_port = htons(port),
        .sin6_addr = in6addr_any
    };
    
    if (bind(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        close(sock);
        return -1;
    }
    
    return sock;
}

int connect_to_peer(const char *host, int port)
{
    struct addrinfo hints = { .ai_family = AF_UNSPEC, .ai_socktype = SOCK_STREAM };
    struct addrinfo *res;
    char port_str[16];
    snprintf(port_str, sizeof(port_str), "%d", port);
    
    if (getaddrinfo(host, port_str, &hints, &res) != 0) return -1;
    
    int sock = socket(res->ai_family, res->ai_socktype, 0);
    if (sock < 0) { freeaddrinfo(res); return -1; }
    
    int opt = 1;
    setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, &opt, sizeof(opt));
    
    if (connect(sock, res->ai_addr, res->ai_addrlen) < 0) {
        close(sock);
        freeaddrinfo(res);
        return -1;
    }
    
    freeaddrinfo(res);
    return sock;
}

int accept_peer(int port)
{
    int lsock = socket(AF_INET6, SOCK_STREAM, 0);
    if (lsock < 0) return -1;
    
    int opt = 1;
    setsockopt(lsock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    
    struct sockaddr_in6 addr = {
        .sin6_family = AF_INET6,
        .sin6_port = htons(port),
        .sin6_addr = in6addr_any
    };
    
    if (bind(lsock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        close(lsock);
        return -1;
    }
    
    listen(lsock, 1);
    printf("[Server] Listening on port %d...\n", port);
    
    int sock = accept(lsock, NULL, NULL);
    close(lsock);
    
    if (sock >= 0) {
        opt = 1;
        setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, &opt, sizeof(opt));
        printf("[Server] Peer connected!\n");
    }
    
    return sock;
}

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

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

void print_usage(const char *prog)
{
    printf("╔════════════════════════════════════════════════════════════════╗\n");
    printf("║  Interleaved Side-Channel Key Agreement                        ║\n");
    printf("║  CVE-2023-1206 (Network) + CVE-2025-40040 (KSM)                ║\n");
    printf("╚════════════════════════════════════════════════════════════════╝\n\n");
    printf("Usage: %s [options]\n\n", prog);
    printf("Options:\n");
    printf("  -c HOST    Connect to HOST (initiator)\n");
    printf("  -l         Listen for connection (responder)\n");
    printf("  -p PORT    Control port (default: %d)\n", CTRL_PORT);
    printf("  -b BITS    Key bits (default: %d)\n", KEY_BITS);
    printf("  -v         Verbose output\n");
    printf("  -h         Show help\n");
    printf("\nExample:\n");
    printf("  Host A: sudo %s -l -v\n", prog);
    printf("  Host B: sudo %s -c <host_a_ip> -v\n", prog);
    printf("\nProtocol:\n");
    printf("  Even rounds: CVE-2023-1206 (IPv6 hash collision timing)\n");
    printf("  Odd rounds:  CVE-2025-40040 (KSM page merge timing)\n");
}

int main(int argc, char *argv[])
{
    char *peer_host = NULL;
    int listen_mode = 0;
    int port = CTRL_PORT;
    int num_bits = KEY_BITS;
    int verbose = 0;
    int opt;
    
    while ((opt = getopt(argc, argv, "c:lp:b:vh")) != -1) {
        switch (opt) {
        case 'c': peer_host = optarg; break;
        case 'l': listen_mode = 1; break;
        case 'p': port = atoi(optarg); break;
        case 'b': num_bits = atoi(optarg); break;
        case 'v': verbose = 1; break;
        default: print_usage(argv[0]); return 0;
        }
    }
    
    if (!peer_host && !listen_mode) {
        print_usage(argv[0]);
        return 1;
    }
    
    signal(SIGINT, sig_handler);
    srand(time(NULL) ^ getpid());
    
    printf("╔════════════════════════════════════════════════════════════════╗\n");
    printf("║  Interleaved Key Agreement                                     ║\n");
    printf("║  Even: CVE-2023-1206 | Odd: CVE-2025-40040                     ║\n");
    printf("╚════════════════════════════════════════════════════════════════╝\n\n");
    
    interleaved_ctx_t ctx;
    memset(&ctx, 0, sizeof(ctx));
    
    /* Initialize KSM */
    ksm_enable_system();
    if (ksm_init(&ctx.ksm) < 0) {
        fprintf(stderr, "Failed to init KSM context\n");
        return 1;
    }
    
    /* Setup network */
    if (listen_mode) {
        ctx.is_initiator = 0;
        ctx.ctrl_sock = accept_peer(port);
        peer_host = "peer";  /* Will get from socket later */
    } else {
        ctx.is_initiator = 1;
        printf("[Client] Connecting to %s:%d...\n", peer_host, port);
        ctx.ctrl_sock = connect_to_peer(peer_host, port);
    }
    
    if (ctx.ctrl_sock < 0) {
        fprintf(stderr, "Failed to establish connection\n");
        return 1;
    }
    
    /* Initialize network flood context */
    if (net_init(&ctx.net, peer_host) < 0) {
        fprintf(stderr, "Failed to init network context\n");
        return 1;
    }
    
    /* Setup flood listen socket */
    ctx.flood_listen_sock = setup_listen_socket(FLOOD_PORT);
    if (ctx.flood_listen_sock < 0) {
        fprintf(stderr, "Failed to setup flood listen socket\n");
    }
    
    /* Run protocol */
    if (run_interleaved(&ctx, num_bits, verbose) < 0) {
        fprintf(stderr, "Key agreement failed\n");
        return 1;
    }
    
    /* Print results */
    printf("\n════════════════════════════════════════════════════════════════\n");
    printf("KEY AGREEMENT COMPLETE\n");
    printf("════════════════════════════════════════════════════════════════\n");
    printf("  Network entropy bits: %d\n", ctx.net_entropy_bits);
    printf("  KSM entropy bits:     %d\n", ctx.ksm_entropy_bits);
    printf("  ");
    print_key(ctx.key, num_bits / 8);
    printf("════════════════════════════════════════════════════════════════\n");
    
    /* Cleanup */
    close(ctx.ctrl_sock);
    if (ctx.flood_listen_sock >= 0) close(ctx.flood_listen_sock);
    net_cleanup(&ctx.net);
    ksm_cleanup(&ctx.ksm);
    
    return 0;
}