spacedls 0.4.0

use core::fmt::{Debug, Formatter};
use hybrid_array::{Array, ArraySize};
use typenum::{Unsigned, consts::*};

use super::mac::Mac;

seq_macro::seq!(N in 0..=64 {
    #(
        #[doc = concat!(stringify!(N), "-byte field length (`typenum::U", stringify!(N), "`).")]
        pub type BYTE~N = U~N;
    )*
});

/// Marker trait for valid sequence number lengths (0, 2..=8 bytes).
pub trait ValidSNLen: ArraySize {}
/// Marker trait for valid initialization vector lengths (0, 1..=32 bytes).
pub trait ValidIVLen: ArraySize {}
/// Marker trait for valid pad length field sizes (0, 1..=2 bytes).
pub trait ValidPLLen: ArraySize {}
/// Marker trait for valid MAC lengths (0, 8..=64 bytes).
pub trait ValidMacLen: ArraySize {}

macro_rules! impl_valid_range {
    ($trait:ident for 0 | $start:literal ..= $end:literal) => {
        impl $trait for U0 {}
        seq_macro::seq!(N in $start..=$end {
            #(impl $trait for U~N {})*
        });
    };
    ($trait:ident for $start:literal ..= $end:literal) => {
        seq_macro::seq!(N in $start..=$end {
            #(impl $trait for ::typenum::consts::U~N {})*
        });
    };
}

impl_valid_range!(ValidIVLen  for 0 | 1..=32);
impl_valid_range!(ValidSNLen  for 0 | 2..=8);
impl_valid_range!(ValidPLLen  for 0 | 1..=2);
impl_valid_range!(ValidMacLen for 0 | 8..=64);

/// Compile-time definition of an SDLS security header/trailer layout
/// (CCSDS 355.0-B-2 Section 4.1.1--4.1.2).
///
/// Implement this trait on a unit struct to define a mission-specific frame format.
/// `HeaderLen` must equal `2 + IVLen + SNLen + PLLen`; this is checked at compile time.
///
/// # Example
///
/// ```rust,ignore
/// use spacedls::consts::*;
///
/// struct TcAeadFmt;
/// impl SDLSFrameFormat for TcAeadFmt {
///     type SNLen  = BYTE4;
///     type IVLen  = BYTE12;
///     type PLLen  = BYTE0;
///     type MacLen = BYTE16;
///     type HeaderLen = BYTE18; // 2 + 12 + 4 + 0
/// }
/// ```
pub trait SDLSFrameFormat {
    type SNLen: ValidSNLen;
    type IVLen: ValidIVLen;
    type PLLen: ValidPLLen;
    type MacLen: ValidMacLen;
    type HeaderLen: ArraySize;

    #[doc(hidden)]
    const _ASSERT_HEADER_LEN: () = assert!(
        Self::HeaderLen::USIZE == 2 + Self::IVLen::USIZE + Self::SNLen::USIZE + Self::PLLen::USIZE
    );
}

/// Parsed SDLS Security Header (CCSDS 355.0-B-2 Section 4.1.1).
///
/// Contains the Security Parameter Index (SPI), IV, sequence number, and pad length.
/// Serializes to/from a fixed-size byte array of length `F::HeaderLen`.
pub struct SecurityHeader<F: SDLSFrameFormat> {
    pub spi: u16,
    pub iv: Array<u8, F::IVLen>,
    pub sn: Array<u8, F::SNLen>,
    pub pad_len: u16,
}

/// SDLS Security Trailer (CCSDS 355.0-B-2 Section 4.1.2).
///
/// Wraps the MAC produced by authentication or authenticated-encryption operations.
pub struct SecurityTrailer<F: SDLSFrameFormat> {
    pub mac: Mac<F::MacLen>,
}

impl<F: SDLSFrameFormat> Clone for SecurityHeader<F> {
    fn clone(&self) -> Self {
        Self { spi: self.spi, iv: self.iv.clone(), sn: self.sn.clone(), pad_len: self.pad_len }
    }
}

impl<F: SDLSFrameFormat> Debug for SecurityHeader<F> {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("SecurityHeader")
            .field("spi", &self.spi)
            .field("iv", &self.iv)
            .field("sn", &self.sn)
            .field("pad_len", &self.pad_len)
            .finish()
    }
}

impl<F: SDLSFrameFormat> Clone for SecurityTrailer<F> {
    fn clone(&self) -> Self { Self { mac: self.mac.clone() } }
}

impl<F: SDLSFrameFormat> Debug for SecurityTrailer<F> {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("SecurityTrailer").field("mac", &self.mac).finish()
    }
}

