1use crate::{CryptoError, SecureKey};
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::time::{SystemTime, UNIX_EPOCH};
7
8#[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 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 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 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 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
70pub struct KeyStore {
72 keys: HashMap<String, (SecureKey, KeyMetadata)>,
73}
74
75impl KeyStore {
76 pub fn new() -> Self {
78 Self {
79 keys: HashMap::new(),
80 }
81 }
82
83 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 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 pub fn remove_key(&mut self, key_id: &str) -> Option<(SecureKey, KeyMetadata)> {
106 self.keys.remove(key_id)
107 }
108
109 pub fn list_keys(&self) -> Vec<String> {
111 self.keys.keys().cloned().collect()
112 }
113
114 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 pub fn count(&self) -> usize {
132 self.keys.len()
133 }
134
135 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 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 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#[derive(Debug, Clone)]
180pub struct RotationPolicy {
181 pub max_age_seconds: u64,
183 pub auto_rotate: bool,
185}
186
187impl RotationPolicy {
188 pub fn ninety_days() -> Self {
190 Self {
191 max_age_seconds: 90 * 24 * 60 * 60,
192 auto_rotate: false,
193 }
194 }
195
196 pub fn thirty_days() -> Self {
198 Self {
199 max_age_seconds: 30 * 24 * 60 * 60,
200 auto_rotate: false,
201 }
202 }
203
204 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 pub fn with_auto_rotate(mut self) -> Self {
214 self.auto_rotate = true;
215 self
216 }
217
218 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); 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 assert!(store.store_key("test-key".to_string(), key, meta).is_ok());
268 assert_eq!(store.count(), 1);
269
270 assert!(store.get_key("test-key").is_some());
272
273 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 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 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 assert!(store.rotate_key("rotate-key").is_ok());
347
348 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); assert_eq!(keys.len(), 1);
420 }
421}