anubis-wormhole 1.0.0

use crate::aead_policy::{AAD, Flags, RecordType, Role, sc, capabilities_hash16};
use crate::frame::Frame;
use crate::subchannel::SubchSend;
use crate::providers::aes_gcm_siv::Aes256GcmSivProvider;
use crate::handshake::{derive_keys, AnubisDerivedKeys};
use zeroize::Zeroize;
use crate::transcript::Transcript;
#[cfg(feature = "policy-l5")]
use crate::policy;
#[cfg(feature = "providers-pqclean")]
use crate::providers::mlkem::MlKem1024;
#[cfg(feature = "providers-pqclean")]
use crate::traits::Kem;
// SAS utilities used by CLI; not needed here
use sha2::{Digest, Sha384};

#[derive(Clone, Debug)]
pub struct Suite {
    pub kdf: &'static str,
    pub kem: &'static str,
    pub aead: &'static str,
}

pub const SUITE_ANUBIS: Suite = Suite {
    kdf: "HKDF-SHA512",
    kem: "ML-KEM-1024",
    aead: "AES-256-GCM-SIV",
};

#[derive(Clone, Debug)]
pub struct Hello {
    pub versions: serde_json::Value,
    pub capabilities: serde_json::Value,
    pub suite: Suite,
}

#[derive(Clone, Debug)]
pub struct HelloAck { pub suite: Suite }

#[derive(Clone, Debug)]
pub enum ControlMsg {
    Hello(Hello),
    HelloAck(HelloAck),
    SpakeA(Vec<u8>),
    SpakeB(Vec<u8>),
    KemPk(Vec<u8>),
    KemCt(Vec<u8>),
    Confirm(Vec<u8>),
    RekeyStart { epoch: u32, kem_pk: Vec<u8> },
    RekeyCt { epoch: u32, kem_ct: Vec<u8> },
    RekeyDone { epoch: u32 },
}

use async_trait::async_trait;

#[async_trait]
pub trait PhaseIO {
    async fn send(&mut self, bytes: &[u8]);
    async fn recv(&mut self) -> Option<Vec<u8>>;
}

pub struct ControlDriver<'a, IO: PhaseIO> {
    pub io: &'a mut IO,
    pub role: Role,
    pub session_id: [u8; 16],
    pub epoch: u32,
    pub suite: Suite,
}

impl<'a, IO: PhaseIO> ControlDriver<'a, IO> {
    pub fn aead_seal(&self, key: &[u8;32], seq: u64, caps_hash: Option<[u8;16]>, rtype: RecordType, flags: Flags, plaintext: &[u8]) -> Vec<u8> {
        let aead = Aes256GcmSivProvider;
        let aad = AAD {
            version: 1,
            role: self.role,
            record_type: rtype,
            subchannel: sc::CONTROL,
            seq,
            flags,
            session_id: self.session_id,
            capabilities_hash16: caps_hash,
        };
        Frame::seal(&aead, key, &aad, plaintext).expect("seal")
    }
    pub fn aead_open_sender(&self, sender_role: Role, key: &[u8;32], seq: u64, caps_hash: Option<[u8;16]>, rtype: RecordType, flags: Flags, ciphertext: &[u8]) -> Vec<u8> {
        let aead = Aes256GcmSivProvider;
        let aad = AAD {
            version: 1,
            role: sender_role,
            record_type: rtype,
            subchannel: sc::CONTROL,
            seq,
            flags,
            session_id: self.session_id,
            capabilities_hash16: caps_hash,
        };
        Frame::open(&aead, key, &aad, ciphertext).expect("open")
    }
}

pub struct InitiatorKeys { pub keys: AnubisDerivedKeys, pub cap_hash16: [u8;16] }

