//! Directory synchronized handshake
//! 1-RTT connection
//!
//! The simplest, fastest handshake supported by Fenrir
//! Downside: It does not offer protection from DDos,
//! no perfect forward secrecy
//!
//! To grant a form of perfect forward secrecy, the server should periodically
//! change the DNSSEC public/private keys

use super::{Error, HandshakeData};
use crate::{
    auth,
    connection::{ProtocolVersion, ID},
    enc::{
        asym::{ExchangePubKey, KeyExchangeKind, KeyID},
        hkdf::HkdfKind,
        sym::{CipherKind, HeadLen, TagLen},
        Random, Secret,
    },
};

// TODO: merge with crate::enc::sym::Nonce
/// random nonce
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Nonce(pub(crate) [u8; 16]);

impl Nonce {
    /// Create a new random Nonce
    pub fn new(rnd: &Random) -> Self {
        use ::core::mem::MaybeUninit;
        let mut out: MaybeUninit<[u8; 16]>;
        #[allow(unsafe_code)]
        unsafe {
            out = MaybeUninit::uninit();
            let _ = rnd.fill(out.assume_init_mut());
            Self(out.assume_init())
        }
    }
    /// Length of the serialized Nonce
    pub const fn len() -> usize {
        16
    }
}
impl From<&[u8; 16]> for Nonce {
    fn from(raw: &[u8; 16]) -> Self {
        Self(raw.clone())
    }
}

/// Parsed handshake
#[derive(Debug, Clone, PartialEq)]
pub enum DirSync {
    /// Directory synchronized handshake: client request
    Req(Req),
    /// Directory synchronized handshake: server response
    Resp(Resp),
}

impl DirSync {
    /// actual length of the dirsync handshake data
    pub fn len(&self, head_len: HeadLen, tag_len: TagLen) -> usize {
        match self {
            DirSync::Req(req) => req.len(),
            DirSync::Resp(resp) => resp.len(head_len, tag_len),
        }
    }
    /// Serialize into raw bytes
    /// NOTE: assumes that there is exactly asa much buffer as needed
    pub fn serialize(
        &self,
        head_len: HeadLen,
        tag_len: TagLen,
        out: &mut [u8],
    ) {
        match self {
            DirSync::Req(req) => req.serialize(head_len, tag_len, out),
            DirSync::Resp(resp) => resp.serialize(head_len, tag_len, out),
        }
    }
}

/// Client request of a directory synchronized handshake
#[derive(Debug, Clone, PartialEq)]
pub struct Req {
    /// Id of the server key used for the key exchange
    pub key_id: KeyID,
    /// Selected key exchange
    pub exchange: KeyExchangeKind,
    /// Selected hkdf
    pub hkdf: HkdfKind,
    /// Selected cipher
    pub cipher: CipherKind,
    /// Client ephemeral public key used for key exchanges
    pub exchange_key: ExchangePubKey,
    /// encrypted data
    pub data: ReqInner,
    // SECURITY: TODO: Add padding to min: 1200 bytes
    //   to avoid amplification attaks
    //   also: 1200 < 1280 to allow better vpn compatibility
}

impl Req {
    /// return the offset of the encrypted data
    /// NOTE: starts from the beginning of the fenrir packet
    pub fn encrypted_offset(&self) -> usize {
        ProtocolVersion::len()
            + crate::handshake::HandshakeID::len()
            + KeyID::len()
            + KeyExchangeKind::len()
            + HkdfKind::len()
            + CipherKind::len()
            + self.exchange_key.kind().pub_len()
    }
    /// return the total length of the cleartext data
    pub fn encrypted_length(&self) -> usize {
        match &self.data {
            ReqInner::ClearText(data) => data.len(),
            _ => 0,
        }
    }
    /// actual length of the directory synchronized request
    pub fn len(&self) -> usize {
        KeyID::len()
            + KeyExchangeKind::len()
            + HkdfKind::len()
            + CipherKind::len()
            + self.exchange_key.kind().pub_len()
            + self.cipher.nonce_len().0
            + self.data.len()
            + self.cipher.tag_len().0
    }
    /// Serialize into raw bytes
    /// NOTE: assumes that there is exactly as much buffer as needed
    pub fn serialize(
        &self,
        head_len: HeadLen,
        tag_len: TagLen,
        out: &mut [u8],
    ) {
        out[0..2].copy_from_slice(&self.key_id.0.to_le_bytes());
        out[2] = self.exchange as u8;
        out[3] = self.hkdf as u8;
        out[4] = self.cipher as u8;
        let key_len = self.exchange_key.len();
        let written_next = 5 + key_len;
        self.exchange_key.serialize_into(&mut out[5..written_next]);
        let written = written_next;
        if let ReqInner::ClearText(data) = &self.data {
            let from = written + head_len.0;
            let to = out.len() - tag_len.0;
            data.serialize(&mut out[from..to]);
        } else {
            unreachable!();
        }
    }
}

