tap_agent/
agent_key_manager.rs

1//! Agent Key Manager for the TAP Agent
2//!
3//! This module provides an implementation of a key manager that uses the agent key abstraction.
4//! It manages keys for signing, verification, encryption, and decryption operations, with support
5//! for different key types (Ed25519, P-256, secp256k1).
6
7use crate::agent_key::{AgentKey, DecryptionKey, EncryptionKey, SigningKey, VerificationKey};
8use crate::did::{DIDGenerationOptions, DIDKeyGenerator, GeneratedKey, KeyType};
9use crate::error::{Error, Result};
10use crate::key_manager::{KeyManager, Secret, SecretMaterial};
11use crate::local_agent_key::{LocalAgentKey, PublicVerificationKey};
12use crate::message::{JweProtected, JwsProtected};
13use crate::message_packing::{KeyManagerPacking, MessageError};
14use crate::storage::{KeyStorage, StoredKey};
15
16use async_trait::async_trait;
17use base64::Engine;
18use std::collections::HashMap;
19use std::path::PathBuf;
20use std::sync::{Arc, RwLock};
21
22/// Agent Key Manager implements the KeyManager trait using the agent key abstraction
23#[derive(Debug, Clone)]
24pub struct AgentKeyManager {
25    /// The DID key generator
26    generator: DIDKeyGenerator,
27    /// The secret storage (legacy)
28    secrets: Arc<RwLock<HashMap<String, Secret>>>,
29    /// Signing keys
30    signing_keys: Arc<RwLock<HashMap<String, Arc<dyn SigningKey + Send + Sync>>>>,
31    /// Encryption keys
32    encryption_keys: Arc<RwLock<HashMap<String, Arc<dyn EncryptionKey + Send + Sync>>>>,
33    /// Decryption keys
34    decryption_keys: Arc<RwLock<HashMap<String, Arc<dyn DecryptionKey + Send + Sync>>>>,
35    /// Verification keys
36    verification_keys: Arc<RwLock<HashMap<String, Arc<dyn VerificationKey + Send + Sync>>>>,
37    /// Storage path
38    storage_path: Option<PathBuf>,
39}
40
41impl AgentKeyManager {
42    /// Create a new key manager
43    pub fn new() -> Self {
44        Self {
45            generator: DIDKeyGenerator::new(),
46            secrets: Arc::new(RwLock::new(HashMap::new())),
47            signing_keys: Arc::new(RwLock::new(HashMap::new())),
48            encryption_keys: Arc::new(RwLock::new(HashMap::new())),
49            decryption_keys: Arc::new(RwLock::new(HashMap::new())),
50            verification_keys: Arc::new(RwLock::new(HashMap::new())),
51            storage_path: None,
52        }
53    }
54
55    /// Get the key type for a signing key (for debugging)
56    pub async fn get_signing_key_type(&self, did: &str) -> Result<String> {
57        // Try to find a signing key for this DID
58        if let Ok(signing_keys) = self.signing_keys.read() {
59            for (kid, key) in signing_keys.iter() {
60                if kid.starts_with(did) {
61                    if let Ok(jwk) = key.public_key_jwk() {
62                        let kty = jwk.get("kty").and_then(|v| v.as_str());
63                        let crv = jwk.get("crv").and_then(|v| v.as_str());
64                        return Ok(format!("kty: {:?}, crv: {:?}", kty, crv));
65                    }
66                }
67            }
68        }
69
70        Err(Error::KeyNotFound(format!(
71            "No signing key found for DID: {}",
72            did
73        )))
74    }
75
76    /// Create a LocalAgentKey from a GeneratedKey
77    pub fn agent_key_from_generated(&self, key: &GeneratedKey) -> Result<LocalAgentKey> {
78        // Create a secret for the key
79        let secret = self.generator.create_secret_from_key(key);
80
81        // Create a LocalAgentKey
82        Ok(LocalAgentKey::new(secret, key.key_type))
83    }
84
85    /// Store an agent key in all collections
86    fn store_agent_key(&self, agent_key: &LocalAgentKey, key_id: &str) -> Result<()> {
87        // Store the agent key as signing, encryption, and decryption keys
88        if let Ok(mut signing_keys) = self.signing_keys.write() {
89            signing_keys.insert(
90                key_id.to_string(),
91                Arc::new(agent_key.clone()) as Arc<dyn SigningKey + Send + Sync>,
92            );
93        } else {
94            return Err(Error::FailedToAcquireResolverWriteLock);
95        }
96
97        if let Ok(mut encryption_keys) = self.encryption_keys.write() {
98            encryption_keys.insert(
99                key_id.to_string(),
100                Arc::new(agent_key.clone()) as Arc<dyn EncryptionKey + Send + Sync>,
101            );
102        } else {
103            return Err(Error::FailedToAcquireResolverWriteLock);
104        }
105
106        if let Ok(mut decryption_keys) = self.decryption_keys.write() {
107            decryption_keys.insert(
108                key_id.to_string(),
109                Arc::new(agent_key.clone()) as Arc<dyn DecryptionKey + Send + Sync>,
110            );
111        } else {
112            return Err(Error::FailedToAcquireResolverWriteLock);
113        }
114
115        // Also store a reference in verification keys
116        if let Ok(mut verification_keys) = self.verification_keys.write() {
117            verification_keys.insert(
118                key_id.to_string(),
119                Arc::new(agent_key.clone()) as Arc<dyn VerificationKey + Send + Sync>,
120            );
121        } else {
122            return Err(Error::FailedToAcquireResolverWriteLock);
123        }
124
125        Ok(())
126    }
127
128    /// Save keys to storage if a storage path is configured
129    pub fn save_to_storage(&self) -> Result<()> {
130        // Skip if no storage path is configured
131        if self.storage_path.is_none() {
132            return Ok(());
133        }
134
135        // Create a KeyStorage object from our secrets
136        let mut key_storage = KeyStorage::new();
137
138        // Add all secrets
139        if let Ok(secrets) = self.secrets.read() {
140            for (did, secret) in secrets.iter() {
141                // Extract key type from the key
142                let key_type = match secret.secret_material {
143                    SecretMaterial::JWK {
144                        ref private_key_jwk,
145                    } => {
146                        let kty = private_key_jwk.get("kty").and_then(|v| v.as_str());
147                        let crv = private_key_jwk.get("crv").and_then(|v| v.as_str());
148
149                        match (kty, crv) {
150                            (Some("OKP"), Some("Ed25519")) => KeyType::Ed25519,
151                            (Some("EC"), Some("P-256")) => KeyType::P256,
152                            (Some("EC"), Some("secp256k1")) => KeyType::Secp256k1,
153                            _ => KeyType::Ed25519, // Default
154                        }
155                    }
156                };
157
158                // Get private and public keys from the JWK
159                let private_key_b64 = match &secret.secret_material {
160                    SecretMaterial::JWK { private_key_jwk } => private_key_jwk
161                        .get("d")
162                        .and_then(|v| v.as_str())
163                        .unwrap_or("")
164                        .to_string(),
165                };
166
167                let public_key_b64 = match &secret.secret_material {
168                    SecretMaterial::JWK { private_key_jwk } => private_key_jwk
169                        .get("x")
170                        .and_then(|v| v.as_str())
171                        .unwrap_or("")
172                        .to_string(),
173                };
174
175                // Create a StoredKey and add to key storage
176                let stored_key = StoredKey {
177                    did: did.clone(),
178                    key_type,
179                    private_key: private_key_b64,
180                    public_key: public_key_b64,
181                    metadata: HashMap::new(),
182                };
183                key_storage.add_key(stored_key);
184            }
185        }
186
187        // Save to storage
188        if let Some(path) = &self.storage_path {
189            key_storage.save_to_path(path)?;
190        } else {
191            key_storage.save_default()?;
192        }
193
194        Ok(())
195    }
196
197    /// Load from default storage location
198    pub fn load_from_default_storage(mut self) -> Result<Self> {
199        self.storage_path = None;
200        self.load_keys_from_storage()
201    }
202
203    /// Load from a specific storage path
204    pub fn load_from_path(mut self, path: PathBuf) -> Result<Self> {
205        self.storage_path = Some(path);
206        self.load_keys_from_storage()
207    }
208
209    /// Load keys from configured storage
210    fn load_keys_from_storage(&self) -> Result<Self> {
211        // Load storage
212        let storage = if let Some(path) = &self.storage_path {
213            KeyStorage::load_from_path(path)?
214        } else {
215            KeyStorage::load_default()?
216        };
217
218        // Process each stored key
219        for (did, stored_key) in storage.keys {
220            // Convert to a legacy secret
221            let secret = KeyStorage::to_secret(&stored_key);
222
223            // Add to secrets
224            if let Ok(mut secrets) = self.secrets.write() {
225                secrets.insert(did.clone(), secret.clone());
226            } else {
227                return Err(Error::FailedToAcquireResolverWriteLock);
228            }
229
230            // Create an agent key
231            let key_type = stored_key.key_type;
232            let agent_key = LocalAgentKey::new(secret, key_type);
233            let key_id = AgentKey::key_id(&agent_key).to_string();
234
235            // Store in all collections
236            self.store_agent_key(&agent_key, &key_id)?;
237        }
238
239        Ok(self.clone())
240    }
241}
242
243impl Default for AgentKeyManager {
244    fn default() -> Self {
245        Self::new()
246    }
247}
248
249#[async_trait]
250impl KeyManager for AgentKeyManager {
251    /// Get access to the secrets storage
252    fn secrets(&self) -> Arc<RwLock<HashMap<String, Secret>>> {
253        Arc::clone(&self.secrets)
254    }
255
256    /// Generate a new key with the specified options
257    fn generate_key(&self, options: DIDGenerationOptions) -> Result<GeneratedKey> {
258        // Generate the key
259        let key = self.generator.generate_did(options)?;
260
261        // Create a LocalAgentKey
262        let agent_key = self.agent_key_from_generated(&key)?;
263        let key_id = AgentKey::key_id(&agent_key).to_string();
264
265        // Store the legacy secret
266        if let Ok(mut secrets) = self.secrets.write() {
267            secrets.insert(key.did.clone(), agent_key.clone().secret);
268        } else {
269            return Err(Error::FailedToAcquireResolverWriteLock);
270        }
271
272        // Store in all collections
273        self.store_agent_key(&agent_key, &key_id)?;
274
275        // Save to storage if configured
276        self.save_to_storage()?;
277
278        Ok(key)
279    }
280
281    /// Generate a new web DID with the specified domain and options
282    fn generate_web_did(
283        &self,
284        domain: &str,
285        options: DIDGenerationOptions,
286    ) -> Result<GeneratedKey> {
287        // Generate the key
288        let key = self.generator.generate_web_did(domain, options)?;
289
290        // Create a LocalAgentKey
291        let agent_key = self.agent_key_from_generated(&key)?;
292        let key_id = AgentKey::key_id(&agent_key).to_string();
293
294        // Store the legacy secret
295        if let Ok(mut secrets) = self.secrets.write() {
296            secrets.insert(key.did.clone(), agent_key.clone().secret);
297        } else {
298            return Err(Error::FailedToAcquireResolverWriteLock);
299        }
300
301        // Store in all collections
302        self.store_agent_key(&agent_key, &key_id)?;
303
304        // Save to storage if configured
305        self.save_to_storage()?;
306
307        Ok(key)
308    }
309
310    /// Add an existing key to the key manager
311    fn add_key(&self, key: &GeneratedKey) -> Result<()> {
312        // Create a LocalAgentKey
313        let agent_key = self.agent_key_from_generated(key)?;
314        let key_id = AgentKey::key_id(&agent_key).to_string();
315
316        // Store the legacy secret
317        if let Ok(mut secrets) = self.secrets.write() {
318            secrets.insert(key.did.clone(), agent_key.clone().secret);
319        } else {
320            return Err(Error::FailedToAcquireResolverWriteLock);
321        }
322
323        // Store in all collections
324        self.store_agent_key(&agent_key, &key_id)?;
325
326        // Save to storage if configured
327        self.save_to_storage()?;
328
329        Ok(())
330    }
331
332    /// Remove a key from the key manager
333    fn remove_key(&self, did: &str) -> Result<()> {
334        // Remove from legacy secrets
335        if let Ok(mut secrets) = self.secrets.write() {
336            secrets.remove(did);
337        } else {
338            return Err(Error::FailedToAcquireResolverWriteLock);
339        }
340
341        // Remove from signing keys
342        if let Ok(mut signing_keys) = self.signing_keys.write() {
343            signing_keys.retain(|k, _| !k.starts_with(did));
344        } else {
345            return Err(Error::FailedToAcquireResolverWriteLock);
346        }
347
348        // Remove from encryption keys
349        if let Ok(mut encryption_keys) = self.encryption_keys.write() {
350            encryption_keys.retain(|k, _| !k.starts_with(did));
351        } else {
352            return Err(Error::FailedToAcquireResolverWriteLock);
353        }
354
355        // Remove from decryption keys
356        if let Ok(mut decryption_keys) = self.decryption_keys.write() {
357            decryption_keys.retain(|k, _| !k.starts_with(did));
358        } else {
359            return Err(Error::FailedToAcquireResolverWriteLock);
360        }
361
362        // Remove from verification keys
363        if let Ok(mut verification_keys) = self.verification_keys.write() {
364            verification_keys.retain(|k, _| !k.starts_with(did));
365        } else {
366            return Err(Error::FailedToAcquireResolverWriteLock);
367        }
368
369        // Save to storage if configured
370        self.save_to_storage()?;
371
372        Ok(())
373    }
374
375    /// Check if the key manager has a key for the given DID
376    fn has_key(&self, did: &str) -> Result<bool> {
377        // Check legacy secrets first
378        if let Ok(secrets) = self.secrets.read() {
379            if secrets.contains_key(did) {
380                return Ok(true);
381            }
382        } else {
383            return Err(Error::FailedToAcquireResolverReadLock);
384        }
385
386        // Check if any signing key has this DID
387        if let Ok(signing_keys) = self.signing_keys.read() {
388            if signing_keys.values().any(|k| k.did() == did) {
389                return Ok(true);
390            }
391        } else {
392            return Err(Error::FailedToAcquireResolverReadLock);
393        }
394
395        Ok(false)
396    }
397
398    /// Get a list of all DIDs in the key manager
399    fn list_keys(&self) -> Result<Vec<String>> {
400        // Collect DIDs from both legacy secrets and new keys
401        let mut dids = Vec::new();
402
403        // Add DIDs from legacy secrets
404        if let Ok(secrets) = self.secrets.read() {
405            dids.extend(secrets.keys().cloned());
406        } else {
407            return Err(Error::FailedToAcquireResolverReadLock);
408        }
409
410        // Add DIDs from signing keys
411        if let Ok(signing_keys) = self.signing_keys.read() {
412            for key in signing_keys.values() {
413                if !dids.contains(&key.did().to_string()) {
414                    dids.push(key.did().to_string());
415                }
416            }
417        } else {
418            return Err(Error::FailedToAcquireResolverReadLock);
419        }
420
421        Ok(dids)
422    }
423
424    /// Add a signing key to the key manager
425    async fn add_signing_key(&self, key: Arc<dyn SigningKey + Send + Sync>) -> Result<()> {
426        if let Ok(mut signing_keys) = self.signing_keys.write() {
427            signing_keys.insert(key.key_id().to_string(), key);
428            Ok(())
429        } else {
430            Err(Error::FailedToAcquireResolverWriteLock)
431        }
432    }
433
434    /// Add an encryption key to the key manager
435    async fn add_encryption_key(&self, key: Arc<dyn EncryptionKey + Send + Sync>) -> Result<()> {
436        if let Ok(mut encryption_keys) = self.encryption_keys.write() {
437            encryption_keys.insert(key.key_id().to_string(), key);
438            Ok(())
439        } else {
440            Err(Error::FailedToAcquireResolverWriteLock)
441        }
442    }
443
444    /// Add a decryption key to the key manager
445    async fn add_decryption_key(&self, key: Arc<dyn DecryptionKey + Send + Sync>) -> Result<()> {
446        if let Ok(mut decryption_keys) = self.decryption_keys.write() {
447            decryption_keys.insert(key.key_id().to_string(), key);
448            Ok(())
449        } else {
450            Err(Error::FailedToAcquireResolverWriteLock)
451        }
452    }
453
454    /// Get a signing key by ID
455    async fn get_signing_key(&self, kid: &str) -> Result<Arc<dyn SigningKey + Send + Sync>> {
456        // Check if we have a signing key with this ID
457        if let Ok(signing_keys) = self.signing_keys.read() {
458            if let Some(key) = signing_keys.get(kid) {
459                return Ok(key.clone());
460            }
461        } else {
462            return Err(Error::FailedToAcquireResolverReadLock);
463        }
464
465        // If not, check legacy secrets
466        if let Ok(secrets) = self.secrets.read() {
467            // Try to find a secret with this DID or kid
468            let did = kid.split('#').next().unwrap_or(kid);
469            if let Some(secret) = secrets.get(did) {
470                // Create a LocalAgentKey
471                let key_type = KeyType::Ed25519; // Default to Ed25519
472                let agent_key = LocalAgentKey::new(secret.clone(), key_type);
473
474                // Add to signing keys for next time
475                if let Ok(mut signing_keys) = self.signing_keys.write() {
476                    let arc_key = Arc::new(agent_key.clone()) as Arc<dyn SigningKey + Send + Sync>;
477                    signing_keys.insert(AgentKey::key_id(&agent_key).to_string(), arc_key.clone());
478                    return Ok(arc_key);
479                }
480            }
481        }
482
483        Err(Error::Cryptography(format!(
484            "No signing key found with ID: {}",
485            kid
486        )))
487    }
488
489    /// Get an encryption key by ID
490    async fn get_encryption_key(&self, kid: &str) -> Result<Arc<dyn EncryptionKey + Send + Sync>> {
491        // Check if we have an encryption key with this ID
492        if let Ok(encryption_keys) = self.encryption_keys.read() {
493            if let Some(key) = encryption_keys.get(kid) {
494                return Ok(key.clone());
495            }
496        } else {
497            return Err(Error::FailedToAcquireResolverReadLock);
498        }
499
500        // If not, check legacy secrets
501        if let Ok(secrets) = self.secrets.read() {
502            // Try to find a secret with this DID or kid
503            let did = kid.split('#').next().unwrap_or(kid);
504            if let Some(secret) = secrets.get(did) {
505                // Create a LocalAgentKey
506                let key_type = KeyType::Ed25519; // Default to Ed25519
507                let agent_key = LocalAgentKey::new(secret.clone(), key_type);
508
509                // Add to encryption keys for next time
510                if let Ok(mut encryption_keys) = self.encryption_keys.write() {
511                    let arc_key =
512                        Arc::new(agent_key.clone()) as Arc<dyn EncryptionKey + Send + Sync>;
513                    encryption_keys
514                        .insert(AgentKey::key_id(&agent_key).to_string(), arc_key.clone());
515                    return Ok(arc_key);
516                }
517            }
518        }
519
520        Err(Error::Cryptography(format!(
521            "No encryption key found with ID: {}",
522            kid
523        )))
524    }
525
526    /// Get a decryption key by ID
527    async fn get_decryption_key(&self, kid: &str) -> Result<Arc<dyn DecryptionKey + Send + Sync>> {
528        // Check if we have a decryption key with this ID
529        if let Ok(decryption_keys) = self.decryption_keys.read() {
530            if let Some(key) = decryption_keys.get(kid) {
531                return Ok(key.clone());
532            }
533        } else {
534            return Err(Error::FailedToAcquireResolverReadLock);
535        }
536
537        // If not, check legacy secrets
538        if let Ok(secrets) = self.secrets.read() {
539            // Try to find a secret with this DID or kid
540            let did = kid.split('#').next().unwrap_or(kid);
541            if let Some(secret) = secrets.get(did) {
542                // Create a LocalAgentKey
543                let key_type = KeyType::Ed25519; // Default to Ed25519
544                let agent_key = LocalAgentKey::new(secret.clone(), key_type);
545
546                // Add to decryption keys for next time
547                if let Ok(mut decryption_keys) = self.decryption_keys.write() {
548                    let arc_key =
549                        Arc::new(agent_key.clone()) as Arc<dyn DecryptionKey + Send + Sync>;
550                    decryption_keys
551                        .insert(AgentKey::key_id(&agent_key).to_string(), arc_key.clone());
552                    return Ok(arc_key);
553                }
554            }
555        }
556
557        Err(Error::Cryptography(format!(
558            "No decryption key found with ID: {}",
559            kid
560        )))
561    }
562
563    /// Resolve a verification key by ID
564    async fn resolve_verification_key(
565        &self,
566        kid: &str,
567    ) -> Result<Arc<dyn VerificationKey + Send + Sync>> {
568        // Check if we have a verification key with this ID
569        if let Ok(verification_keys) = self.verification_keys.read() {
570            if let Some(key) = verification_keys.get(kid) {
571                return Ok(key.clone());
572            }
573        } else {
574            return Err(Error::FailedToAcquireResolverReadLock);
575        }
576
577        // If not found locally, try to derive from a signing key
578        let signing_key = KeyManager::get_signing_key(self, kid).await;
579        if let Ok(key) = signing_key {
580            // Create a verification key from the signing key
581            let public_jwk = key.public_key_jwk()?;
582            let verification_key = Arc::new(PublicVerificationKey::new(kid.to_string(), public_jwk))
583                as Arc<dyn VerificationKey + Send + Sync>;
584
585            // Add to verification keys for next time
586            if let Ok(mut verification_keys) = self.verification_keys.write() {
587                verification_keys.insert(kid.to_string(), verification_key.clone());
588            }
589
590            return Ok(verification_key);
591        }
592
593        // In a full implementation, we would use a DID Resolver here
594        Err(Error::Cryptography(format!(
595            "No verification key found with ID: {}",
596            kid
597        )))
598    }
599
600    /// Sign data with a key
601    async fn sign_jws(
602        &self,
603        kid: &str,
604        payload: &[u8],
605        protected_header: Option<JwsProtected>,
606    ) -> Result<String> {
607        // Get the signing key
608        let signing_key = KeyManager::get_signing_key(self, kid).await?;
609
610        // Sign the payload
611        let jws = signing_key
612            .create_jws(payload, protected_header)
613            .await
614            .map_err(|e| Error::Cryptography(e.to_string()))?;
615
616        // Serialize the JWS
617        serde_json::to_string(&jws).map_err(|e| Error::Serialization(e.to_string()))
618    }
619
620    /// Verify a JWS
621    async fn verify_jws(&self, jws: &str, expected_kid: Option<&str>) -> Result<Vec<u8>> {
622        // Parse the JWS
623        let jws: crate::message::Jws = serde_json::from_str(jws)
624            .map_err(|e| Error::Serialization(format!("Failed to parse JWS: {}", e)))?;
625
626        // Find the signature to verify
627        let signature = if let Some(kid) = expected_kid {
628            jws.signatures
629                .iter()
630                .find(|s| s.header.kid == kid)
631                .ok_or_else(|| {
632                    Error::Cryptography(format!("No signature found with kid: {}", kid))
633                })?
634        } else {
635            // Use the first signature
636            jws.signatures
637                .first()
638                .ok_or_else(|| Error::Cryptography("No signatures in JWS".to_string()))?
639        };
640
641        // Decode the protected header
642        let protected_bytes = base64::engine::general_purpose::STANDARD
643            .decode(&signature.protected)
644            .map_err(|e| {
645                Error::Cryptography(format!("Failed to decode protected header: {}", e))
646            })?;
647
648        // Parse the protected header
649        let protected: JwsProtected = serde_json::from_slice(&protected_bytes).map_err(|e| {
650            Error::Serialization(format!("Failed to parse protected header: {}", e))
651        })?;
652
653        // Resolve the verification key
654        let verification_key =
655            KeyManager::resolve_verification_key(self, &signature.header.kid).await?;
656
657        // Decode the signature
658        let signature_bytes = base64::engine::general_purpose::STANDARD
659            .decode(&signature.signature)
660            .map_err(|e| Error::Cryptography(format!("Failed to decode signature: {}", e)))?;
661
662        // Create the signing input (protected.payload)
663        let signing_input = format!("{}.{}", signature.protected, jws.payload);
664
665        // Verify the signature
666        let verified = verification_key
667            .verify_signature(signing_input.as_bytes(), &signature_bytes, &protected)
668            .await
669            .map_err(|e| Error::Cryptography(e.to_string()))?;
670
671        if !verified {
672            return Err(Error::Cryptography(
673                "Signature verification failed".to_string(),
674            ));
675        }
676
677        // Decode the payload
678        let payload_bytes = base64::engine::general_purpose::STANDARD
679            .decode(&jws.payload)
680            .map_err(|e| Error::Cryptography(format!("Failed to decode payload: {}", e)))?;
681
682        Ok(payload_bytes)
683    }
684
685    /// Encrypt data for a recipient
686    async fn encrypt_jwe(
687        &self,
688        sender_kid: &str,
689        recipient_kid: &str,
690        plaintext: &[u8],
691        protected_header: Option<JweProtected>,
692    ) -> Result<String> {
693        // Get the encryption key
694        let encryption_key = KeyManager::get_encryption_key(self, sender_kid).await?;
695
696        // Resolve the recipient's verification key
697        let recipient_key = KeyManager::resolve_verification_key(self, recipient_kid).await?;
698
699        // Encrypt the plaintext
700        let jwe = encryption_key
701            .create_jwe(plaintext, &[recipient_key], protected_header)
702            .await
703            .map_err(|e| Error::Cryptography(e.to_string()))?;
704
705        // Serialize the JWE
706        serde_json::to_string(&jwe).map_err(|e| Error::Serialization(e.to_string()))
707    }
708
709    /// Decrypt a JWE
710    async fn decrypt_jwe(&self, jwe: &str, expected_kid: Option<&str>) -> Result<Vec<u8>> {
711        // Parse the JWE
712        let jwe: crate::message::Jwe = serde_json::from_str(jwe)
713            .map_err(|e| Error::Serialization(format!("Failed to parse JWE: {}", e)))?;
714
715        // Find the recipient if expected_kid is provided
716        if let Some(kid) = expected_kid {
717            // Just verify recipient exists, we don't need the actual instance
718            jwe.recipients
719                .iter()
720                .find(|r| r.header.kid == kid)
721                .ok_or_else(|| {
722                    Error::Cryptography(format!("No recipient found with kid: {}", kid))
723                })?;
724
725            // Get the decryption key
726            let decryption_key = KeyManager::get_decryption_key(self, kid).await?;
727
728            // Decrypt the JWE
729            decryption_key
730                .unwrap_jwe(&jwe)
731                .await
732                .map_err(|e| Error::Cryptography(e.to_string()))
733        } else {
734            // Try each recipient
735            for recipient in &jwe.recipients {
736                // Try to get the decryption key
737                if let Ok(decryption_key) =
738                    KeyManager::get_decryption_key(self, &recipient.header.kid).await
739                {
740                    // Try to decrypt
741                    if let Ok(plaintext) = decryption_key.unwrap_jwe(&jwe).await {
742                        return Ok(plaintext);
743                    }
744                }
745            }
746
747            Err(Error::Cryptography(
748                "Failed to decrypt JWE for any recipient".to_string(),
749            ))
750        }
751    }
752}
753
754/// A builder for AgentKeyManager
755#[derive(Debug, Clone)]
756pub struct AgentKeyManagerBuilder {
757    /// Legacy secrets
758    secrets: HashMap<String, Secret>,
759    /// Signing keys
760    signing_keys: HashMap<String, Arc<dyn SigningKey + Send + Sync>>,
761    /// Encryption keys
762    encryption_keys: HashMap<String, Arc<dyn EncryptionKey + Send + Sync>>,
763    /// Decryption keys
764    decryption_keys: HashMap<String, Arc<dyn DecryptionKey + Send + Sync>>,
765    /// Verification keys
766    verification_keys: HashMap<String, Arc<dyn VerificationKey + Send + Sync>>,
767    /// Load from storage
768    load_from_storage: bool,
769    /// Storage path
770    storage_path: Option<PathBuf>,
771}
772
773impl Default for AgentKeyManagerBuilder {
774    fn default() -> Self {
775        Self::new()
776    }
777}
778
779impl AgentKeyManagerBuilder {
780    /// Create a new KeyManagerBuilder
781    pub fn new() -> Self {
782        Self {
783            secrets: HashMap::new(),
784            signing_keys: HashMap::new(),
785            encryption_keys: HashMap::new(),
786            decryption_keys: HashMap::new(),
787            verification_keys: HashMap::new(),
788            load_from_storage: false,
789            storage_path: None,
790        }
791    }
792
793    /// Load keys from default storage location
794    pub fn load_from_default_storage(mut self) -> Self {
795        self.load_from_storage = true;
796        self.storage_path = None;
797        self
798    }
799
800    /// Load keys from a specific storage path
801    pub fn load_from_path(mut self, path: PathBuf) -> Self {
802        self.load_from_storage = true;
803        self.storage_path = Some(path);
804        self
805    }
806
807    /// Add a legacy secret
808    pub fn add_secret(mut self, did: String, secret: Secret) -> Self {
809        self.secrets.insert(did, secret);
810        self
811    }
812
813    /// Add a signing key
814    pub fn add_signing_key(mut self, key: Arc<dyn SigningKey + Send + Sync>) -> Self {
815        self.signing_keys.insert(key.key_id().to_string(), key);
816        self
817    }
818
819    /// Add an encryption key
820    pub fn add_encryption_key(mut self, key: Arc<dyn EncryptionKey + Send + Sync>) -> Self {
821        self.encryption_keys.insert(key.key_id().to_string(), key);
822        self
823    }
824
825    /// Add a decryption key
826    pub fn add_decryption_key(mut self, key: Arc<dyn DecryptionKey + Send + Sync>) -> Self {
827        self.decryption_keys.insert(key.key_id().to_string(), key);
828        self
829    }
830
831    /// Add a verification key
832    pub fn add_verification_key(mut self, key: Arc<dyn VerificationKey + Send + Sync>) -> Self {
833        self.verification_keys.insert(key.key_id().to_string(), key);
834        self
835    }
836
837    /// Build the KeyManager
838    pub fn build(self) -> Result<AgentKeyManager> {
839        let mut key_manager = AgentKeyManager {
840            generator: DIDKeyGenerator::new(),
841            secrets: Arc::new(RwLock::new(self.secrets)),
842            signing_keys: Arc::new(RwLock::new(self.signing_keys)),
843            encryption_keys: Arc::new(RwLock::new(self.encryption_keys)),
844            decryption_keys: Arc::new(RwLock::new(self.decryption_keys)),
845            verification_keys: Arc::new(RwLock::new(self.verification_keys)),
846            storage_path: self.storage_path.clone(),
847        };
848
849        // Load keys from storage if requested
850        if self.load_from_storage {
851            key_manager = if let Some(path) = self.storage_path {
852                key_manager.load_from_path(path)?
853            } else {
854                key_manager.load_from_default_storage()?
855            };
856        }
857
858        Ok(key_manager)
859    }
860}
861
862#[async_trait]
863impl KeyManagerPacking for AgentKeyManager {
864    async fn get_signing_key(&self, kid: &str) -> Result<Arc<dyn SigningKey + Send + Sync>> {
865        KeyManager::get_signing_key(self, kid)
866            .await
867            .map_err(|e| Error::from(MessageError::KeyManager(e.to_string())))
868    }
869
870    async fn get_encryption_key(&self, kid: &str) -> Result<Arc<dyn EncryptionKey + Send + Sync>> {
871        KeyManager::get_encryption_key(self, kid)
872            .await
873            .map_err(|e| Error::from(MessageError::KeyManager(e.to_string())))
874    }
875
876    async fn get_decryption_key(&self, kid: &str) -> Result<Arc<dyn DecryptionKey + Send + Sync>> {
877        KeyManager::get_decryption_key(self, kid)
878            .await
879            .map_err(|e| Error::from(MessageError::KeyManager(e.to_string())))
880    }
881
882    async fn resolve_verification_key(
883        &self,
884        kid: &str,
885    ) -> Result<Arc<dyn VerificationKey + Send + Sync>> {
886        KeyManager::resolve_verification_key(self, kid)
887            .await
888            .map_err(|e| Error::from(MessageError::KeyManager(e.to_string())))
889    }
890}
tap_agent/agent_key_manager.rs

tap_agent/
agent_key_manager.rs