pub async fn initiator_handshake_over<IO: PhaseIO + Send + Sync>(io: &mut IO, code: &str, versions: serde_json::Value, capabilities: serde_json::Value) -> (InitiatorKeys, [u8;16]) {
    // 1) Hello + Ack (plaintext under mailbox; not confidential but minimal)
    let hello = Hello { versions: versions.clone(), capabilities: capabilities.clone(), suite: SUITE_ANUBIS };
    let hello_bytes = serde_json::to_vec(&serde_json::json!({"hello":{"suite":hello.suite.aead,"versions":hello.versions,"capabilities":hello.capabilities}})).unwrap();
    io.send(&hello_bytes).await;
    // For loopback IO, consume the ack immediately
    let _ack = HelloAck { suite: SUITE_ANUBIS };

    // 2) SPAKE2 exchange (still unencrypted or mailbox-level until AEAD is ready)
    use spake2::{Ed25519Group, Identity, Password, Spake2};
    let pw = Password::new(code.as_bytes());
    let (state_a, msg_a) = Spake2::<Ed25519Group>::start_a(&pw, &Identity::new(&[]), &Identity::new(&[]));
    io.send(&msg_a).await;
    let msg_b = io.recv().await.expect("spake msg_b");
    let pake_key = state_a.finish(&msg_b).expect("spake finish");

    // 3) KEM inside PAKE transcript
    let mut tr = Transcript::new();
    tr.absorb("spake.msg_a", &msg_a);
    tr.absorb("spake.msg_b", &msg_b);
    let th = tr.finish();
    // 3b) KEM and derive keys (if available)
    #[cfg(feature = "providers-pqclean")]
    let keys = {
        let kem = MlKem1024;
        let (pk_a, sk_a) = kem.keypair();
        io.send(&pk_a).await;
        let ct_b = io.recv().await.expect("kem ct");
        let ss_a = kem.decapsulate(&sk_a, &ct_b);
        derive_keys(&pake_key, ss_a.as_ref(), &th)
    };
    #[cfg(not(feature = "providers-pqclean"))]
    let keys = {
        // Derive from PAKE-only if KEM is not compiled in
        derive_keys(&pake_key, &[], &th)
    };
    let mut key_arr = [0u8;32]; key_arr.copy_from_slice(&keys.k_ctrl[..32]);
    // 5) Confirm + SAS with subchannel sequencing
    let cap_obj = serde_json::json!({
        "suite": SUITE_ANUBIS.aead,
        "versions": versions,
        "capabilities": capabilities,
        "sec_cat": 5,
        "kem": SUITE_ANUBIS.kem,
        "aead": SUITE_ANUBIS.aead,
        "kdf": SUITE_ANUBIS.kdf
    });
    #[cfg(feature = "policy-l5")]
    { if policy::l5_enforcement_enabled() { let _ = policy::assert_level5_caps(&cap_obj); } }
    let cap = serde_json::to_vec(&cap_obj).unwrap();
    let cap_hash16 = capabilities_hash16(&cap);
    let mut h = Sha384::new(); h.update(&keys.k_verify); let sas_raw = h.finalize(); let sas = &sas_raw[..8];
    let mut sc = SubchSend::new(sc::CONTROL).with_capacity(16);
    sc.enqueue(&key_arr, session_id_from_th(&th), Role::Initiator, RecordType::Control, sas, false).expect("enqueue confirm");
    if let Some(ct) = sc.dequeue() { io.send(&ct).await; }

    #[cfg(feature = "policy-l5")]
    { if policy::l5_enforcement_enabled() { let _ = policy::assert_level5(SUITE_ANUBIS.kem, None); } }
    let out = (InitiatorKeys { keys, cap_hash16 }, session_id_from_th(&th));
    // Zeroize sensitive intermediates
    {
        let mut mb = msg_b.clone(); mb.zeroize();
    }
    // pake_key and th live within keys/session derivation; no separate copy retained here
    out
}

pub struct ResponderKeys { pub keys: AnubisDerivedKeys, pub cap_hash16: [u8;16] }

