spacedls 0.4.0

use core::cell::Cell;
use core::fmt::{Debug, Formatter};
use hybrid_array::{Array, ArraySize};

use super::error::KeyManagementError;
use crate::core::Key;

/// Lifecycle state of a [`ManagedKey`] (CCSDS 355.1-B-1 Section 2.3.2).
///
/// ```text
/// PreActive --activate()--> Active --deactivate()--> Deactivated
///                             |                          |
///                             +--------destroy()-------->+-> Destroyed
///                             |
///                             +--------rekey()---------> PreActive
/// ```
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum KeyState {
    PreActive,
    Active,
    Deactivated,
    Destroyed,
}

/// Key with lifecycle state management (CCSDS 355.1-B-1 Section 2.3.2).
///
/// Key material is zeroized on [`destroy`](Self::destroy).
/// Implements [`Key`]: returns `Some` only in the `Active` state.
pub struct ManagedKey<N: ArraySize> {
    id: u16,
    state: Cell<KeyState>,
    material: Cell<Array<u8, N>>,
}

impl<N: ArraySize> Debug for ManagedKey<N> {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("ManagedKey")
            .field("id", &self.id)
            .field("state", &self.state.get())
            .field("material", &"***")
            .finish()
    }
}

impl<N: ArraySize> ManagedKey<N> {
    #[inline]
    pub const fn new(id: u16, material: Array<u8, N>) -> Self {
        Self { id, state: Cell::new(KeyState::PreActive), material: Cell::new(material) }
    }
    #[inline]
    #[must_use]
    pub fn id(&self) -> u16 { self.id }
    #[inline]
    #[must_use]
    pub fn state(&self) -> KeyState { self.state.get() }
    #[inline]
    #[must_use]
    pub fn is_active(&self) -> bool { matches!(self.state.get(), KeyState::Active) }
    #[inline]
    #[must_use]
    pub fn is_deactivated(&self) -> bool { matches!(self.state.get(), KeyState::Deactivated) }
    #[inline]
    #[must_use]
    pub fn is_destroyed(&self) -> bool { matches!(self.state.get(), KeyState::Destroyed) }
    #[inline]
    #[must_use]
    pub fn is_preactive(&self) -> bool { matches!(self.state.get(), KeyState::PreActive) }
}

impl<N: ArraySize> ManagedKey<N>
where N::ArrayType<u8>: Copy
{
    #[inline]
    pub fn activate(&self) -> Result<(), KeyManagementError> {
        if self.is_preactive() {
            self.state.set(KeyState::Active);
            Ok(())
        } else {
            Err(KeyManagementError::InvalidTransition {
                from: self.state.get(),
                to: KeyState::Active,
            })
        }
    }

    #[inline]
    pub fn deactivate(&self) -> Result<(), KeyManagementError> {
        if self.is_active() {
            self.state.set(KeyState::Deactivated);
            Ok(())
        } else {
            Err(KeyManagementError::InvalidTransition {
                from: self.state.get(),
                to: KeyState::Deactivated,
            })
        }
    }

    #[inline]
    pub fn destroy(&self) -> Result<(), KeyManagementError> {
        if self.is_active() | self.is_deactivated() {
            self.material.set(Array::default());
            self.state.set(KeyState::Destroyed);
            Ok(())
        } else {
            Err(KeyManagementError::InvalidTransition {
                from: self.state.get(),
                to: KeyState::Destroyed,
            })
        }
    }

    #[inline]
    pub fn rekey(&self, material: Array<u8, N>) -> Result<(), KeyManagementError> {
        if self.is_active() {
            self.material.set(material);
            self.state.set(KeyState::PreActive);
            Ok(())
        } else {
            Err(KeyManagementError::NotActive)
        }
    }

    #[inline]
    pub fn get_active(&self) -> Result<Array<u8, N>, KeyManagementError> {
        if self.is_active() { Ok(self.material.get()) } else { Err(KeyManagementError::NotActive) }
    }
}

impl<N: ArraySize> Key<N> for ManagedKey<N>
where N::ArrayType<u8>: Copy
{
    fn size(&self) -> usize { N::USIZE }

    fn get(&self) -> Option<Array<u8, N>> {
        if matches!(self.state.get(), KeyState::Active) { Some(self.material.get()) } else { None }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::BIT128;

    type TestKey = ManagedKey<BIT128>;

    fn test_key() -> TestKey { TestKey::new(0x0001, [0xAA; 16].into()) }

    #[test]
    fn new_key_is_preactive() {
        let k = test_key();
        assert_eq!(k.id(), 0x0001);
        assert_eq!(k.state(), KeyState::PreActive);
    }

    #[test]
    fn activate_from_preactive() {
        let k = test_key();
        assert!(k.activate().is_ok());
        assert_eq!(k.state(), KeyState::Active);
    }

    #[test]
    fn activate_from_active_fails() {
        let k = test_key();
        k.activate().unwrap();
        assert_eq!(
            k.activate(),
            Err(KeyManagementError::InvalidTransition {
                from: KeyState::Active,
                to: KeyState::Active
            })
        );
    }

    #[test]
    fn deactivate_from_active() {
        let k = test_key();
        k.activate().unwrap();
        assert!(k.deactivate().is_ok());
        assert_eq!(k.state(), KeyState::Deactivated);
    }

    #[test]
    fn deactivate_from_preactive_fails() {
        let k = test_key();
        assert_eq!(
            k.deactivate(),
            Err(KeyManagementError::InvalidTransition {
                from: KeyState::PreActive,
                to: KeyState::Deactivated
            })
        );
    }

    #[test]
    fn destroy_from_active() {
        let k = test_key();
        k.activate().unwrap();
        assert!(k.destroy().is_ok());
        assert_eq!(k.state(), KeyState::Destroyed);
    }

    #[test]
    fn destroy_from_deactivated() {
        let k = test_key();
        k.activate().unwrap();
        k.deactivate().unwrap();
        assert!(k.destroy().is_ok());
        assert_eq!(k.state(), KeyState::Destroyed);
    }

    #[test]
    fn destroy_zeroizes_material() {
        let k = test_key();
        k.activate().unwrap();
        k.destroy().unwrap();
        assert_eq!(k.material.get(), Array::default());
    }

    #[test]
    fn get_active_returns_material() {
        let k = test_key();
        k.activate().unwrap();
        assert_eq!(k.get_active().unwrap(), Array::from([0xAA; 16]));
    }

    #[test]
    fn get_active_fails_when_not_active() {
        let k = test_key();
        assert_eq!(k.get_active(), Err(KeyManagementError::NotActive));
    }

    #[test]
    fn rekey_replaces_material_and_sets_preactive() {
        let k = test_key();
        k.activate().unwrap();
        k.rekey([0xBB; 16].into()).unwrap();
        assert_eq!(k.state(), KeyState::PreActive);
        k.activate().unwrap();
        assert_eq!(k.get_active().unwrap(), Array::from([0xBB; 16]));
    }

    #[test]
    fn rekey_fails_when_not_active() {
        let k = test_key();
        assert_eq!(k.rekey([0xCC; 16].into()), Err(KeyManagementError::NotActive));
    }

    #[test]
    fn key_trait_impl_returns_material() {
        let k = test_key();
        k.activate().unwrap();
        let val: Array<u8, BIT128> = Key::get(&k).unwrap();
        assert_eq!(val, Array::from([0xAA; 16]));
    }
}