//! `pq_wire_tap.rs` — наглядное доказательство (для отчёта по практике), что Aura собирает
//! постквантовый туннель и что трафик после хендшейка реально зашифрован.
//!
//! Тест прокачивает один полный клиент↔сервер обмен (handshake + одна Data-запись) поверх
//! in-memory duplex-пайпа, обёрнутого «отводом» ([`TeeWriter`]), который сохраняет каждый байт
//! на проводе. Затем по сохранённому потоку байтов проверяются четыре утверждения, каждое из
//! которых соответствует тест-кейсу в `TEST_CASES.md`:
//!
//! 1. **Туннель собран.** Хендшейк завершается успешно, каждая сторона распознала Common Name
//!    другой стороны по сертификату.
//! 2. **Размеры PQ-полей соответствуют FIPS 203.** В ClientHello payload ровно
//!    32 (X25519 pub) + 1184 (ML-KEM-768 encapsulation key) + 32 (nonce); в ServerHello payload
//!    ровно 32 (X25519 эфемеральный) + 1088 (ML-KEM-768 ciphertext) + 32 (nonce).
//! 3. **Открытого текста на проводе нет.** Уникальный 56-байтовый маркер, посланный в Data-кадре,
//!    не встречается ни в одном из двух направлений.
//! 4. **Шифротекст похож на случайный.** Тело ServerAuth (первый зашифрованный кадр после
//!    ServerHello) имеет Shannon-энтропию ≥ 7.0 бит на байт — характерная подпись AEAD-вывода.

mod common;

use std::io;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};

use aura_proto::{client_handshake, server_handshake, Frame, MsgType};
use bytes::Bytes;
use tokio::io::{split, AsyncWrite};

/// Уникальная 56-байтовая строка, которую мы шлём в Data-кадре. После шифрования ChaCha20-Poly1305
/// её не должно остаться на проводе ни в одном направлении — иначе крипты нет.
const PLAINTEXT_MARKER: &[u8] = b"AURA_PQ_PRACTICE_PROOF_MARKER_NEVER_APPEARS_ON_WIRE_2026";

/// Размер «дополнительного» payload'а, который мы шлём вместе с маркером, чтобы получить
/// достаточно большой блок AEAD-вывода для энтропийной проверки. ChaCha20-Poly1305 на этих
/// данных даст ≈ 1 КБ шифротекста; на такой выборке энтропия уверенно близка к 8 бит/байт.
const ENTROPY_PADDING_LEN: usize = 1024;

/// Длина 8-байтного `seq`-префикса перед AEAD-телом в Data-записи (см. `aura_proto::session`).
/// Это часть открытых метаданных, а не вывод шифра, поэтому из энтропийной оценки её исключаем.
const SEQ_LEN: usize = 8;

// ===== Помощник: writer, дублирующий каждый байт в общий буфер ===============================

/// Прозрачно проксирует все вызовы [`AsyncWrite`] на `inner`, попутно складывая успешно
/// записанные байты в общий `log`. Используется, чтобы перехватить полный набор байтов на
/// проводе без необходимости лезть в реальный сетевой стек.
struct TeeWriter<W> {
    inner: W,
    log: Arc<Mutex<Vec<u8>>>,
}

impl<W> TeeWriter<W> {
    fn new(inner: W, log: Arc<Mutex<Vec<u8>>>) -> Self {
        Self { inner, log }
    }
}

impl<W: AsyncWrite + Unpin> AsyncWrite for TeeWriter<W> {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        let res = Pin::new(&mut self.inner).poll_write(cx, buf);
        if let Poll::Ready(Ok(n)) = &res {
            self.log
                .lock()
                .expect("tap log mutex not poisoned")
                .extend_from_slice(&buf[..*n]);
        }
        res
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_shutdown(cx)
    }
}

// ===== Помощники для разбора заголовков и анализа байтов =====================================

/// Длина протокольного заголовка Aura (см. `aura_proto::frame`).
const HEADER_LEN: usize = 5;

/// Распаковать u24-be длину payload из 5-байтового заголовка по смещению `off`.
fn read_payload_len(buf: &[u8], off: usize) -> usize {
    ((buf[off + 1] as usize) << 16) | ((buf[off + 2] as usize) << 8) | (buf[off + 3] as usize)
}

