artisan_keystore 2.1.1

//! In-memory key store with persistence and auditing capabilities.
//!
//! The [`KeyStore`] struct manages active and archived keys, handles periodic
//! updates to key status and provides helper methods for migrating and storing
//! keys on disk.

use crate::shared::consts::STOREPATH;
use dusa_collection_utils::{
    core::functions::current_timestamp,
    core::logger::LogLevel,
    core::types::{pathtype::PathType, stringy::Stringy},
    log,
};
use rand::Rng;
use std::{collections::HashMap, fs};

use super::{
    audit::{AuditEvent, AuditLog},
    key::{Key, KeyStatus},
};
use crate::encryption::enc::{decrypt_with_embedded_key, encrypt_with_embedded_key};

#[derive(Clone)]
/// Central structure holding all keys and related metadata.
pub struct KeyStore {
    pub keys: HashMap<String, Key>,               // Active keys
    pub archived_keys: HashMap<String, Vec<Key>>, // Archived keys grouped by ID
    pub audit_logs: Vec<AuditLog>,
    pub file_path: PathType,
}

impl KeyStore {
    /// Initialise a new [`KeyStore`] loading any stored keys from disk.
    pub async fn new() -> Self {
        let keys: HashMap<String, Key> = match Self::load_keys().await {
            Ok(data) => data,
            Err(err) => {
                log!(LogLevel::Error, "{}", err);
                HashMap::new()
            }
        };
        let file_path: PathType = PathType::Str(STOREPATH.into());
        Self {
            keys,
            archived_keys: HashMap::new(),
            audit_logs: vec![],
            file_path,
        }
    }

    /// Load keys from the persistent store at startup.
    pub async fn load_keys() -> Result<HashMap<String, Key>, Box<dyn std::error::Error>> {
        let encrypted_content: Stringy =
            fs::read_to_string(PathType::Str(STOREPATH.into()))?.into();
        let content = decrypt_with_embedded_key(encrypted_content.as_bytes()).map_err(|_| {
            std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "Decryption failed".to_string(),
            )
        })?;

