//! `Noise_IKpsk2` handshake initiator packets

use rand::{Rng, thread_rng};
use x25519_dalek::PublicKey;

use crate::noise::error::NoiseError;
use crate::noise::handshake::{HANDSHAKE_INITIATOR_CHAIN_KEY, HANDSHAKE_INITIATOR_CHAIN_KEY_HASH, HandshakeState, needs_cookie};
use crate::qcrypto::{LABEL_MAC1, timestamp};
use crate::qcrypto::aead::{qcrypto_aead, qcrypto_aead_decrypt};
use crate::qcrypto::hashes::{qcrypto_hash_twice, qcrypto_mac};
use crate::qcrypto::hkdf::qcrypto_hkdf;
use crate::qcrypto::pki::{qcrypto_dh_generate_longterm, qcrypto_dh_longterm};

struct HandshakeInitiatorRaw {
    sender: [u8; 4],
    ephemeral: [u8; 32],
    static_pub: [u8; 32 + 16],
    timestamp: [u8; 12 + 16],
    mac1: [u8; 16],
    mac2: [u8; 16]
}
impl HandshakeInitiatorRaw {
    fn to_bytes(&self, session: &HandshakeState) -> [u8; 148] {
        let mut output = [0u8; 148];

        output[0] = 1u8;
        output[4..8].copy_from_slice(&self.sender);
        output[8..40].copy_from_slice(&self.ephemeral);
        output[40..88].copy_from_slice(&self.static_pub);
        output[88..116].copy_from_slice(&self.timestamp);

        let mac1: [u8; 16] = qcrypto_mac(&qcrypto_hash_twice(LABEL_MAC1.as_bytes(), session.s_pub_i.as_bytes()), &output[..116]);

        output[116..132].copy_from_slice(&mac1);

        let mac2 = if needs_cookie(session) { qcrypto_mac(&session.cookies[session.cookies.len() - 1].cookie, &output[..132]) } else { [0u8; 16] };

        output[132..148].copy_from_slice(&mac2);

        output
    }

    #[allow(clippy::unwrap_used)] // Only used for type conversions in known safe ways
    fn from_bytes(bytes: [u8; 148]) -> Self {
        Self {
            sender: bytes[4..8].try_into().unwrap(),
            ephemeral: bytes[8..40].try_into().unwrap(),
            static_pub: bytes[40..88].try_into().unwrap(),
            timestamp: bytes[88..116].try_into().unwrap(),
            mac1: bytes[116..132].try_into().unwrap(),
            mac2: bytes[132..148].try_into().unwrap()
        }
    }
}

/// Generate a handshake initiator packet and encrypt it using the given session state, starting a new handshake state
/// # Errors
/// This function will error if encryption was unsuccessful
/// # Panics
/// While containing unwraps, this function will never panic.
#[allow(clippy::module_name_repetitions)]
#[allow(clippy::unwrap_used)] // Safe because it is only used for type conversions in known safe ways
pub fn generate_handshake_init(session: &mut HandshakeState) -> Result<[u8; 148], NoiseError> {
    session.we_are_initiator = true;
    session.s_pub_i = PublicKey::from(session.s_priv_me);
    session.s_pub_r = session.s_pub_them;

    session.i_i = thread_rng().gen();

    let mut msg = HandshakeInitiatorRaw {
        sender: session.i_i.to_le_bytes(),
        ephemeral: [0u8; 32],
        static_pub: [0u8; 32 + 16],
        timestamp: [0u8; 12 + 16],
        mac1: [0u8; 16],
        mac2: [0u8; 16]
    };

    session.ck = HANDSHAKE_INITIATOR_CHAIN_KEY;
    session.h = HANDSHAKE_INITIATOR_CHAIN_KEY_HASH;
    session.h = qcrypto_hash_twice(&session.h, session.s_pub_r.as_bytes());

    let eph_keypair = qcrypto_dh_generate_longterm();

    session.ck = qcrypto_hkdf::<1>(&session.ck, eph_keypair.1.as_bytes())[0];
    session.e_pub_i = eph_keypair.1;
    session.e_priv_me = eph_keypair.0;

    msg.ephemeral = eph_keypair.1.to_bytes();

    session.h = qcrypto_hash_twice(&session.h, &msg.ephemeral);

    let ci_k_pair = qcrypto_hkdf::<2>(&session.ck, qcrypto_dh_longterm(&session.e_priv_me, &session.s_pub_r).as_bytes());
    session.ck = ci_k_pair[0];
    let k = ci_k_pair[1];

    // This unwrap is safe because the output length is a known constant with these inputs
    msg.static_pub = match qcrypto_aead(&k, 0, session.s_pub_i.as_bytes(), &session.h) {
        Ok(s) => s,
        Err(e) => return Err(NoiseError::ChaCha20Error(e))
    }.try_into().unwrap();

    session.h = qcrypto_hash_twice(&session.h, &msg.static_pub);

    let ci_k_pair = qcrypto_hkdf::<2>(&session.ck, qcrypto_dh_longterm(session.s_priv_me, &session.s_pub_r).as_bytes());
    session.ck = ci_k_pair[0];
    let k = ci_k_pair[1];

    // This unwrap is safe because the output length is a known constant with these inputs
    msg.timestamp = match qcrypto_aead(&k, 0, &timestamp().to_bytes(), &session.h) {
        Ok(s) => s,
        Err(e) => return Err(NoiseError::ChaCha20Error(e))
    }.try_into().unwrap();

    session.h = qcrypto_hash_twice(&session.h, &msg.timestamp);

    Ok(msg.to_bytes(session))
}

