rust_crypto_utils/
keymanagement.rs

1//! Secure key management and storage
2
3use crate::{CryptoError, SecureKey};
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::time::{SystemTime, UNIX_EPOCH};
7
8/// Key metadata
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct KeyMetadata {
11    pub key_id: String,
12    pub created_at: u64,
13    pub expires_at: Option<u64>,
14    pub algorithm: String,
15    pub purpose: String,
16    pub rotation_count: u32,
17}
18
19impl KeyMetadata {
20    /// Create new metadata
21    pub fn new(key_id: String, algorithm: String, purpose: String) -> Self {
22        let created_at = SystemTime::now()
23            .duration_since(UNIX_EPOCH)
24            .unwrap()
25            .as_secs();
26
27        Self {
28            key_id,
29            created_at,
30            expires_at: None,
31            algorithm,
32            purpose,
33            rotation_count: 0,
34        }
35    }
36
37    /// Set expiration (seconds from now)
38    pub fn with_expiration(mut self, seconds: u64) -> Self {
39        let now = SystemTime::now()
40            .duration_since(UNIX_EPOCH)
41            .unwrap()
42            .as_secs();
43        self.expires_at = Some(now + seconds);
44        self
45    }
46
47    /// Check if key has expired
48    pub fn is_expired(&self) -> bool {
49        if let Some(expires_at) = self.expires_at {
50            let now = SystemTime::now()
51                .duration_since(UNIX_EPOCH)
52                .unwrap()
53                .as_secs();
54            now >= expires_at
55        } else {
56            false
57        }
58    }
59
60    /// Get age in seconds
61    pub fn age_seconds(&self) -> u64 {
62        let now = SystemTime::now()
63            .duration_since(UNIX_EPOCH)
64            .unwrap()
65            .as_secs();
66        now.saturating_sub(self.created_at)
67    }
68}
69
70/// In-memory key store with automatic cleanup
71pub struct KeyStore {
72    keys: HashMap<String, (SecureKey, KeyMetadata)>,
73}
74
75impl KeyStore {
76    /// Create new key store
77    pub fn new() -> Self {
78        Self {
79            keys: HashMap::new(),
80        }
81    }
82
83    /// Store a key with metadata
84    pub fn store_key(
85        &mut self,
86        key_id: String,
87        key: SecureKey,
88        metadata: KeyMetadata,
89    ) -> Result<(), CryptoError> {
90        if self.keys.contains_key(&key_id) {
91            return Err(CryptoError::HashingError(
92                "Key ID already exists".to_string(),
93            ));
94        }
95        self.keys.insert(key_id, (key, metadata));
96        Ok(())
97    }
98
99    /// Retrieve a key (returns clone of metadata but reference to key)
100    pub fn get_key(&self, key_id: &str) -> Option<(&SecureKey, KeyMetadata)> {
101        self.keys.get(key_id).map(|(key, meta)| (key, meta.clone()))
102    }
103
104    /// Remove a key
105    pub fn remove_key(&mut self, key_id: &str) -> Option<(SecureKey, KeyMetadata)> {
106        self.keys.remove(key_id)
107    }
108
109    /// List all key IDs
110    pub fn list_keys(&self) -> Vec<String> {
111        self.keys.keys().cloned().collect()
112    }
113
114    /// Clean up expired keys
115    pub fn cleanup_expired(&mut self) -> usize {
116        let expired: Vec<String> = self
117            .keys
118            .iter()
119            .filter(|(_, (_, meta))| meta.is_expired())
120            .map(|(id, _)| id.clone())
121            .collect();
122
123        let count = expired.len();
124        for key_id in expired {
125            self.keys.remove(&key_id);
126        }
127        count
128    }
129
130    /// Get count of stored keys
131    pub fn count(&self) -> usize {
132        self.keys.len()
133    }
134
135    /// Rotate a key (generate new key, increment rotation counter)
136    pub fn rotate_key(&mut self, key_id: &str) -> Result<(), CryptoError> {
137        if let Some((_, meta)) = self.keys.get(key_id) {
138            let new_key = SecureKey::generate();
139            let mut new_meta = meta.clone();
140            new_meta.rotation_count += 1;
141            new_meta.created_at = SystemTime::now()
142                .duration_since(UNIX_EPOCH)
143                .unwrap()
144                .as_secs();
145
146            self.keys.insert(key_id.to_string(), (new_key, new_meta));
147            Ok(())
148        } else {
149            Err(CryptoError::HashingError("Key not found".to_string()))
150        }
151    }
152
153    /// Find keys by purpose
154    pub fn find_by_purpose(&self, purpose: &str) -> Vec<(String, KeyMetadata)> {
155        self.keys
156            .iter()
157            .filter(|(_, (_, meta))| meta.purpose == purpose)
158            .map(|(id, (_, meta))| (id.clone(), meta.clone()))
159            .collect()
160    }
161
162    /// Get keys requiring rotation (older than specified age in seconds)
163    pub fn keys_requiring_rotation(&self, max_age_seconds: u64) -> Vec<String> {
164        self.keys
165            .iter()
166            .filter(|(_, (_, meta))| meta.age_seconds() > max_age_seconds)
167            .map(|(id, _)| id.clone())
168            .collect()
169    }
170}
171
172impl Default for KeyStore {
173    fn default() -> Self {
174        Self::new()
175    }
176}
177
178/// Key rotation policy
179#[derive(Debug, Clone)]
180pub struct RotationPolicy {
181    /// Maximum key age in seconds before rotation required
182    pub max_age_seconds: u64,
183    /// Automatically rotate keys when max age is reached
184    pub auto_rotate: bool,
185}
186
187impl RotationPolicy {
188    /// Create policy with 90-day rotation
189    pub fn ninety_days() -> Self {
190        Self {
191            max_age_seconds: 90 * 24 * 60 * 60,
192            auto_rotate: false,
193        }
194    }
195
196    /// Create policy with 30-day rotation
197    pub fn thirty_days() -> Self {
198        Self {
199            max_age_seconds: 30 * 24 * 60 * 60,
200            auto_rotate: false,
201        }
202    }
203
204    /// Create policy with custom days
205    pub fn custom_days(days: u64) -> Self {
206        Self {
207            max_age_seconds: days * 24 * 60 * 60,
208            auto_rotate: false,
209        }
210    }
211
212    /// Enable auto-rotation
213    pub fn with_auto_rotate(mut self) -> Self {
214        self.auto_rotate = true;
215        self
216    }
217
218    /// Check if key needs rotation
219    pub fn needs_rotation(&self, metadata: &KeyMetadata) -> bool {
220        metadata.age_seconds() > self.max_age_seconds
221    }
222}
223
224#[cfg(test)]
225mod tests {
226    use super::*;
227
228    #[test]
229    fn test_key_metadata_creation() {
230        let meta = KeyMetadata::new(
231            "key-001".to_string(),
232            "AES-256-GCM".to_string(),
233            "encryption".to_string(),
234        );
235
236        assert_eq!(meta.key_id, "key-001");
237        assert_eq!(meta.algorithm, "AES-256-GCM");
238        assert_eq!(meta.rotation_count, 0);
239        assert!(!meta.is_expired());
240    }
241
242    #[test]
243    fn test_key_metadata_expiration() {
244        let meta = KeyMetadata::new(
245            "key-002".to_string(),
246            "AES-256-GCM".to_string(),
247            "encryption".to_string(),
248        )
249        .with_expiration(1); // Expires in 1 second
250
251        assert!(!meta.is_expired());
252        std::thread::sleep(std::time::Duration::from_secs(2));
253        assert!(meta.is_expired());
254    }
255
256    #[test]
257    fn test_key_store_operations() {
258        let mut store = KeyStore::new();
259        let key = SecureKey::generate();
260        let meta = KeyMetadata::new(
261            "test-key".to_string(),
262            "AES-256".to_string(),
263            "encryption".to_string(),
264        );
265
266        // Store key
267        assert!(store.store_key("test-key".to_string(), key, meta).is_ok());
268        assert_eq!(store.count(), 1);
269
270        // Retrieve key
271        assert!(store.get_key("test-key").is_some());
272
273        // List keys
274        let keys = store.list_keys();
275        assert_eq!(keys.len(), 1);
276        assert!(keys.contains(&"test-key".to_string()));
277    }
278
279    #[test]
280    fn test_key_store_duplicate_prevention() {
281        let mut store = KeyStore::new();
282        let key1 = SecureKey::generate();
283        let key2 = SecureKey::generate();
284        let meta = KeyMetadata::new(
285            "dup-key".to_string(),
286            "AES-256".to_string(),
287            "encryption".to_string(),
288        );
289
290        assert!(store
291            .store_key("dup-key".to_string(), key1, meta.clone())
292            .is_ok());
293        assert!(store.store_key("dup-key".to_string(), key2, meta).is_err());
294    }
295
296    #[test]
297    fn test_key_cleanup() {
298        let mut store = KeyStore::new();
299
300        // Add expired key
301        let meta_expired = KeyMetadata::new(
302            "expired".to_string(),
303            "AES-256".to_string(),
304            "encryption".to_string(),
305        )
306        .with_expiration(1);
307        store
308            .store_key("expired".to_string(), SecureKey::generate(), meta_expired)
309            .unwrap();
310
311        // Add non-expired key
312        let meta_valid = KeyMetadata::new(
313            "valid".to_string(),
314            "AES-256".to_string(),
315            "encryption".to_string(),
316        )
317        .with_expiration(3600);
318        store
319            .store_key("valid".to_string(), SecureKey::generate(), meta_valid)
320            .unwrap();
321
322        std::thread::sleep(std::time::Duration::from_secs(2));
323
324        let removed = store.cleanup_expired();
325        assert_eq!(removed, 1);
326        assert_eq!(store.count(), 1);
327        assert!(store.get_key("valid").is_some());
328        assert!(store.get_key("expired").is_none());
329    }
330
331    #[test]
332    fn test_key_rotation() {
333        let mut store = KeyStore::new();
334        let key = SecureKey::generate();
335        let meta = KeyMetadata::new(
336            "rotate-key".to_string(),
337            "AES-256".to_string(),
338            "encryption".to_string(),
339        );
340
341        store
342            .store_key("rotate-key".to_string(), key, meta)
343            .unwrap();
344
345        // Rotate
346        assert!(store.rotate_key("rotate-key").is_ok());
347
348        // Check rotation count increased
349        let (_, meta) = store.get_key("rotate-key").unwrap();
350        assert_eq!(meta.rotation_count, 1);
351    }
352
353    #[test]
354    fn test_find_by_purpose() {
355        let mut store = KeyStore::new();
356
357        let meta1 = KeyMetadata::new(
358            "enc-1".to_string(),
359            "AES-256".to_string(),
360            "encryption".to_string(),
361        );
362        let meta2 = KeyMetadata::new(
363            "sig-1".to_string(),
364            "HMAC".to_string(),
365            "signing".to_string(),
366        );
367        let meta3 = KeyMetadata::new(
368            "enc-2".to_string(),
369            "AES-256".to_string(),
370            "encryption".to_string(),
371        );
372
373        store
374            .store_key("enc-1".to_string(), SecureKey::generate(), meta1)
375            .unwrap();
376        store
377            .store_key("sig-1".to_string(), SecureKey::generate(), meta2)
378            .unwrap();
379        store
380            .store_key("enc-2".to_string(), SecureKey::generate(), meta3)
381            .unwrap();
382
383        let encryption_keys = store.find_by_purpose("encryption");
384        assert_eq!(encryption_keys.len(), 2);
385
386        let signing_keys = store.find_by_purpose("signing");
387        assert_eq!(signing_keys.len(), 1);
388    }
389
390    #[test]
391    fn test_rotation_policy() {
392        let policy = RotationPolicy::ninety_days();
393        assert_eq!(policy.max_age_seconds, 90 * 24 * 60 * 60);
394
395        let meta = KeyMetadata::new(
396            "key".to_string(),
397            "AES-256".to_string(),
398            "encryption".to_string(),
399        );
400        assert!(!policy.needs_rotation(&meta));
401    }
402
403    #[test]
404    fn test_keys_requiring_rotation() {
405        let mut store = KeyStore::new();
406
407        let meta = KeyMetadata::new(
408            "old-key".to_string(),
409            "AES-256".to_string(),
410            "encryption".to_string(),
411        );
412        store
413            .store_key("old-key".to_string(), SecureKey::generate(), meta)
414            .unwrap();
415
416        std::thread::sleep(std::time::Duration::from_secs(2));
417
418        let keys = store.keys_requiring_rotation(1); // Keys older than 1 second
419        assert_eq!(keys.len(), 1);
420    }
421}