pub async fn responder_handshake_over<IO: PhaseIO + Send + Sync>(io: &mut IO, code: &str, versions: serde_json::Value, capabilities: serde_json::Value) -> (ResponderKeys, [u8;16]) {
    use spake2::{Ed25519Group, Identity, Password, Spake2};
    let pw = Password::new(code.as_bytes());
    let (state_b, msg_b) = Spake2::<Ed25519Group>::start_b(&pw, &Identity::new(&[]), &Identity::new(&[]));
    // receive msg_a, send msg_b
    let msg_a = io.recv().await.expect("spake msg_a");
    io.send(&msg_b).await;
    let pake_key = state_b.finish(&msg_a).expect("spake finish");

    let mut tr = Transcript::new();
    tr.absorb("spake.msg_a", &msg_a);
    tr.absorb("spake.msg_b", &msg_b);
    let th = tr.finish();

    // KEM if available, otherwise PAKE-only derivation
    #[cfg(feature = "providers-pqclean")]
    let keys = {
        let kem = MlKem1024;
        // We are responder: receive pk_a from initiator then send ct_b
        let pk_from_initiator = io.recv().await.expect("kem pk");
        let (ct_b, ss_b) = kem.encapsulate(&pk_from_initiator);
        io.send(&ct_b).await;
        derive_keys(&pake_key, ss_b.as_ref(), &th)
    };
    #[cfg(not(feature = "providers-pqclean"))]
    let keys = {
        derive_keys(&pake_key, &[], &th)
    };
    let cap_obj = serde_json::json!({
        "suite": SUITE_ANUBIS.aead,
        "versions": versions,
        "capabilities": capabilities,
        "sec_cat": 5,
        "kem": SUITE_ANUBIS.kem,
        "aead": SUITE_ANUBIS.aead,
        "kdf": SUITE_ANUBIS.kdf
    });
    #[cfg(feature = "policy-l5")]
    { if policy::l5_enforcement_enabled() { let _ = policy::assert_level5_caps(&cap_obj); } }
    let cap = serde_json::to_vec(&cap_obj).unwrap();
    let cap_hash16 = capabilities_hash16(&cap);

    // Receive confirm ct and verify SAS (enforced by caller)
    let sess = session_id_from_th(&th);
    let mut key_arr = [0u8;32]; key_arr.copy_from_slice(&keys.k_ctrl[..32]);
    let ct = io.recv().await.expect("confirm ct");
    let driver = ControlDriver { io, role: Role::Responder, session_id: sess, epoch: 0, suite: SUITE_ANUBIS };
    // First CONTROL record from initiator must be seq==0.
    // Under L5 enforcement, the sender set L5_POLICY on seq==0; reflect it here.
    #[allow(unused_mut)]
    let mut flags = Flags::empty();
    #[cfg(feature = "policy-l5")]
    {
        if policy::l5_enforcement_enabled() {
            flags |= Flags::L5_POLICY;
        }
    }
    let _sas_plain = driver.aead_open_sender(Role::Initiator, &key_arr, 0, Some(cap_hash16), RecordType::Control, flags, &ct);
    // Zeroize sensitive intermediates
    {
        let mut ma = msg_a.clone(); ma.zeroize();
        let mut mb = msg_b.clone(); mb.zeroize();
    }
    // SAS enforcement should happen outside (CLI/UI)
    #[cfg(feature = "policy-l5")]
    { if policy::l5_enforcement_enabled() { let _ = policy::assert_level5(SUITE_ANUBIS.kem, None); } }
    (ResponderKeys { keys, cap_hash16 }, sess)
}

fn session_id_from_th(th: &[u8]) -> [u8;16] {
    let sid_raw = sha2::Sha512::digest([b"anubis/session_id".as_ref(), th].concat());
    let mut session_id = [0u8; 16];
    session_id.copy_from_slice(&sid_raw[..16]);
    session_id
}

/// Exchange transit abilities/hints JSON over CONTROL after handshake.
/// Initiator sends first (seq=2), then receives responder info (seq=1 from responder side).
/// Responder receives first (seq=2 from initiator), then sends (seq=1 from responder).
pub async fn exchange_transit_info_over<IO: PhaseIO + Send + Sync>(io: &mut IO, keys: &AnubisDerivedKeys, role: Role, session_id: [u8;16], my_info_json: serde_json::Value) -> serde_json::Value {
    let mut key = [0u8;32]; key.copy_from_slice(&keys.k_ctrl[..32]);
    let driver = ControlDriver { io, role, session_id, epoch: 0, suite: SUITE_ANUBIS };
    match role {
        Role::Initiator => {
            // send
            let pt = serde_json::to_vec(&my_info_json).expect("json");
            let ct = driver.aead_seal(&key, 2, None, RecordType::Control, Flags::empty(), &pt);
            driver.io.send(&ct).await;
            // recv
            let peer_ct = driver.io.recv().await.expect("peer transit info");
            let peer_pt = driver.aead_open_sender(Role::Responder, &key, 1, None, RecordType::Control, Flags::empty(), &peer_ct);
            serde_json::from_slice(&peer_pt).expect("peer json")
        }
        Role::Responder => {
            // recv
            let peer_ct = driver.io.recv().await.expect("peer transit info");
            let peer_pt = driver.aead_open_sender(Role::Initiator, &key, 2, None, RecordType::Control, Flags::empty(), &peer_ct);
            // send
            let pt = serde_json::to_vec(&my_info_json).expect("json");
            let ct = driver.aead_seal(&key, 1, None, RecordType::Control, Flags::empty(), &pt);
            driver.io.send(&ct).await;
            serde_json::from_slice(&peer_pt).expect("peer json")
        }
    }
}