        let keys: HashMap<String, Key> = serde_json::from_slice(&content)?;
        Ok(keys)
    }

    /// Persist all active keys to the configured storage file.
    pub async fn save_keys(&self) -> Result<(), Box<dyn std::error::Error>> {
        let data: Stringy = serde_json::to_string(&self.keys)?.into();

        let key_data = encrypt_with_embedded_key(data.as_bytes())
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;

        fs::write(&self.file_path, key_data.to_string())?;
        Ok(())
    }

    /// Insert a new key into the store.
    pub fn add_key(&mut self, id: String, value: Vec<u8>, ttl: Option<u64>) {
        let ttl = ttl.unwrap_or(24 * 60 * 60);
        let created_at = current_timestamp();
        let expires_at = current_timestamp() + ttl;
        let legnth = value.len();
        let key = Key {
            id: id.clone(),
            value,
            created_at,
            expires_at,
            version: 1,
            status: KeyStatus::Active,
            legnth,
            ttl,
        };
        self.keys.insert(id.clone(), key);
        self.log_event(AuditEvent::Creation, &id, Some("New key added".to_string()));
    }

    /// Retrieve a key by ID and optionally by a specific version.
    pub fn get_key(&mut self, key_id: &str, version: Option<u32>) -> Option<Key> {
        // Search active keys first
        if let Some(key) = self.clone().keys.get(key_id) {
            if version.is_none() || version == Some(key.version) {
                self.log_event(
                    AuditEvent::Access,
                    key_id,
                    Some("Retrieved active key".to_string()),
                );
                return Some(key.clone());
            }
        }

        // Search archived keys if version is specified or not found in active keys
        if let Some(archived_versions) = self.clone().archived_keys.get(key_id) {
            for archived_key in archived_versions {
                if version == Some(archived_key.version) {
                    self.log_event(
                        AuditEvent::Access,
                        key_id,
                        Some(format!(
                            "Retrieved archived key, version {}",
                            archived_key.version
                        )),
                    );
                    return Some(archived_key.clone());
                }
            }
        }

        None
    }

    /// Move a key from the active set into the archive when it is replaced.
    pub fn archive_key(&mut self, key_id: &str) {
        if let Some(key) = self.keys.remove(key_id) {
            self.archived_keys
                .entry(key_id.to_string())
                .or_default()
                .push(key);
        }
    }

    /// Remove a key entirely from the archive by its ID.
    pub fn _remove_key(&mut self, key_id: &str) {
        self.archived_keys.retain(|id, _| id != key_id);
    }

    /// Remove archived keys older than the configured retention period.
    pub fn purge_old_archives(&mut self, retention_days: i64) {
        let retention_time: i64 = retention_days * 24 * 60 * 60;
        let cutoff_date: u64 = current_timestamp() - retention_time as u64;
        for keys in self.archived_keys.values_mut() {
            keys.retain(|key| key.created_at > cutoff_date);
        }
    }

    /// Update the [`KeyStatus`] of all active keys based on their expiration
    /// timestamp.
    pub fn update_key_status(&mut self) {
        let now = current_timestamp();
        for key in self.keys.values_mut() {
            if key.expires_at <= now {
                key.status = KeyStatus::Expired;
            } else if key.expires_at - (60 * 60 * 24) <= now {
                key.status = KeyStatus::NearExpiration;
            } else {
                key.status = KeyStatus::Active;
            }
        }
    }

    /// Refresh key statuses and migrate any key that is now due for
    /// rotation (`NearExpiration` or `Expired`) to a new version.
    ///
    /// Returns the ids of the keys that were rotated, so callers can decide
    /// whether to persist the store immediately rather than waiting for the
    /// next scheduled save.
    pub fn rotate_expiring_keys(&mut self) -> Vec<String> {
        self.update_key_status();

        let due: Vec<String> = self
            .keys
            .iter()
            .filter(|(_, key)| {
                matches!(key.status, KeyStatus::NearExpiration | KeyStatus::Expired)
            })
            .map(|(id, _)| id.clone())
            .collect();

        for id in &due {
            self.migrate_key(id);
        }

        due
    }

    /// Replace a key with a newly generated version when it is near expiration.
    pub fn migrate_key(&mut self, key_id: &str) {
        if let Some(key) = self.clone().keys.get_mut(key_id) {
            if matches!(key.status, KeyStatus::NearExpiration | KeyStatus::Expired) {
                // Create a new version
                let new_key = Key {
                    id: key.id.clone(),
                    value: generate_new_key(Some(key.legnth)),
                    created_at: current_timestamp(),
                    expires_at: current_timestamp() + key.ttl, // expiration set 30 days out
                    version: key.version + 1,
                    status: KeyStatus::Active,
                    legnth: key.legnth,
                    ttl: key.ttl,
                };

                // Archive the old key
                self.archive_key(key_id);

                // Replace with the new key
                self.keys.insert(key_id.to_string(), new_key);
                self.log_event(
                    AuditEvent::Migration,
                    key_id,
                    Some("Key migrated to new version".to_string()),
                );
            }
        }
    }

    /// Record an [`AuditLog`] entry describing an operation performed against a
    /// key.
    pub fn log_event(&mut self, event: AuditEvent, key_id: &str, details: Option<String>) {
        self.audit_logs.push(AuditLog {
            event,
            key_id: key_id.to_string(),
            timestamp: current_timestamp(),
            details,
        });
    }
}

/// Generate a random sequence of bytes with an optional length.
pub fn generate_new_key(length: Option<usize>) -> Vec<u8> {
    let key_length = length.unwrap_or(32); // Default length is 32 characters

    let mut encrypted_key_buffer: Vec<u8> = Vec::with_capacity(key_length); // Allocate buffer for key
    let mut rng = rand::thread_rng(); // Create a random number generator

    for _ in 0..key_length {
        encrypted_key_buffer.push(rng.r#gen());
    }

    // // Convert the buffer into a hexadecimal string
    encrypted_key_buffer
}