tap_agent/
storage.rs

1//! Key storage functionality for TAP Agent
2//!
3//! This module provides utilities for persisting agent keys to disk
4//! and loading them later. This allows for persistent agent identities
5//! across multiple runs.
6
7use crate::did::{GeneratedKey, KeyType};
8use crate::error::{Error, Result};
9use crate::key_manager::{Secret, SecretMaterial, SecretType};
10use base64::Engine;
11use dirs::home_dir;
12use serde::{Deserialize, Serialize};
13use std::collections::HashMap;
14use std::fs;
15use std::path::{Path, PathBuf};
16
17/// Default directory for TAP configuration and keys
18pub const DEFAULT_TAP_DIR: &str = ".tap";
19/// Default filename for the keys file
20pub const DEFAULT_KEYS_FILE: &str = "keys.json";
21
22/// A structure representing a stored key
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct StoredKey {
25    /// The DID for this key
26    pub did: String,
27    /// Human-friendly label for this key
28    #[serde(default)]
29    pub label: String,
30    /// The key type (e.g., Ed25519, P256)
31    #[serde(with = "key_type_serde")]
32    pub key_type: KeyType,
33    /// Base64-encoded private key
34    pub private_key: String,
35    /// Base64-encoded public key
36    pub public_key: String,
37    /// Optional metadata for this key
38    #[serde(default)]
39    pub metadata: HashMap<String, String>,
40}
41
42/// Serialization helper for KeyType
43mod key_type_serde {
44    use super::KeyType;
45    use serde::{Deserialize, Deserializer, Serializer};
46
47    pub fn serialize<S>(key_type: &KeyType, serializer: S) -> Result<S::Ok, S::Error>
48    where
49        S: Serializer,
50    {
51        let s = match key_type {
52            KeyType::Ed25519 => "Ed25519",
53            KeyType::P256 => "P256",
54            KeyType::Secp256k1 => "Secp256k1",
55        };
56        serializer.serialize_str(s)
57    }
58
59    pub fn deserialize<'de, D>(deserializer: D) -> Result<KeyType, D::Error>
60    where
61        D: Deserializer<'de>,
62    {
63        let s = String::deserialize(deserializer)?;
64        match s.as_str() {
65            "Ed25519" => Ok(KeyType::Ed25519),
66            "P256" => Ok(KeyType::P256),
67            "Secp256k1" => Ok(KeyType::Secp256k1),
68            _ => Err(serde::de::Error::custom(format!("Unknown key type: {}", s))),
69        }
70    }
71}
72
73/// A collection of stored keys
74#[derive(Debug, Clone, Serialize, Deserialize, Default)]
75pub struct KeyStorage {
76    /// A map of DIDs to their stored keys
77    pub keys: HashMap<String, StoredKey>,
78    /// The default DID to use when not specified
79    pub default_did: Option<String>,
80    /// Creation timestamp
81    #[serde(default = "chrono::Utc::now")]
82    pub created_at: chrono::DateTime<chrono::Utc>,
83    /// Last update timestamp
84    #[serde(default = "chrono::Utc::now")]
85    pub updated_at: chrono::DateTime<chrono::Utc>,
86    /// Base directory for agent subdirectories (not serialized, runtime only)
87    #[serde(skip)]
88    base_directory: Option<PathBuf>,
89}
90
91impl KeyStorage {
92    /// Create a new empty key storage
93    pub fn new() -> Self {
94        Default::default()
95    }
96
97    /// Add a key to the storage
98    pub fn add_key(&mut self, mut key: StoredKey) {
99        // Generate default label if not provided
100        if key.label.is_empty() {
101            key.label = self.generate_default_label();
102        }
103
104        // Ensure label is unique
105        let final_label = self.ensure_unique_label(&key.label, Some(&key.did));
106        key.label = final_label.clone();
107
108        // If this is the first key, make it the default
109        if self.keys.is_empty() {
110            self.default_did = Some(key.did.clone());
111        }
112
113        self.keys.insert(key.did.clone(), key);
114        self.updated_at = chrono::Utc::now();
115    }
116
117    /// Generate a default label in the format agent-{n}
118    fn generate_default_label(&self) -> String {
119        let mut counter = 1;
120        loop {
121            let label = format!("agent-{}", counter);
122            if !self.keys.values().any(|key| key.label == label) {
123                return label;
124            }
125            counter += 1;
126        }
127    }
128
129    /// Ensure a label is unique, modifying it if necessary
130    fn ensure_unique_label(&self, desired_label: &str, exclude_did: Option<&str>) -> String {
131        // Check if the label exists in any key
132        if let Some(existing_key) = self.keys.values().find(|key| key.label == desired_label) {
133            // If it belongs to the same DID we're updating, it's fine
134            if exclude_did.is_some() && existing_key.did == exclude_did.unwrap() {
135                return desired_label.to_string();
136            }
137        } else {
138            // Label doesn't exist, so it's available
139            return desired_label.to_string();
140        }
141
142        // Generate a unique label by appending a number
143        let mut counter = 2;
144        loop {
145            let new_label = format!("{}-{}", desired_label, counter);
146            if !self.keys.values().any(|key| key.label == new_label) {
147                return new_label;
148            }
149            counter += 1;
150        }
151    }
152
153    /// Find a key by label
154    pub fn find_by_label(&self, label: &str) -> Option<&StoredKey> {
155        self.keys.values().find(|key| key.label == label)
156    }
157
158    /// Update a key's label
159    pub fn update_label(&mut self, did: &str, new_label: &str) -> Result<()> {
160        // First ensure the key exists
161        if !self.keys.contains_key(did) {
162            return Err(Error::Storage(format!("Key with DID '{}' not found", did)));
163        }
164
165        // Ensure new label is unique
166        let final_label = self.ensure_unique_label(new_label, Some(did));
167
168        // Update the key's label
169        if let Some(key) = self.keys.get_mut(did) {
170            key.label = final_label;
171        }
172
173        self.updated_at = chrono::Utc::now();
174        Ok(())
175    }
176
177    /// Get the default key path
178    pub fn default_key_path() -> Option<PathBuf> {
179        home_dir().map(|home| home.join(DEFAULT_TAP_DIR).join(DEFAULT_KEYS_FILE))
180    }
181
182    /// Load keys from the default location
183    pub fn load_default() -> Result<Self> {
184        let path = Self::default_key_path().ok_or_else(|| {
185            Error::Storage("Could not determine home directory for default key path".to_string())
186        })?;
187        Self::load_from_path(&path)
188    }
189
190    /// Load keys from a specific path
191    pub fn load_from_path(path: &Path) -> Result<Self> {
192        let mut storage = if !path.exists() {
193            Self::new()
194        } else {
195            let contents = fs::read_to_string(path)
196                .map_err(|e| Error::Storage(format!("Failed to read key storage file: {}", e)))?;
197
198            let mut storage: KeyStorage = serde_json::from_str(&contents)
199                .map_err(|e| Error::Storage(format!("Failed to parse key storage file: {}", e)))?;
200
201            // Ensure all keys have labels (for backward compatibility)
202            storage.ensure_all_keys_have_labels();
203
204            storage
205        };
206
207        // Set base directory from the path for agent subdirectories
208        if let Some(parent) = path.parent() {
209            storage.base_directory = Some(parent.to_path_buf());
210        }
211
212        Ok(storage)
213    }
214
215    /// Ensure all keys have labels (for backward compatibility)
216    fn ensure_all_keys_have_labels(&mut self) {
217        let mut keys_to_update = Vec::new();
218
219        for (did, key) in &self.keys {
220            if key.label.is_empty() {
221                keys_to_update.push(did.clone());
222            }
223        }
224
225        for did in keys_to_update {
226            let new_label = self.generate_default_label();
227            if let Some(key) = self.keys.get_mut(&did) {
228                key.label = new_label;
229            }
230        }
231    }
232
233    /// Save keys to the default location
234    pub fn save_default(&self) -> Result<()> {
235        let path = Self::default_key_path().ok_or_else(|| {
236            Error::Storage("Could not determine home directory for default key path".to_string())
237        })?;
238
239        // Ensure directory exists
240        if let Some(parent) = path.parent() {
241            fs::create_dir_all(parent).map_err(|e| {
242                Error::Storage(format!("Failed to create key storage directory: {}", e))
243            })?;
244        }
245
246        self.save_to_path(&path)
247    }
248
249    /// Save keys to a specific path
250    pub fn save_to_path(&self, path: &Path) -> Result<()> {
251        let contents = serde_json::to_string_pretty(self)
252            .map_err(|e| Error::Storage(format!("Failed to serialize key storage: {}", e)))?;
253
254        fs::write(path, contents)
255            .map_err(|e| Error::Storage(format!("Failed to write key storage file: {}", e)))?;
256
257        Ok(())
258    }
259
260    /// Convert a GeneratedKey to a StoredKey
261    pub fn from_generated_key(key: &GeneratedKey) -> StoredKey {
262        StoredKey {
263            did: key.did.clone(),
264            label: String::new(), // Will be set when added to storage
265            key_type: key.key_type,
266            private_key: base64::engine::general_purpose::STANDARD.encode(&key.private_key),
267            public_key: base64::engine::general_purpose::STANDARD.encode(&key.public_key),
268            metadata: HashMap::new(),
269        }
270    }
271
272    /// Convert a GeneratedKey to a StoredKey with a specific label
273    pub fn from_generated_key_with_label(key: &GeneratedKey, label: &str) -> StoredKey {
274        StoredKey {
275            did: key.did.clone(),
276            label: label.to_string(),
277            key_type: key.key_type,
278            private_key: base64::engine::general_purpose::STANDARD.encode(&key.private_key),
279            public_key: base64::engine::general_purpose::STANDARD.encode(&key.public_key),
280            metadata: HashMap::new(),
281        }
282    }
283
284    /// Convert a StoredKey to a Secret
285    pub fn to_secret(key: &StoredKey) -> Secret {
286        Secret {
287            id: key.did.clone(),
288            type_: SecretType::JsonWebKey2020,
289            secret_material: SecretMaterial::JWK {
290                private_key_jwk: generate_jwk_for_key(key),
291            },
292        }
293    }
294    /// Create agent directory and save policies/metadata files
295    pub fn create_agent_directory(
296        &self,
297        did: &str,
298        policies: &[String],
299        metadata: &HashMap<String, String>,
300    ) -> Result<()> {
301        let sanitized_did = sanitize_did(did);
302        let agent_dir = self.get_agent_directory(&sanitized_did)?;
303
304        // Create the agent directory
305        fs::create_dir_all(&agent_dir).map_err(|e| {
306            Error::Storage(format!(
307                "Failed to create agent directory {}: {}",
308                agent_dir.display(),
309                e
310            ))
311        })?;
312
313        // Save policies.json
314        let policies_file = agent_dir.join("policies.json");
315        let policies_json = serde_json::to_string_pretty(policies)
316            .map_err(|e| Error::Storage(format!("Failed to serialize policies: {}", e)))?;
317        fs::write(&policies_file, policies_json)
318            .map_err(|e| Error::Storage(format!("Failed to write policies file: {}", e)))?;
319
320        // Save metadata.json
321        let metadata_file = agent_dir.join("metadata.json");
322        let metadata_json = serde_json::to_string_pretty(metadata)
323            .map_err(|e| Error::Storage(format!("Failed to serialize metadata: {}", e)))?;
324        fs::write(&metadata_file, metadata_json)
325            .map_err(|e| Error::Storage(format!("Failed to write metadata file: {}", e)))?;
326
327        Ok(())
328    }
329
330    /// Load policies from agent directory
331    pub fn load_agent_policies(&self, did: &str) -> Result<Vec<String>> {
332        let sanitized_did = sanitize_did(did);
333        let agent_dir = self.get_agent_directory(&sanitized_did)?;
334        let policies_file = agent_dir.join("policies.json");
335
336        if !policies_file.exists() {
337            return Ok(vec![]);
338        }
339
340        let content = fs::read_to_string(&policies_file)
341            .map_err(|e| Error::Storage(format!("Failed to read policies file: {}", e)))?;
342
343        serde_json::from_str(&content)
344            .map_err(|e| Error::Storage(format!("Failed to parse policies file: {}", e)))
345    }
346
347    /// Load metadata from agent directory
348    pub fn load_agent_metadata(&self, did: &str) -> Result<HashMap<String, String>> {
349        let sanitized_did = sanitize_did(did);
350        let agent_dir = self.get_agent_directory(&sanitized_did)?;
351        let metadata_file = agent_dir.join("metadata.json");
352
353        if !metadata_file.exists() {
354            return Ok(HashMap::new());
355        }
356
357        let content = fs::read_to_string(&metadata_file)
358            .map_err(|e| Error::Storage(format!("Failed to read metadata file: {}", e)))?;
359
360        serde_json::from_str(&content)
361            .map_err(|e| Error::Storage(format!("Failed to parse metadata file: {}", e)))
362    }
363
364    /// Get the agent directory path for a given DID
365    fn get_agent_directory(&self, sanitized_did: &str) -> Result<PathBuf> {
366        let base_dir = if let Some(ref base) = self.base_directory {
367            base.clone()
368        } else {
369            let home = home_dir()
370                .ok_or_else(|| Error::Storage("Could not determine home directory".to_string()))?;
371            home.join(DEFAULT_TAP_DIR)
372        };
373        Ok(base_dir.join(sanitized_did))
374    }
375}
376
377/// Sanitize a DID for use as a directory name (same as TAP Node)
378fn sanitize_did(did: &str) -> String {
379    did.replace(':', "_")
380}
381
382/// Generate a JWK for a stored key
383fn generate_jwk_for_key(key: &StoredKey) -> serde_json::Value {
384    match key.key_type {
385        KeyType::Ed25519 => {
386            serde_json::json!({
387                "kty": "OKP",
388                "crv": "Ed25519",
389                "x": key.public_key,
390                "d": key.private_key,
391                "kid": format!("{}#keys-1", key.did)
392            })
393        }
394        KeyType::P256 => {
395            serde_json::json!({
396                "kty": "EC",
397                "crv": "P-256",
398                "x": key.public_key,
399                "d": key.private_key,
400                "kid": format!("{}#keys-1", key.did)
401            })
402        }
403        KeyType::Secp256k1 => {
404            serde_json::json!({
405                "kty": "EC",
406                "crv": "secp256k1",
407                "x": key.public_key,
408                "d": key.private_key,
409                "kid": format!("{}#keys-1", key.did)
410            })
411        }
412    }
413}
tap_agent/storage.rs

tap_agent/
storage.rs