/// Send an offer (file/dir metadata) from initiator to responder at CONTROL seq=3.
pub async fn send_offer_over_control<IO: PhaseIO + Send + Sync>(io: &mut IO, keys: &AnubisDerivedKeys, role: Role, session_id: [u8;16], offer_json: serde_json::Value) -> Option<serde_json::Value> {
    let mut key = [0u8;32]; key.copy_from_slice(&keys.k_ctrl[..32]);
    let driver = ControlDriver { io, role, session_id, epoch: 0, suite: SUITE_ANUBIS };
    match role {
        Role::Initiator => {
            let pt = serde_json::to_vec(&offer_json).ok()?;
            let ct = driver.aead_seal(&key, 3, None, RecordType::Control, Flags::empty(), &pt);
            driver.io.send(&ct).await;
            None
        }
        Role::Responder => {
            let peer_ct = driver.io.recv().await?;
            let pt = driver.aead_open_sender(Role::Initiator, &key, 3, None, RecordType::Control, Flags::empty(), &peer_ct);
            serde_json::from_slice(&pt).ok()
        }
    }
}

pub async fn initiator_rekey<IO: PhaseIO + Send + Sync>(io: &mut IO, old_keys: &AnubisDerivedKeys, epoch: u32) -> AnubisDerivedKeys {
    #[cfg(feature = "providers-pqclean")]
    {
        let kem = MlKem1024;
        let (pk, sk) = kem.keypair();
        // Send RekeyStart{epoch+1, kem_pk}
        let msg = serde_json::to_vec(&serde_json::json!({"rk_start":{"epoch": epoch+1, "pk": hex::encode(&pk)}})).unwrap();
        io.send(&msg).await;
        // Receive RekeyCt{epoch+1, kem_ct}
        let ct_msg = io.recv().await.expect("rk ct");
        let val: serde_json::Value = serde_json::from_slice(&ct_msg).unwrap();
        let ct_hex = val["rk_ct"]["ct"].as_str().unwrap();
        let ct = hex::decode(ct_hex).unwrap();
        let ss = kem.decapsulate(&sk, &ct);
        // derive new master keys from old K_verify and ss (bind epochs)
        let mut t = Transcript::new();
        t.absorb("epoch", &(epoch+1).to_be_bytes());
        t.absorb("kverify", &old_keys.k_verify);
        let th = t.finish();
        return derive_keys(&old_keys.k_verify, ss.as_ref(), &th);
    }
    #[cfg(not(feature = "providers-pqclean"))]
    {
        // Fallback: keep old keys when KEM unavailable
        return old_keys.clone();
    }
}

pub async fn responder_rekey<IO: PhaseIO + Send + Sync>(io: &mut IO, old_keys: &AnubisDerivedKeys, _epoch: u32) -> AnubisDerivedKeys {
    #[cfg(feature = "providers-pqclean")]
    {
        let kem = MlKem1024;
        // Receive RekeyStart with peer pk
        let msg = io.recv().await.expect("rk start");
        let v: serde_json::Value = serde_json::from_slice(&msg).unwrap();
        let nxt = v["rk_start"]["epoch"].as_u64().unwrap() as u32;
        let pk_hex = v["rk_start"]["pk"].as_str().unwrap();
        let pk = hex::decode(pk_hex).unwrap();
        let (ct, ss) = kem.encapsulate(&pk);
        let ct_msg = serde_json::to_vec(&serde_json::json!({"rk_ct":{"epoch": nxt, "ct": hex::encode(&ct)}})).unwrap();
        io.send(&ct_msg).await;
        // derive new keys
        let mut t = Transcript::new();
        t.absorb("epoch", &nxt.to_be_bytes());
        t.absorb("kverify", &old_keys.k_verify);
        let th = t.finish();
        return derive_keys(&old_keys.k_verify, ss.as_ref(), &th);
    }
    #[cfg(not(feature = "providers-pqclean"))]
    {
        // Fallback: keep old keys
        return old_keys.clone();
    }
}