impl super::HandshakeParsing for Req {
    fn deserialize(raw: &[u8]) -> Result<HandshakeData, Error> {
        const MIN_PKT_LEN: usize = 10;
        if raw.len() < MIN_PKT_LEN {
            return Err(Error::NotEnoughData);
        }
        let key_id: KeyID =
            KeyID(u16::from_le_bytes(raw[0..2].try_into().unwrap()));
        use ::num_traits::FromPrimitive;
        let exchange: KeyExchangeKind = match KeyExchangeKind::from_u8(raw[2]) {
            Some(exchange) => exchange,
            None => return Err(Error::Parsing),
        };
        let hkdf: HkdfKind = match HkdfKind::from_u8(raw[3]) {
            Some(exchange) => exchange,
            None => return Err(Error::Parsing),
        };
        let cipher: CipherKind = match CipherKind::from_u8(raw[4]) {
            Some(cipher) => cipher,
            None => return Err(Error::Parsing),
        };
        let (exchange_key, len) = match ExchangePubKey::deserialize(&raw[5..]) {
            Ok(exchange_key) => exchange_key,
            Err(e) => return Err(e.into()),
        };
        let data = ReqInner::CipherText(raw.len() - (5 + len));
        Ok(HandshakeData::DirSync(DirSync::Req(Self {
            key_id,
            exchange,
            hkdf,
            cipher,
            exchange_key,
            data,
        })))
    }
}

/// Quick way to avoid mixing cipher and clear text
#[derive(Debug, Clone, PartialEq)]
pub enum ReqInner {
    /// Data is still encrytped, we only keep the length
    CipherText(usize),
    /// Client data, decrypted and parsed
    ClearText(ReqData),
}
impl ReqInner {
    /// The length of the data
    pub fn len(&self) -> usize {
        match self {
            ReqInner::CipherText(len) => *len,
            ReqInner::ClearText(data) => data.len(),
        }
    }
    /// parse the cleartext
    pub fn deserialize_as_cleartext(
        &mut self,
        raw: &[u8],
    ) -> Result<(), Error> {
        let clear = match self {
            ReqInner::CipherText(len) => {
                assert!(
                    *len > raw.len(),
                    "DirSync::ReqInner::CipherText length mismatch"
                );
                match ReqData::deserialize(raw) {
                    Ok(clear) => clear,
                    Err(e) => return Err(e),
                }
            }
            _ => return Err(Error::Parsing),
        };
        *self = ReqInner::ClearText(clear);
        Ok(())
    }
}

/// Informations needed for authentication
#[derive(Debug, Clone, PartialEq)]
pub struct AuthInfo {
    /// User of the domain
    pub user: auth::UserID,
    /// Authentication token
    pub token: auth::Token,
    /// service ID that we want to use on the domain
    pub service_id: auth::ServiceID,
    /// subdomain on which we authenticate
    // SECURITY: TODO: this should be padded to multiples of 32 bytes or so
    // reason: packet length can let you infer the authentication domain
    pub domain: auth::Domain,
}

impl AuthInfo {
    /// Minimum length of the authentication info
    pub const MIN_PKT_LEN: usize =
        auth::UserID::len() + auth::Token::len() + auth::ServiceID::len() + 1;
    /// Actual length of the authentication info
    pub fn len(&self) -> usize {
        Self::MIN_PKT_LEN + self.domain.len()
    }
    /// serialize into a buffer
    /// Note: assumes there is enough space
    pub fn serialize(&self, out: &mut [u8]) {
        out[..auth::UserID::len()].copy_from_slice(&self.user.0);
        const WRITTEN_TOKEN: usize = auth::UserID::len() + auth::Token::len();
        out[auth::UserID::len()..WRITTEN_TOKEN].copy_from_slice(&self.token.0);
        const WRITTEN_SERVICE_ID: usize =
            WRITTEN_TOKEN + auth::ServiceID::len();
        out[WRITTEN_TOKEN..WRITTEN_SERVICE_ID]
            .copy_from_slice(&self.service_id.0);
        let domain_len = self.domain.0.as_bytes().len() as u8;
        out[WRITTEN_SERVICE_ID] = domain_len;
        const WRITTEN_DOMAIN_LEN: usize = WRITTEN_SERVICE_ID + 1;
        out[WRITTEN_DOMAIN_LEN..].copy_from_slice(&self.domain.0.as_bytes());
    }
    /// deserialize from raw bytes
    pub fn deserialize(raw: &[u8]) -> Result<Self, Error> {
        if raw.len() < Self::MIN_PKT_LEN {
            return Err(Error::NotEnoughData);
        }
        let raw_user_id: [u8; auth::UserID::len()];
        let raw_token: [u8; auth::Token::len()];
        let raw_service_id: [u8; auth::ServiceID::len()];
        let mut start = 0;
        let mut end = auth::UserID::len();
        raw_user_id = raw[start..end].try_into().unwrap();
        let user: auth::UserID = raw_user_id.into();
        start = end;
        end = end + auth::Token::len();
        raw_token = raw[start..end].try_into().unwrap();
        let token: auth::Token = raw_token.into();
        start = end;
        end = end + auth::ServiceID::len();
        raw_service_id = raw[start..end].try_into().unwrap();
        let service_id: auth::ServiceID = raw_service_id.into();
        start = end;
        let domain_len = raw[start];
        start = start + 1;
        end = start + domain_len as usize;
        if raw.len() < end {
            return Err(Error::NotEnoughData);
        }
        let domain: auth::Domain = match raw[start..end].try_into() {
            Ok(domain) => domain,
            Err(_) => return Err(Error::Parsing),
        };
        Ok(Self {
            user,
            token,
            service_id,
            domain,
        })
    }
}