/// Подсчитать классическую Shannon-энтропию (бит/байт) последовательности.
///
/// Для равномерно случайных байт энтропия стремится к 8.0; AEAD-вывод на практике даёт
/// > 7.0 даже на коротких блобах в сотни байт.
fn shannon_entropy(b: &[u8]) -> f64 {
    if b.is_empty() {
        return 0.0;
    }
    let mut counts = [0u64; 256];
    for &x in b {
        counts[x as usize] += 1;
    }
    let n = b.len() as f64;
    let mut h = 0.0;
    for &c in &counts {
        if c == 0 {
            continue;
        }
        let p = c as f64 / n;
        h -= p * p.log2();
    }
    h
}

/// Линейный поиск `needle` в `hay`.
fn contains_subslice(hay: &[u8], needle: &[u8]) -> bool {
    if needle.is_empty() {
        return true;
    }
    hay.windows(needle.len()).any(|w| w == needle)
}

// ===== Основной интеграционный тест ==========================================================

/// Полный цикл: handshake + одно зашифрованное сообщение + ответ. По завершении исследуем
/// сохранённые на «отводах» байты и убеждаемся, что:
///
/// * формат ClientHello / ServerHello точно соответствует профилю Aura (X25519 + ML-KEM-768);
/// * маркерная строка не утекла на провод ни в одну сторону;
/// * первый постхендшейковый кадр (ServerAuth) выглядит случайным.
///
/// Эти три проверки в сумме — наблюдаемое доказательство, что туннель действительно работает
/// поверх гибридного PQ-KEM и применяет AEAD-шифрование ко всему, что идёт после ServerHello.
#[tokio::test]
async fn pq_handshake_and_data_wire_capture() {
    let pki = common::mint_pki("vpn.aura.example", "client-pq-proof");
    let client_cfg = pki.client_config();
    let server_cfg = pki.server_config();

    // Связанный in-memory транспорт. Половинки для чтения отдаём «как есть»; половинки для
    // записи оборачиваем TeeWriter'ом, чтобы скопить полный поток байтов на проводе.
    let (client_end, server_end) = tokio::io::duplex(256 * 1024);
    let (c_read, c_write_raw) = split(client_end);
    let (s_read, s_write_raw) = split(server_end);

    let c_to_s_log: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new()));
    let s_to_c_log: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new()));

    let c_write = TeeWriter::new(c_write_raw, c_to_s_log.clone());
    let s_write = TeeWriter::new(s_write_raw, s_to_c_log.clone());

    // Собираем payload: маркер + 1 КБ нулей. Нули в plaintext'е после ChaCha20 превращаются в
    // чистый поток ключа — это даёт нам объективно большой и репрезентативный AEAD-блок для
    // энтропийной проверки, при этом сохраняя возможность искать маркер на проводе.
    let mut data_payload = Vec::with_capacity(PLAINTEXT_MARKER.len() + ENTROPY_PADDING_LEN);
    data_payload.extend_from_slice(PLAINTEXT_MARKER);
    data_payload.extend_from_slice(&vec![0u8; ENTROPY_PADDING_LEN]);
    let data_payload_for_client = data_payload.clone();

    // Запускаем обе стороны параллельно.
    let client = tokio::spawn(async move {
        let mut session = client_handshake(c_read, c_write, &client_cfg)
            .await
            .expect("client handshake");
        // Шлём один Data-кадр с заведомо узнаваемым маркером + большим зануленным хвостом.
        session
            .send_frame(Frame::Data {
                stream_id: 0xC0FF_EEBB,
                payload: Bytes::from(data_payload_for_client),
            })
            .await
            .expect("send_frame");
        // Дожидаемся встречного Pong, чтобы исключить TOCTOU между записью и проверкой логов.
        let reply = session.recv_frame().await.expect("recv reply");
        assert!(matches!(reply, Frame::Pong { seq: 42 }));
        session.peer_id().map(str::to_string)
    });
    let server = tokio::spawn(async move {
        let mut session = server_handshake(s_read, s_write, &server_cfg)
            .await
            .expect("server handshake");
        let incoming = session.recv_frame().await.expect("recv data");
        // Сервер видит plaintext в чистом виде (после AEAD-open), но на проводе его не было.
        if let Frame::Data { ref payload, .. } = incoming {
            assert_eq!(
                payload.as_ref(),
                data_payload.as_slice(),
                "plaintext after decryption matches what client sent"
            );
        } else {
            panic!("expected Data frame, got {incoming:?}");
        }
        session
            .send_frame(Frame::Pong { seq: 42 })
            .await
            .expect("send Pong");
        session.peer_id().map(str::to_string)
    });

    let client_peer = client.await.expect("client task").expect("client peer id");
    let server_peer = server.await.expect("server task").expect("server peer id");

    // === ТК-1: туннель собран — обе стороны взаимно аутентифицированы. =======================
    assert_eq!(
        server_peer, "client-pq-proof",
        "server learned client CN from verified leaf certificate"
    );
    assert_eq!(
        client_peer, "vpn.aura.example",
        "client recorded the server name it authenticated"
    );

    // Снимаем итоговые буферы байтов.
    let c_to_s = c_to_s_log.lock().expect("c->s log").clone();
    let s_to_c = s_to_c_log.lock().expect("s->c log").clone();

    // === ТК-2: ClientHello имеет точный PQ-гибридный layout. ================================
    // FIPS 203: ML-KEM-768 encapsulation key = 1184 bytes; X25519 public key = 32 bytes;
    // handshake nonce = 32 bytes. Версионный байт = 0x01 (project §6.1).
    const ML_KEM_EK_LEN: usize = 1184;
    const ML_KEM_CT_LEN: usize = 1088;
    const X25519_LEN: usize = 32;
    const NONCE_LEN: usize = 32;
    const CH_PAYLOAD_LEN: usize = X25519_LEN + ML_KEM_EK_LEN + NONCE_LEN;
    const SH_PAYLOAD_LEN: usize = X25519_LEN + ML_KEM_CT_LEN + NONCE_LEN;

    assert!(
        c_to_s.len() >= HEADER_LEN + CH_PAYLOAD_LEN,
        "client must have written at least the ClientHello frame ({} bytes), got {}",
        HEADER_LEN + CH_PAYLOAD_LEN,
        c_to_s.len(),
    );
    assert_eq!(
        c_to_s[0],
        MsgType::ClientHello as u8,
        "first byte on wire is the ClientHello msg-type tag (0x01)",
    );
    assert_eq!(
        read_payload_len(&c_to_s, 0),
        CH_PAYLOAD_LEN,
        "ClientHello payload = X25519(32) || ML-KEM-768 ek(1184) || nonce(32)",
    );
    assert_eq!(c_to_s[4], 0x01, "protocol version byte 0x01");

    // === ТК-2': ServerHello точно так же. ====================================================
    assert!(s_to_c.len() >= HEADER_LEN + SH_PAYLOAD_LEN);
    assert_eq!(
        s_to_c[0],
        MsgType::ServerHello as u8,
        "first byte from server is the ServerHello tag (0x02)",
    );
    assert_eq!(
        read_payload_len(&s_to_c, 0),
        SH_PAYLOAD_LEN,
        "ServerHello payload = X25519_eph(32) || ML-KEM-768 ct(1088) || nonce(32)",
    );

    // === ТК-3: открытого маркера нет ни в одну сторону. ======================================
    // Это самое прямое доказательство: если бы AEAD не работал, plaintext «PROOF_MARKER...»
    // лежал бы прямо в шифрованной части Data-кадра.
    assert!(
        !contains_subslice(&c_to_s, PLAINTEXT_MARKER),
        "plaintext marker LEAKED into c->s wire bytes ({} bytes captured)",
        c_to_s.len(),
    );
    assert!(
        !contains_subslice(&s_to_c, PLAINTEXT_MARKER),
        "plaintext marker LEAKED into s->c wire bytes",
    );

    // === ТК-4: тело ServerAuth выглядит случайным (Shannon entropy ≥ 7 бит/байт). ============
    // Сразу после ServerHello сервер шлёт ServerAuth (зашифрованный AEAD'ом сессионным ключом
    // s2c). Если бы шифра не было, мы бы увидели DER-сертификат и подпись — низкая энтропия.
    let server_auth_off = HEADER_LEN + SH_PAYLOAD_LEN;
    assert!(
        s_to_c.len() > server_auth_off + HEADER_LEN,
        "server must have sent ServerAuth right after ServerHello",
    );
    assert_eq!(
        s_to_c[server_auth_off],
        MsgType::ServerAuth as u8,
        "next frame after ServerHello is ServerAuth (0x04)",
    );
    let sa_payload_len = read_payload_len(&s_to_c, server_auth_off);
    let body_start = server_auth_off + HEADER_LEN;
    let body_end = body_start + sa_payload_len;
    assert!(s_to_c.len() >= body_end);
    let body = &s_to_c[body_start..body_end];
    let ent = shannon_entropy(body);
    assert!(
        ent >= 7.0,
        "ServerAuth body must look like AEAD ciphertext (entropy = {ent:.3} bits/byte over {} bytes; \
         clear DER would be < 5)",
        body.len(),
    );

    // === ТК-4': аналогично для Data-кадра, который шёл с клиента на сервер. ==================
    // Data всегда последний кадр в c_to_s после ClientHello, ClientAuth, Finished.
    // Найдём его, просканировав c_to_s по заголовкам.
    let mut off = 0;
    let mut last_data: Option<(usize, usize)> = None;
    while off + HEADER_LEN <= c_to_s.len() {
        let ty = c_to_s[off];
        let len = read_payload_len(&c_to_s, off);
        let end = off + HEADER_LEN + len;
        if end > c_to_s.len() {
            break;
        }
        if ty == MsgType::Data as u8 {
            last_data = Some((off + HEADER_LEN, end));
        }
        off = end;
    }
    let (ds, de) = last_data.expect("c->s must contain at least one Data record");
    // Тело Data-записи = seq(8) || AEAD(frame_bytes, ...). Первые 8 байт — открытый счётчик,
    // именно поэтому считать энтропию надо ТОЛЬКО по AEAD-части, иначе нули в seq её занижают.
    assert!(de - ds > SEQ_LEN, "Data body must include AEAD-ciphertext");
    let data_ciphertext = &c_to_s[ds + SEQ_LEN..de];
    let data_ent = shannon_entropy(data_ciphertext);
    assert!(
        data_ent >= 7.0,
        "Data-record AEAD body must look encrypted (entropy = {data_ent:.3} bits/byte over {} bytes; \
         clear text padded with zeros would be near 0)",
        data_ciphertext.len(),
    );

    // === Резюме (печатается на --nocapture, удобно для отчёта). ==============================
    eprintln!("=== Aura PQ wire-tap test summary ===");
    eprintln!(
        "client_peer = {client_peer:?}, server_peer = {server_peer:?}"
    );
    eprintln!(
        "captured c->s = {} bytes, s->c = {} bytes",
        c_to_s.len(),
        s_to_c.len()
    );
    eprintln!(
        "ClientHello payload = {} bytes (= 32 + 1184 + 32, X25519 + ML-KEM-768 ek + nonce)",
        read_payload_len(&c_to_s, 0)
    );
    eprintln!(
        "ServerHello payload = {} bytes (= 32 + 1088 + 32, X25519_eph + ML-KEM-768 ct + nonce)",
        read_payload_len(&s_to_c, 0)
    );
    eprintln!(
        "ServerAuth body Shannon entropy = {ent:.3} bits/byte over {} bytes",
        body.len()
    );
    eprintln!(
        "Data record AEAD body Shannon entropy = {data_ent:.3} bits/byte over {} bytes \
         (plaintext was marker + {} zero bytes; zeros become keystream after ChaCha20)",
        data_ciphertext.len(),
        ENTROPY_PADDING_LEN
    );
    eprintln!("Plaintext marker present on wire? c->s: NO, s->c: NO");
}

/// Микро-тест: вспомогательная функция `shannon_entropy` ведёт себя как ожидается.
/// Это не часть основного отчёта, но защищает от регрессий в самом проверочном коде.
#[test]
fn shannon_entropy_baseline() {
    // Полностью одинаковые байты → 0 бит.
    assert!((shannon_entropy(&[0xAAu8; 1024]) - 0.0).abs() < 1e-9);
    // 256 различных значений по одному разу → ровно 8 бит.
    let uniform: Vec<u8> = (0u32..256).map(|i| i as u8).collect();
    let h = shannon_entropy(&uniform);
    assert!((h - 8.0).abs() < 1e-9, "uniform entropy = {h}");
    // ASCII-текст обычно даёт < 5 бит.
    let text = b"This is some readable English text written here for the entropy baseline.";
    let ht = shannon_entropy(text);
    assert!(ht < 5.0, "ASCII entropy = {ht}");
}