/*
 * remote_key_agreement.c - Remote Key Agreement via Network Timing
 * 
 * Uses CVE-2023-1206 (IPv6 hash collision) timing side-channel for
 * establishing a shared secret between two hosts over the network.
 * 
 * Protocol:
 * 1. Both parties connect to each other (or a shared server)
 * 2. For each key bit, both parties simultaneously:
 *    - Generate a random bit locally
 *    - If bit=1: flood the network (cause congestion)
 *    - If bit=0: stay quiet
 * 3. Both measure network latency during each round
 * 4. High latency = at least one party sent 1
 * 5. Low latency = both sent 0
 * 6. Use XOR of measurements as shared entropy
 * 
 * This creates shared randomness because:
 * - If both flood: high latency (both know: likely both sent 1)
 * - If one floods: medium latency (uncertainty)
 * - If neither floods: low latency (both know: both sent 0)
 * 
 * The timing measurements become correlated, providing shared entropy
 * even without revealing individual bits directly.
 *
 * Alternative protocol (more robust):
 * - Use NIST randomness beacon as seed
 * - Use network RTT measurements to derive additional entropy
 * - Combine with Diffie-Hellman for hybrid approach
 *
 * 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/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <poll.h>

#define KEY_BITS 256
#define KEY_BYTES (KEY_BITS / 8)
#define SYNC_PORT 31337
#define DATA_PORT 31338
#define FLOOD_PORT 31339

#define ROUND_DURATION_MS 100
#define FLOOD_PACKETS 1000
#define LATENCY_THRESHOLD_US 500  /* Threshold for detecting flood */

static volatile int running = 1;

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

static inline uint64_t get_us(void)
{
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return ts.tv_sec * 1000000ULL + ts.tv_nsec / 1000;
}

/*
 * ============================================================================
 * Network Timing Measurement
 * ============================================================================
 */

/* Measure RTT to remote host */
uint64_t measure_rtt(int sock)
{
    uint8_t ping = 0xAA;
    uint8_t pong;
    
    uint64_t t1 = get_us();
    
    if (send(sock, &ping, 1, 0) != 1) return 0;
    if (recv(sock, &pong, 1, 0) != 1) return 0;
    
    uint64_t t2 = get_us();
    return t2 - t1;
}

/* Send flood packets to cause congestion */
void send_flood(const char *target_ip, int duration_ms)
{
    int sock = socket(AF_INET6, SOCK_DGRAM, 0);
    if (sock < 0) {
        sock = socket(AF_INET, SOCK_DGRAM, 0);
        if (sock < 0) return;
        
        struct sockaddr_in addr = {
            .sin_family = AF_INET,
            .sin_port = htons(FLOOD_PORT),
        };
        inet_pton(AF_INET, target_ip, &addr.sin_addr);
        
        uint8_t junk[64];
        memset(junk, 0x41, sizeof(junk));
        
        uint64_t end_time = get_us() + duration_ms * 1000;
        while (get_us() < end_time) {
            for (int i = 0; i < FLOOD_PACKETS; i++) {
                sendto(sock, junk, sizeof(junk), 0, 
                       (struct sockaddr *)&addr, sizeof(addr));
            }
        }
        close(sock);
        return;
    }
    
    /* IPv6 flood for CVE-2023-1206 */
    struct sockaddr_in6 addr6 = {
        .sin6_family = AF_INET6,
        .sin6_port = htons(FLOOD_PORT),
    };
    inet_pton(AF_INET6, target_ip, &addr6.sin6_addr);
    
    uint8_t junk[64];
    memset(junk, 0x41, sizeof(junk));
    
    uint64_t end_time = get_us() + duration_ms * 1000;
    while (get_us() < end_time) {
        for (int i = 0; i < FLOOD_PACKETS; i++) {
            sendto(sock, junk, sizeof(junk), 0,
                   (struct sockaddr *)&addr6, sizeof(addr6));
        }
    }
    
    close(sock);
}

/*
 * ============================================================================
 * Key Agreement Protocol
 * ============================================================================
 */

typedef struct {
    int ctrl_sock;           /* Control socket for coordination */
    char peer_ip[64];        /* Peer IP address */
    uint8_t key[KEY_BYTES];  /* Derived key */
    int is_initiator;        /* Are we the initiator? */
} key_ctx_t;

/* Synchronize with peer for next round */
int sync_round(key_ctx_t *ctx, int round)
{
    uint32_t msg = htonl(round);
    uint32_t peer_round;
    
    if (ctx->is_initiator) {
        send(ctx->ctrl_sock, &msg, sizeof(msg), 0);
        recv(ctx->ctrl_sock, &peer_round, sizeof(peer_round), 0);
    } else {
        recv(ctx->ctrl_sock, &peer_round, sizeof(peer_round), 0);
        send(ctx->ctrl_sock, &msg, sizeof(msg), 0);
    }
    
    return (ntohl(peer_round) == (uint32_t)round) ? 0 : -1;
}