impl<F: SDLSFrameFormat> SecurityTrailer<F> {
    #[inline]
    pub fn from_mac_bytes(mac_bytes: &[u8]) -> Self {
        let mut mac = Array::<u8, F::MacLen>::default();
        let copy_len = mac_bytes.len().min(F::MacLen::USIZE);
        mac[..copy_len].copy_from_slice(&mac_bytes[..copy_len]);
        Self { mac: mac.into() }
    }
}

impl<F: SDLSFrameFormat> SecurityHeader<F> {
    #[inline]
    pub fn serialize(&self) -> Array<u8, F::HeaderLen> {
        let mut out = Array::default();
        let mut pos = 0;

        out[pos..pos + 2].copy_from_slice(&self.spi.to_be_bytes());
        pos += 2;

        let iv_len = F::IVLen::USIZE;
        out[pos..pos + iv_len].copy_from_slice(&self.iv);
        pos += iv_len;

        let sn_len = F::SNLen::USIZE;
        out[pos..pos + sn_len].copy_from_slice(&self.sn);
        pos += sn_len;

        let pl_len = F::PLLen::USIZE;
        let pl_be = self.pad_len.to_be_bytes();
        out[pos..pos + pl_len].copy_from_slice(&pl_be[2 - pl_len..]);

        out
    }

    #[inline]
    pub fn deserialize(bytes: &Array<u8, F::HeaderLen>) -> Self {
        let mut pos = 0;

        let spi = u16::from_be_bytes([bytes[pos], bytes[pos + 1]]);
        pos += 2;

        let iv_len = F::IVLen::USIZE;
        let mut iv = Array::default();
        iv.copy_from_slice(&bytes[pos..pos + iv_len]);
        pos += iv_len;

        let sn_len = F::SNLen::USIZE;
        let mut sn = Array::default();
        sn.copy_from_slice(&bytes[pos..pos + sn_len]);
        pos += sn_len;

        let pl_len = F::PLLen::USIZE;
        let mut pl_buf = [0u8; 2];
        pl_buf[2 - pl_len..].copy_from_slice(&bytes[pos..pos + pl_len]);
        let pad_len = u16::from_be_bytes(pl_buf);

        Self { spi, iv, sn, pad_len }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use hybrid_array::Array;
    use typenum::{U0, U2, U4, U8, U12, U16, U18};

    struct TestFmt;
    impl SDLSFrameFormat for TestFmt {
        type SNLen = U4;
        type IVLen = U12;
        type PLLen = U0;
        type MacLen = U16;
        type HeaderLen = U18;
    }

    struct TestFmtWithPad;
    impl SDLSFrameFormat for TestFmtWithPad {
        type SNLen = U4;
        type IVLen = U8;
        type PLLen = U2;
        type MacLen = U16;
        type HeaderLen = U16;
    }

    #[test]
    fn security_header_serialize_deserialize() {
        let hdr = SecurityHeader::<TestFmt> {
            spi: 0xABCD,
            iv: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12].into(),
            sn: [0x00, 0x01, 0x02, 0x03].into(),
            pad_len: 0,
        };
        let bytes = hdr.serialize();
        let recovered = SecurityHeader::<TestFmt>::deserialize(&bytes);

        assert_eq!(recovered.spi, 0xABCD);
        assert_eq!(recovered.iv, Array::from([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]));
        assert_eq!(recovered.sn, Array::from([0x00, 0x01, 0x02, 0x03]));
        assert_eq!(recovered.pad_len, 0);
    }

    #[test]
    fn security_header_round_trip_with_pad_len() {
        let hdr = SecurityHeader::<TestFmtWithPad> {
            spi: 0x1234,
            iv: [0xA0, 0xA1, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6, 0xA7].into(),
            sn: [0x00, 0x01, 0x02, 0x03].into(),
            pad_len: 7,
        };
        let bytes = hdr.serialize();
        let recovered = SecurityHeader::<TestFmtWithPad>::deserialize(&bytes);

        assert_eq!(recovered.spi, 0x1234);
        assert_eq!(recovered.sn, Array::from([0x00, 0x01, 0x02, 0x03]));
        assert_eq!(recovered.pad_len, 7);
    }

    struct SmallMacFmt;
    impl SDLSFrameFormat for SmallMacFmt {
        type SNLen = U0;
        type IVLen = U0;
        type PLLen = U0;
        type MacLen = U8;
        type HeaderLen = U2;
    }

    #[test]
    fn security_trailer_from_mac_truncates() {
        let long_mac = [0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xAA];
        let trlr = SecurityTrailer::<SmallMacFmt>::from_mac_bytes(&long_mac);
        assert_eq!(&*trlr.mac, &[0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88]);
    }

    #[test]
    fn security_trailer_from_mac_exact() {
        let exact_mac = [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08];
        let trlr = SecurityTrailer::<SmallMacFmt>::from_mac_bytes(&exact_mac);
        assert_eq!(&*trlr.mac, &[0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]);
    }
}