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    /// The key type (e.g., Ed25519, P256)
28    #[serde(with = "key_type_serde")]
29    pub key_type: KeyType,
30    /// Base64-encoded private key
31    pub private_key: String,
32    /// Base64-encoded public key
33    pub public_key: String,
34    /// Optional metadata for this key
35    #[serde(default)]
36    pub metadata: HashMap<String, String>,
37}
38
39/// Serialization helper for KeyType
40mod key_type_serde {
41    use super::KeyType;
42    use serde::{Deserialize, Deserializer, Serializer};
43
44    pub fn serialize<S>(key_type: &KeyType, serializer: S) -> Result<S::Ok, S::Error>
45    where
46        S: Serializer,
47    {
48        let s = match key_type {
49            KeyType::Ed25519 => "Ed25519",
50            KeyType::P256 => "P256",
51            KeyType::Secp256k1 => "Secp256k1",
52        };
53        serializer.serialize_str(s)
54    }
55
56    pub fn deserialize<'de, D>(deserializer: D) -> Result<KeyType, D::Error>
57    where
58        D: Deserializer<'de>,
59    {
60        let s = String::deserialize(deserializer)?;
61        match s.as_str() {
62            "Ed25519" => Ok(KeyType::Ed25519),
63            "P256" => Ok(KeyType::P256),
64            "Secp256k1" => Ok(KeyType::Secp256k1),
65            _ => Err(serde::de::Error::custom(format!("Unknown key type: {}", s))),
66        }
67    }
68}
69
70/// A collection of stored keys
71#[derive(Debug, Clone, Serialize, Deserialize, Default)]
72pub struct KeyStorage {
73    /// A map of DIDs to their stored keys
74    pub keys: HashMap<String, StoredKey>,
75    /// The default DID to use when not specified
76    pub default_did: Option<String>,
77    /// Creation timestamp
78    #[serde(default = "chrono::Utc::now")]
79    pub created_at: chrono::DateTime<chrono::Utc>,
80    /// Last update timestamp
81    #[serde(default = "chrono::Utc::now")]
82    pub updated_at: chrono::DateTime<chrono::Utc>,
83}
84
85impl KeyStorage {
86    /// Create a new empty key storage
87    pub fn new() -> Self {
88        Default::default()
89    }
90
91    /// Add a key to the storage
92    pub fn add_key(&mut self, key: StoredKey) {
93        // If this is the first key, make it the default
94        if self.keys.is_empty() {
95            self.default_did = Some(key.did.clone());
96        }
97
98        self.keys.insert(key.did.clone(), key);
99        self.updated_at = chrono::Utc::now();
100    }
101
102    /// Get the default key path
103    pub fn default_key_path() -> Option<PathBuf> {
104        home_dir().map(|home| home.join(DEFAULT_TAP_DIR).join(DEFAULT_KEYS_FILE))
105    }
106
107    /// Load keys from the default location
108    pub fn load_default() -> Result<Self> {
109        let path = Self::default_key_path().ok_or_else(|| {
110            Error::Storage("Could not determine home directory for default key path".to_string())
111        })?;
112        Self::load_from_path(&path)
113    }
114
115    /// Load keys from a specific path
116    pub fn load_from_path(path: &Path) -> Result<Self> {
117        if !path.exists() {
118            return Ok(Self::new());
119        }
120
121        let contents = fs::read_to_string(path)
122            .map_err(|e| Error::Storage(format!("Failed to read key storage file: {}", e)))?;
123
124        let storage: KeyStorage = serde_json::from_str(&contents)
125            .map_err(|e| Error::Storage(format!("Failed to parse key storage file: {}", e)))?;
126
127        Ok(storage)
128    }
129
130    /// Save keys to the default location
131    pub fn save_default(&self) -> Result<()> {
132        let path = Self::default_key_path().ok_or_else(|| {
133            Error::Storage("Could not determine home directory for default key path".to_string())
134        })?;
135
136        // Ensure directory exists
137        if let Some(parent) = path.parent() {
138            fs::create_dir_all(parent).map_err(|e| {
139                Error::Storage(format!("Failed to create key storage directory: {}", e))
140            })?;
141        }
142
143        self.save_to_path(&path)
144    }
145
146    /// Save keys to a specific path
147    pub fn save_to_path(&self, path: &Path) -> Result<()> {
148        let contents = serde_json::to_string_pretty(self)
149            .map_err(|e| Error::Storage(format!("Failed to serialize key storage: {}", e)))?;
150
151        fs::write(path, contents)
152            .map_err(|e| Error::Storage(format!("Failed to write key storage file: {}", e)))?;
153
154        Ok(())
155    }
156
157    /// Convert a GeneratedKey to a StoredKey
158    pub fn from_generated_key(key: &GeneratedKey) -> StoredKey {
159        StoredKey {
160            did: key.did.clone(),
161            key_type: key.key_type,
162            private_key: base64::engine::general_purpose::STANDARD.encode(&key.private_key),
163            public_key: base64::engine::general_purpose::STANDARD.encode(&key.public_key),
164            metadata: HashMap::new(),
165        }
166    }
167
168    /// Convert a StoredKey to a Secret
169    pub fn to_secret(key: &StoredKey) -> Secret {
170        Secret {
171            id: key.did.clone(),
172            type_: SecretType::JsonWebKey2020,
173            secret_material: SecretMaterial::JWK {
174                private_key_jwk: generate_jwk_for_key(key),
175            },
176        }
177    }
178}
179
180/// Generate a JWK for a stored key
181fn generate_jwk_for_key(key: &StoredKey) -> serde_json::Value {
182    match key.key_type {
183        KeyType::Ed25519 => {
184            serde_json::json!({
185                "kty": "OKP",
186                "crv": "Ed25519",
187                "x": key.public_key,
188                "d": key.private_key,
189                "kid": format!("{}#keys-1", key.did)
190            })
191        }
192        KeyType::P256 => {
193            serde_json::json!({
194                "kty": "EC",
195                "crv": "P-256",
196                "x": key.public_key,
197                "d": key.private_key,
198                "kid": format!("{}#keys-1", key.did)
199            })
200        }
201        KeyType::Secp256k1 => {
202            serde_json::json!({
203                "kty": "EC",
204                "crv": "secp256k1",
205                "x": key.public_key,
206                "d": key.private_key,
207                "kid": format!("{}#keys-1", key.did)
208            })
209        }
210    }
211}