1use crate::keys::bls::{BlsKeypair, BlsSecretKey};
7use aes::cipher::{KeyIvInit, StreamCipher};
8use pbkdf2::pbkdf2_hmac;
9use scrypt::{scrypt, Params as ScryptParams};
10use serde::{Deserialize, Serialize};
11use sha2::{Digest, Sha256};
12use std::fs;
13use std::path::Path;
14use thiserror::Error;
15use uuid::Uuid;
16
17type Aes128Ctr = ctr::Ctr128BE<aes::Aes128>;
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct Keystore {
22 pub crypto: KeystoreCrypto,
23 #[serde(default)]
24 pub description: Option<String>,
25 pub pubkey: String,
26 pub path: String,
27 pub uuid: Uuid,
28 pub version: u32,
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct KeystoreCrypto {
34 pub kdf: KeystoreKdf,
35 pub checksum: KeystoreChecksum,
36 pub cipher: KeystoreCipher,
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct KeystoreKdf {
42 pub function: String,
43 pub params: KdfParams,
44 pub message: String,
45}
46
47#[derive(Debug, Clone, Serialize, Deserialize)]
49#[serde(untagged)]
50pub enum KdfParams {
51 Scrypt(ScryptKdfParams),
52 Pbkdf2(Pbkdf2KdfParams),
53}
54
55#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct ScryptKdfParams {
58 pub dklen: u32,
59 pub n: u32,
60 pub p: u32,
61 pub r: u32,
62 pub salt: String,
63}
64
65#[derive(Debug, Clone, Serialize, Deserialize)]
67pub struct Pbkdf2KdfParams {
68 pub dklen: u32,
69 pub c: u32,
70 pub prf: String,
71 pub salt: String,
72}
73
74#[derive(Debug, Clone, Serialize, Deserialize)]
76pub struct KeystoreChecksum {
77 pub function: String,
78 pub params: serde_json::Value,
79 pub message: String,
80}
81
82#[derive(Debug, Clone, Serialize, Deserialize)]
84pub struct KeystoreCipher {
85 pub function: String,
86 pub params: CipherParams,
87 pub message: String,
88}
89
90#[derive(Debug, Clone, Serialize, Deserialize)]
92pub struct CipherParams {
93 pub iv: String,
94}
95
96impl Keystore {
97 pub fn load(path: impl AsRef<Path>) -> Result<Self, KeystoreError> {
99 let contents = fs::read_to_string(path.as_ref())
100 .map_err(|e| KeystoreError::Io(e.to_string()))?;
101
102 serde_json::from_str(&contents)
103 .map_err(|e| KeystoreError::Parse(e.to_string()))
104 }
105
106 pub fn decrypt(&self, password: &str) -> Result<BlsSecretKey, KeystoreError> {
108 let decryption_key = self.derive_key(password)?;
110
111 self.verify_checksum(&decryption_key)?;
113
114 let secret_key_bytes = self.decrypt_secret(&decryption_key)?;
116
117 BlsSecretKey::from_bytes(&secret_key_bytes)
119 .map_err(|e| KeystoreError::InvalidKey(e.to_string()))
120 }
121
122 pub fn decrypt_keypair(&self, password: &str) -> Result<BlsKeypair, KeystoreError> {
124 let secret = self.decrypt(password)?;
125 Ok(BlsKeypair::from_secret(secret))
126 }
127
128 fn derive_key(&self, password: &str) -> Result<[u8; 32], KeystoreError> {
130 let kdf = &self.crypto.kdf;
131
132 match kdf.function.as_str() {
133 "scrypt" => {
134 let params = match &kdf.params {
135 KdfParams::Scrypt(p) => p,
136 _ => return Err(KeystoreError::InvalidKdf("Expected scrypt params".into())),
137 };
138
139 let salt = hex::decode(¶ms.salt)
140 .map_err(|e| KeystoreError::Parse(format!("Invalid salt hex: {}", e)))?;
141
142 let log_n = (params.n as f64).log2() as u8;
144 let scrypt_params = ScryptParams::new(log_n, params.r, params.p, params.dklen as usize)
145 .map_err(|e| KeystoreError::InvalidKdf(format!("Invalid scrypt params: {:?}", e)))?;
146
147 let mut dk = [0u8; 32];
148 scrypt(password.as_bytes(), &salt, &scrypt_params, &mut dk)
149 .map_err(|e| KeystoreError::InvalidKdf(format!("Scrypt failed: {:?}", e)))?;
150
151 Ok(dk)
152 }
153 "pbkdf2" => {
154 let params = match &kdf.params {
155 KdfParams::Pbkdf2(p) => p,
156 _ => return Err(KeystoreError::InvalidKdf("Expected pbkdf2 params".into())),
157 };
158
159 let salt = hex::decode(¶ms.salt)
160 .map_err(|e| KeystoreError::Parse(format!("Invalid salt hex: {}", e)))?;
161
162 let mut dk = [0u8; 32];
163 pbkdf2_hmac::<Sha256>(password.as_bytes(), &salt, params.c, &mut dk);
164
165 Ok(dk)
166 }
167 other => Err(KeystoreError::UnsupportedKdf(other.to_string())),
168 }
169 }
170
171 fn verify_checksum(&self, decryption_key: &[u8; 32]) -> Result<(), KeystoreError> {
173 let cipher_message = hex::decode(&self.crypto.cipher.message)
174 .map_err(|e| KeystoreError::Parse(format!("Invalid cipher message hex: {}", e)))?;
175
176 let mut hasher = Sha256::new();
178 hasher.update(&decryption_key[16..32]);
179 hasher.update(&cipher_message);
180 let computed_checksum = hasher.finalize();
181
182 let expected_checksum = hex::decode(&self.crypto.checksum.message)
183 .map_err(|e| KeystoreError::Parse(format!("Invalid checksum hex: {}", e)))?;
184
185 if computed_checksum.as_slice() != expected_checksum.as_slice() {
186 return Err(KeystoreError::InvalidPassword);
187 }
188
189 Ok(())
190 }
191
192 fn decrypt_secret(&self, decryption_key: &[u8; 32]) -> Result<[u8; 32], KeystoreError> {
194 let cipher = &self.crypto.cipher;
195
196 if cipher.function != "aes-128-ctr" {
197 return Err(KeystoreError::UnsupportedCipher(cipher.function.clone()));
198 }
199
200 let iv = hex::decode(&cipher.params.iv)
201 .map_err(|e| KeystoreError::Parse(format!("Invalid IV hex: {}", e)))?;
202
203 let ciphertext = hex::decode(&cipher.message)
204 .map_err(|e| KeystoreError::Parse(format!("Invalid ciphertext hex: {}", e)))?;
205
206 let aes_key: [u8; 16] = decryption_key[0..16].try_into().unwrap();
208 let iv_array: [u8; 16] = iv.try_into()
209 .map_err(|_| KeystoreError::Parse("Invalid IV length".into()))?;
210
211 let mut cipher = Aes128Ctr::new(&aes_key.into(), &iv_array.into());
212 let mut plaintext = ciphertext;
213 cipher.apply_keystream(&mut plaintext);
214
215 if plaintext.len() != 32 {
216 return Err(KeystoreError::InvalidKey(format!(
217 "Expected 32 bytes, got {}",
218 plaintext.len()
219 )));
220 }
221
222 let mut secret = [0u8; 32];
223 secret.copy_from_slice(&plaintext);
224 Ok(secret)
225 }
226
227 pub fn public_key_bytes(&self) -> Result<[u8; 48], KeystoreError> {
229 let pubkey = self.pubkey.strip_prefix("0x").unwrap_or(&self.pubkey);
230 let bytes = hex::decode(pubkey)
231 .map_err(|e| KeystoreError::Parse(format!("Invalid pubkey hex: {}", e)))?;
232
233 if bytes.len() != 48 {
234 return Err(KeystoreError::Parse(format!(
235 "Expected 48 bytes for pubkey, got {}",
236 bytes.len()
237 )));
238 }
239
240 let mut arr = [0u8; 48];
241 arr.copy_from_slice(&bytes);
242 Ok(arr)
243 }
244}
245
246pub fn load_keystores_from_dir(
248 dir: impl AsRef<Path>,
249 password: &str,
250) -> Result<Vec<BlsKeypair>, KeystoreError> {
251 let dir = dir.as_ref();
252
253 if !dir.exists() {
254 return Err(KeystoreError::Io(format!(
255 "Directory does not exist: {}",
256 dir.display()
257 )));
258 }
259
260 let mut keypairs = Vec::new();
261
262 for entry in fs::read_dir(dir).map_err(|e| KeystoreError::Io(e.to_string()))? {
263 let entry = entry.map_err(|e| KeystoreError::Io(e.to_string()))?;
264 let path = entry.path();
265
266 if path.extension().and_then(|s| s.to_str()) == Some("json") {
267 match Keystore::load(&path) {
268 Ok(keystore) => {
269 match keystore.decrypt_keypair(password) {
270 Ok(keypair) => {
271 tracing::info!(
272 pubkey = %keystore.pubkey,
273 path = %path.display(),
274 "Loaded validator key"
275 );
276 keypairs.push(keypair);
277 }
278 Err(e) => {
279 tracing::warn!(
280 path = %path.display(),
281 error = %e,
282 "Failed to decrypt keystore"
283 );
284 }
285 }
286 }
287 Err(e) => {
288 tracing::debug!(
289 path = %path.display(),
290 error = %e,
291 "Failed to load keystore"
292 );
293 }
294 }
295 }
296 }
297
298 Ok(keypairs)
299}
300
301#[derive(Debug, Error)]
303pub enum KeystoreError {
304 #[error("I/O error: {0}")]
305 Io(String),
306
307 #[error("Parse error: {0}")]
308 Parse(String),
309
310 #[error("Invalid password")]
311 InvalidPassword,
312
313 #[error("Unsupported KDF: {0}")]
314 UnsupportedKdf(String),
315
316 #[error("Invalid KDF parameters: {0}")]
317 InvalidKdf(String),
318
319 #[error("Unsupported cipher: {0}")]
320 UnsupportedCipher(String),
321
322 #[error("Invalid key: {0}")]
323 InvalidKey(String),
324}
325
326#[cfg(test)]
327mod tests {
328 use super::*;
329 use std::io::Write;
330 use tempfile::TempDir;
331
332 const TEST_KEYSTORE_SCRYPT: &str = r#"{
334 "crypto": {
335 "kdf": {
336 "function": "scrypt",
337 "params": {
338 "dklen": 32,
339 "n": 262144,
340 "p": 1,
341 "r": 8,
342 "salt": "d4e56740f876aef8c010b86a40d5f56745a118d0906a34e69aec8c0db1cb8fa3"
343 },
344 "message": ""
345 },
346 "checksum": {
347 "function": "sha256",
348 "params": {},
349 "message": "d2217fe5f3e9a1e34581ef8a78f7c9928e436d36dacc5e846690a5581e8ea484"
350 },
351 "cipher": {
352 "function": "aes-128-ctr",
353 "params": {
354 "iv": "264daa3f303d7259501c93d997d84fe6"
355 },
356 "message": "06ae90d55fe0a6e9c5c3bc5b170827b2e5cce3929ed3f116c2811e6366dfe20f"
357 }
358 },
359 "description": "Test keystore",
360 "pubkey": "0x9612d7a727c9d0a22e185a1c768478dfe919cada9266988cb32359c11f2b7b27f4ae4040902382ae2910c15e2b420d07",
361 "path": "m/12381/3600/0/0/0",
362 "uuid": "1d85ae20-35c5-4611-8e4c-00a8c6c5fb55",
363 "version": 4
364 }"#;
365
366 #[test]
367 fn test_parse_keystore() {
368 let keystore: Keystore = serde_json::from_str(TEST_KEYSTORE_SCRYPT).unwrap();
369 assert_eq!(keystore.version, 4);
370 assert!(keystore.pubkey.starts_with("0x"));
371 }
372
373 #[test]
374 fn test_load_keystore_from_file() {
375 let dir = TempDir::new().unwrap();
376 let path = dir.path().join("keystore.json");
377
378 let mut file = fs::File::create(&path).unwrap();
379 file.write_all(TEST_KEYSTORE_SCRYPT.as_bytes()).unwrap();
380
381 let keystore = Keystore::load(&path).unwrap();
382 assert_eq!(keystore.version, 4);
383 }
384
385 #[test]
386 fn test_public_key_bytes() {
387 let keystore: Keystore = serde_json::from_str(TEST_KEYSTORE_SCRYPT).unwrap();
388 let pubkey = keystore.public_key_bytes().unwrap();
389 assert_eq!(pubkey.len(), 48);
390 }
391}