/* Exchange our random bit (encrypted with temporary value) */
int exchange_commitment(key_ctx_t *ctx, int our_bit, int *peer_bit)
{
    /* 
     * Commitment scheme:
     * 1. Generate random nonce
     * 2. Send hash(bit || nonce)
     * 3. After both committed, reveal bit and nonce
     * 4. Verify commitments
     * 
     * Simplified for demo: just exchange bits directly
     * (In real implementation, use proper commitment)
     */
    uint8_t msg = our_bit ? 1 : 0;
    uint8_t peer_msg;
    
    if (ctx->is_initiator) {
        send(ctx->ctrl_sock, &msg, 1, 0);
        recv(ctx->ctrl_sock, &peer_msg, 1, 0);
    } else {
        recv(ctx->ctrl_sock, &peer_msg, 1, 0);
        send(ctx->ctrl_sock, &msg, 1, 0);
    }
    
    *peer_bit = peer_msg ? 1 : 0;
    return 0;
}

/* Single round of key agreement */
int key_round(key_ctx_t *ctx, int bit_index, int verbose)
{
    /* Generate random bit */
    int our_bit = rand() & 1;
    
    /* Sync with peer */
    if (sync_round(ctx, bit_index) < 0) {
        fprintf(stderr, "Sync failed at bit %d\n", bit_index);
        return -1;
    }
    
    /* Measure baseline RTT */
    uint64_t baseline = measure_rtt(ctx->ctrl_sock);
    
    /* Both parties: if bit=1, flood; if bit=0, quiet */
    pthread_t flood_thread;
    int flooding = 0;
    
    if (our_bit) {
        flooding = 1;
        /* Start flood in background */
        pthread_create(&flood_thread, NULL, 
            (void *(*)(void *))send_flood, 
            (void *)ctx->peer_ip);
    }
    
    /* Wait for round duration and measure latency */
    usleep(ROUND_DURATION_MS * 1000 / 2);
    uint64_t latency = measure_rtt(ctx->ctrl_sock);
    usleep(ROUND_DURATION_MS * 1000 / 2);
    
    if (flooding) {
        pthread_cancel(flood_thread);
        pthread_join(flood_thread, NULL);
    }
    
    /* Exchange bits (in real protocol, use commitment scheme) */
    int peer_bit;
    exchange_commitment(ctx, our_bit, &peer_bit);
    
    /* 
     * Key derivation:
     * - XOR of both bits gives unpredictable result
     * - Latency measurement provides additional entropy
     * - Combined: key_bit = (our_bit XOR peer_bit) XOR (latency > threshold)
     */
    int latency_bit = (latency > baseline + LATENCY_THRESHOLD_US) ? 1 : 0;
    int key_bit = our_bit ^ peer_bit ^ latency_bit;
    
    /* Store key bit */
    int byte_idx = bit_index / 8;
    int bit_pos = bit_index % 8;
    if (key_bit) {
        ctx->key[byte_idx] |= (1 << bit_pos);
    }
    
    if (verbose) {
        printf("[Round %3d] our=%d peer=%d latency=%lu (base=%lu) -> key_bit=%d\n",
               bit_index, our_bit, peer_bit, latency, baseline, key_bit);
    }
    
    return 0;
}

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

/* Connect to peer (initiator mode) */
int connect_to_peer(const char *host, int port)
{
    struct addrinfo hints = {
        .ai_family = AF_UNSPEC,
        .ai_socktype = SOCK_STREAM,
    };
    struct addrinfo *result;
    char port_str[16];
    snprintf(port_str, sizeof(port_str), "%d", port);
    
    if (getaddrinfo(host, port_str, &hints, &result) != 0) {
        perror("getaddrinfo");
        return -1;
    }
    
    int sock = socket(result->ai_family, result->ai_socktype, result->ai_protocol);
    if (sock < 0) {
        perror("socket");
        freeaddrinfo(result);
        return -1;
    }
    
    /* Set TCP_NODELAY for accurate timing */
    int flag = 1;
    setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, &flag, sizeof(flag));
    
    if (connect(sock, result->ai_addr, result->ai_addrlen) < 0) {
        perror("connect");
        close(sock);
        freeaddrinfo(result);
        return -1;
    }
    
    freeaddrinfo(result);
    return sock;
}