/// Decrypted request data
#[derive(Debug, Clone, PartialEq)]
pub struct ReqData {
    /// Random nonce, the client can use this to track multiple key exchanges
    pub nonce: Nonce,
    /// Client key id so the client can use and rotate keys
    pub client_key_id: KeyID,
    /// Receiving connection id for the client
    pub id: ID,
    /// Authentication data
    pub auth: AuthInfo,
}
impl ReqData {
    /// actual length of the request data
    pub fn len(&self) -> usize {
        Nonce::len() + KeyID::len() + ID::len() + self.auth.len()
    }
    /// Minimum byte length of the request data
    pub const MIN_PKT_LEN: usize =
        16 + KeyID::len() + ID::len() + AuthInfo::MIN_PKT_LEN;
    /// serialize into a buffer
    /// Note: assumes there is enough space
    pub fn serialize(&self, out: &mut [u8]) {
        out[..Nonce::len()].copy_from_slice(&self.nonce.0);
        const WRITTEN_KEY: usize = Nonce::len() + KeyID::len();
        out[Nonce::len()..WRITTEN_KEY]
            .copy_from_slice(&self.client_key_id.0.to_le_bytes());
        const WRITTEN: usize = WRITTEN_KEY;
        const WRITTEN_ID: usize = WRITTEN + 8;
        out[WRITTEN..WRITTEN_ID]
            .copy_from_slice(&self.id.as_u64().to_le_bytes());
        self.auth.serialize(&mut out[WRITTEN_ID..]);
    }
    /// Parse the cleartext raw data
    pub fn deserialize(raw: &[u8]) -> Result<Self, Error> {
        if raw.len() < Self::MIN_PKT_LEN {
            return Err(Error::NotEnoughData);
        }
        let mut start = 0;
        let mut end = 16;
        let raw_sized: &[u8; 16] = raw[start..end].try_into().unwrap();
        let nonce: Nonce = raw_sized.into();
        start = end;
        end = end + KeyID::len();
        let client_key_id =
            KeyID(u16::from_le_bytes(raw[start..end].try_into().unwrap()));
        start = end;
        end = end + ID::len();
        let id: ID =
            u64::from_le_bytes(raw[start..end].try_into().unwrap()).into();
        if id.is_handshake() {
            return Err(Error::Parsing);
        }
        start = end;
        //end = raw.len();
        let auth = AuthInfo::deserialize(&raw[start..])?;
        Ok(Self {
            nonce,
            client_key_id,
            id,
            auth,
        })
    }
}

/// Quick way to avoid mixing cipher and clear text
#[derive(Debug, Clone, PartialEq)]
pub enum RespInner {
    /// Server data, still in ciphertext
    CipherText(usize),
    /// Parsed, cleartext server data
    ClearText(RespData),
}
impl RespInner {
    /// The length of the data
    pub fn len(&self) -> usize {
        match self {
            RespInner::CipherText(len) => *len,
            RespInner::ClearText(_) => RespData::len(),
        }
    }
    /// parse the cleartext
    pub fn deserialize_as_cleartext(
        &mut self,
        raw: &[u8],
    ) -> Result<(), Error> {
        let clear = match self {
            RespInner::CipherText(len) => {
                assert!(
                    *len > raw.len(),
                    "DirSync::RespInner::CipherText length mismatch"
                );
                match RespData::deserialize(raw) {
                    Ok(clear) => clear,
                    Err(e) => return Err(e),
                }
            }
            _ => return Err(Error::Parsing),
        };
        *self = RespInner::ClearText(clear);
        Ok(())
    }
    /// Serialize the still cleartext data
    pub fn serialize(&self, out: &mut [u8]) {
        if let RespInner::ClearText(clear) = &self {
            clear.serialize(out);
        }
    }
}