/// Parse a handshake initiator packet and decrypt it using the given session state, updating the session state with decrypted and authenticated values
/// # Errors
/// This function will error if decryption was unsuccessful
/// # Panics
/// While containing unwraps, this function will never panic.
#[allow(clippy::module_name_repetitions)]
#[allow(clippy::unwrap_used)] // Only used for type conversions in known safe ways
pub fn parse_handshake_init(session: &mut HandshakeState, packet: [u8; 148]) -> Result<(), NoiseError> {
    let s_pub_r = PublicKey::from(session.s_priv_me);

    let msg = HandshakeInitiatorRaw::from_bytes(packet);

    let i_i = u32::from_le_bytes(msg.sender);

    let ck = HANDSHAKE_INITIATOR_CHAIN_KEY;
    let h = HANDSHAKE_INITIATOR_CHAIN_KEY_HASH;
    let h = qcrypto_hash_twice(&h, s_pub_r.as_bytes());

    let ephemeral_public = msg.ephemeral;
    let eph_pub = PublicKey::from(ephemeral_public);
    let e_pub_i = eph_pub;

    let ck = qcrypto_hkdf::<1>(&ck, eph_pub.as_bytes())[0];

    let h = qcrypto_hash_twice(&h, &msg.ephemeral);

    let ci_k_pair = qcrypto_hkdf::<2>(&ck, qcrypto_dh_longterm(session.s_priv_me, &eph_pub).as_bytes());
    let ck = ci_k_pair[0];
    let k = ci_k_pair[1];

    // This unwrap is safe because the output length is a known constant with these inputs


    let s_pub_i = PublicKey::from(<Vec<u8> as TryInto<[u8; 32]>>::try_into(match qcrypto_aead_decrypt(&k, 0, &msg.static_pub, &h) {
        Ok(s) => s,
        Err(e) => return Err(NoiseError::ChaCha20Error(e))
    }).unwrap());

    let h = qcrypto_hash_twice(&h, &msg.static_pub);

    let ci_k_pair = qcrypto_hkdf::<2>(&ck, qcrypto_dh_longterm(session.s_priv_me, &s_pub_i).as_bytes());
    let ck = ci_k_pair[0];
    let k = ci_k_pair[1];

    // This unwrap is safe because the output length is a known constant with these inputs
    let _sent_timestamp: [u8; 12] = match qcrypto_aead_decrypt(&k, 0, &msg.timestamp, &h) {
        Ok(s) => s,
        Err(e) => return Err(NoiseError::ChaCha20Error(e))
    }.try_into().unwrap();

    let h = qcrypto_hash_twice(&h, &msg.timestamp);

    // we need to check mac1 and mac2

    let mac1: [u8; 16] = qcrypto_mac(&qcrypto_hash_twice(LABEL_MAC1.as_bytes(), s_pub_i.as_bytes()), &packet[..116]);
    let mac2 = if needs_cookie(session) { qcrypto_mac(&session.cookies[session.cookies.len() - 1].cookie, &packet[..132]) } else { [0u8; 16] };

    if mac1 != msg.mac1 || mac2 != msg.mac2 {
        return Err(NoiseError::PacketUnauthenticated)
    }

    session.h = h;
    session.ck = ck;
    session.e_pub_i = e_pub_i;
    session.s_pub_i = s_pub_i;
    session.s_pub_r = s_pub_r;
    session.i_i = i_i;
    session.we_are_initiator = false;

    Ok(())
}