/* Listen for peer (responder mode) */
int accept_peer(int port)
{
    int listen_sock = socket(AF_INET6, SOCK_STREAM, 0);
    if (listen_sock < 0) {
        listen_sock = socket(AF_INET, SOCK_STREAM, 0);
        if (listen_sock < 0) {
            perror("socket");
            return -1;
        }
        
        struct sockaddr_in addr = {
            .sin_family = AF_INET,
            .sin_port = htons(port),
            .sin_addr.s_addr = INADDR_ANY,
        };
        
        int opt = 1;
        setsockopt(listen_sock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
        
        if (bind(listen_sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
            perror("bind");
            close(listen_sock);
            return -1;
        }
    } else {
        struct sockaddr_in6 addr6 = {
            .sin6_family = AF_INET6,
            .sin6_port = htons(port),
            .sin6_addr = in6addr_any,
        };
        
        int opt = 1;
        setsockopt(listen_sock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
        setsockopt(listen_sock, IPPROTO_IPV6, IPV6_V6ONLY, &opt, sizeof(opt));
        
        if (bind(listen_sock, (struct sockaddr *)&addr6, sizeof(addr6)) < 0) {
            perror("bind");
            close(listen_sock);
            return -1;
        }
    }
    
    listen(listen_sock, 1);
    printf("[Server] Listening on port %d...\n", port);
    
    struct sockaddr_storage peer_addr;
    socklen_t peer_len = sizeof(peer_addr);
    int sock = accept(listen_sock, (struct sockaddr *)&peer_addr, &peer_len);
    
    close(listen_sock);
    
    if (sock < 0) {
        perror("accept");
        return -1;
    }
    
    /* Set TCP_NODELAY */
    int flag = 1;
    setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, &flag, sizeof(flag));
    
    printf("[Server] Peer connected\n");
    return sock;
}

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

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

void print_usage(const char *prog)
{
    printf("Usage: %s [options]\n", prog);
    printf("\nRemote Key Agreement via Network Timing Side-Channel\n");
    printf("\nOptions:\n");
    printf("  -c HOST     Connect to HOST (initiator mode)\n");
    printf("  -l          Listen for connection (responder mode)\n");
    printf("  -p PORT     Port number (default: %d)\n", SYNC_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: %s -l -v\n", prog);
    printf("  Host B: %s -c <host_a_ip> -v\n", prog);
}

int main(int argc, char *argv[])
{
    int mode = 0;  /* 0=help, 1=initiator, 2=responder */
    char *peer_host = NULL;
    int port = SYNC_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':
            mode = 1;
            peer_host = optarg;
            break;
        case 'l':
            mode = 2;
            break;
        case 'p':
            port = atoi(optarg);
            break;
        case 'b':
            num_bits = atoi(optarg);
            break;
        case 'v':
            verbose = 1;
            break;
        case 'h':
        default:
            print_usage(argv[0]);
            return 0;
        }
    }
    
    printf("╔════════════════════════════════════════════════════════════════╗\n");
    printf("║  Remote Key Agreement via Network Timing                       ║\n");
    printf("║  Uses CVE-2023-1206 timing side-channel                        ║\n");
    printf("╚════════════════════════════════════════════════════════════════╝\n\n");
    
    if (mode == 0) {
        print_usage(argv[0]);
        return 0;
    }
    
    signal(SIGINT, signal_handler);
    srand(time(NULL) ^ getpid());
    
    key_ctx_t ctx;
    memset(&ctx, 0, sizeof(ctx));
    
    /* Establish connection */
    if (mode == 1) {
        printf("[Init] Connecting to %s:%d...\n", peer_host, port);
        ctx.ctrl_sock = connect_to_peer(peer_host, port);
        ctx.is_initiator = 1;
        strncpy(ctx.peer_ip, peer_host, sizeof(ctx.peer_ip) - 1);
    } else {
        ctx.ctrl_sock = accept_peer(port);
        ctx.is_initiator = 0;
        /* Get peer IP from socket */
        struct sockaddr_storage peer;
        socklen_t len = sizeof(peer);
        getpeername(ctx.ctrl_sock, (struct sockaddr *)&peer, &len);
        if (peer.ss_family == AF_INET) {
            inet_ntop(AF_INET, &((struct sockaddr_in *)&peer)->sin_addr,
                      ctx.peer_ip, sizeof(ctx.peer_ip));
        } else {
            inet_ntop(AF_INET6, &((struct sockaddr_in6 *)&peer)->sin6_addr,
                      ctx.peer_ip, sizeof(ctx.peer_ip));
        }
    }
    
    if (ctx.ctrl_sock < 0) {
        fprintf(stderr, "Failed to establish connection\n");
        return 1;
    }
    
    printf("[Setup] Connected! Starting key agreement (%d bits)...\n\n", num_bits);
    
    /* Run key agreement rounds */
    for (int bit = 0; bit < num_bits && running; bit++) {
        if (key_round(&ctx, bit, verbose) < 0) {
            break;
        }
        
        if (!verbose && bit % 32 == 31) {
            printf("[Progress] %d/%d bits derived\n", bit + 1, num_bits);
        }
    }
    
    printf("\n[Complete] Key agreement finished!\n");
    print_key(ctx.key, num_bits / 8);
    
    close(ctx.ctrl_sock);
    return 0;
}