/// Server response in a directory synchronized handshake
#[derive(Debug, Clone, PartialEq)]
pub struct Resp {
    /// Tells the client with which key the exchange was done
    pub client_key_id: KeyID,
    /// actual response data, might be encrypted
    pub data: RespInner,
}

impl super::HandshakeParsing for Resp {
    fn deserialize(raw: &[u8]) -> Result<HandshakeData, Error> {
        const MIN_PKT_LEN: usize = 68;
        if raw.len() < MIN_PKT_LEN {
            return Err(Error::NotEnoughData);
        }
        let client_key_id: KeyID =
            KeyID(u16::from_le_bytes(raw[0..2].try_into().unwrap()));
        Ok(HandshakeData::DirSync(DirSync::Resp(Self {
            client_key_id,
            data: RespInner::CipherText(raw[KeyID::len()..].len()),
        })))
    }
}

impl Resp {
    /// return the offset of the encrypted data
    /// NOTE: starts from the beginning of the fenrir packet
    pub fn encrypted_offset(&self) -> usize {
        ProtocolVersion::len()
            + crate::connection::handshake::HandshakeID::len()
            + KeyID::len()
    }
    /// return the total length of the cleartext data
    pub fn encrypted_length(&self) -> usize {
        match &self.data {
            RespInner::ClearText(_data) => RespData::len(),
            _ => 0,
        }
    }
    /// Total length of the response handshake
    pub fn len(&self, head_len: HeadLen, tag_len: TagLen) -> usize {
        KeyID::len() + head_len.0 + self.data.len() + tag_len.0
    }
    /// Serialize into raw bytes
    /// NOTE: assumes that there is exactly as much buffer as needed
    pub fn serialize(
        &self,
        head_len: HeadLen,
        _tag_len: TagLen,
        out: &mut [u8],
    ) {
        out[0..2].copy_from_slice(&self.client_key_id.0.to_le_bytes());
        let start_data = 2 + head_len.0;
        let end_data = start_data + self.data.len();
        self.data.serialize(&mut out[start_data..end_data]);
    }
}

/// Decrypted response data
#[derive(Debug, Clone, PartialEq)]
pub struct RespData {
    /// Client nonce, copied from the request
    pub client_nonce: Nonce,
    /// Server Connection ID
    pub id: ID,
    /// Service Connection ID
    pub service_connection_id: ID,
    /// Service encryption key
    pub service_key: Secret,
}

impl RespData {
    /// Return the expected length for buffer allocation
    pub fn len() -> usize {
        Nonce::len() + ID::len() + ID::len() + Secret::len()
    }
    /// Serialize the data into a buffer
    /// NOTE: assumes that there is exactly asa much buffer as needed
    pub fn serialize(&self, out: &mut [u8]) {
        let mut start = 0;
        let mut end = Nonce::len();
        out[start..end].copy_from_slice(&self.client_nonce.0);
        start = end;
        end = end + ID::len();
        self.id.serialize(&mut out[start..end]);
        start = end;
        end = end + ID::len();
        self.service_connection_id.serialize(&mut out[start..end]);
        start = end;
        end = end + Secret::len();
        out[start..end].copy_from_slice(self.service_key.as_ref());
    }
    /// Parse the cleartext raw data
    pub fn deserialize(raw: &[u8]) -> Result<Self, Error> {
        let raw_sized: &[u8; 16] = raw[..Nonce::len()].try_into().unwrap();
        let client_nonce: Nonce = raw_sized.into();
        let end = Nonce::len() + ID::len();
        let id: ID =
            u64::from_le_bytes(raw[Nonce::len()..end].try_into().unwrap())
                .into();
        if id.is_handshake() {
            return Err(Error::Parsing);
        }
        let parsed = end;
        let end = parsed + ID::len();
        let service_connection_id: ID =
            u64::from_le_bytes(raw[parsed..end].try_into().unwrap()).into();
        if service_connection_id.is_handshake() {
            return Err(Error::Parsing);
        }
        let parsed = end;
        let end = parsed + Secret::len();
        let raw_secret: &[u8; 32] = raw[parsed..end].try_into().unwrap();
        let service_key = raw_secret.into();

        Ok(Self {
            client_nonce,
            id,
            service_connection_id,
            service_key,
        })
    }
}