1use std::{
2 collections::HashMap,
3 fs,
4 io::{self, Read, Write},
5 path::{Path, PathBuf},
6 sync::{Arc, RwLock},
7};
8
9use rand::RngCore;
10use serde::{Deserialize, Serialize};
11use sha2::{Digest as Sha2Digest, Sha256};
12
13use crate::attestation::{Ed25519Signer, Signer};
14
15pub type KeyId = String;
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct KeyInfo {
22 pub id: KeyId,
23 pub algorithm: String, pub is_default: bool,
25 pub created_at: String, pub fingerprint: String,
28 pub public_key: Vec<u8>, }
30
31#[derive(Debug)]
33pub enum KeyError {
34 Io(io::Error),
35 Json(serde_json::Error),
36 Crypto(String),
37 NotFound(KeyId),
38 EmptyKeyId,
39 NoDefaultKey,
40}
41
42impl std::fmt::Display for KeyError {
43 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44 match self {
45 Self::Io(e) => write!(f, "keys io: {}", e),
46 Self::Json(e) => write!(f, "keys json: {}", e),
47 Self::Crypto(e) => write!(f, "keys crypto: {}", e),
48 Self::NotFound(k) => write!(f, "key not found: {}", k),
49 Self::EmptyKeyId => write!(f, "key id must not be empty"),
50 Self::NoDefaultKey => write!(f, "no default key — run treeship init"),
51 }
52 }
53}
54
55impl std::error::Error for KeyError {}
56impl From<io::Error> for KeyError { fn from(e: io::Error) -> Self { Self::Io(e) } }
57impl From<serde_json::Error> for KeyError { fn from(e: serde_json::Error) -> Self { Self::Json(e) } }
58
59#[derive(Serialize, Deserialize)]
63struct EncryptedEntry {
64 id: KeyId,
65 algorithm: String,
66 created_at: String,
67 public_key: Vec<u8>,
68 enc_priv_key: Vec<u8>,
70 nonce: Vec<u8>,
72}
73
74#[derive(Serialize, Deserialize, Default)]
76struct Manifest {
77 default_key_id: Option<KeyId>,
78 key_ids: Vec<KeyId>,
79}
80
81pub struct Store {
91 dir: PathBuf,
92 machine_key: [u8; 32],
93 cache: Arc<RwLock<HashMap<KeyId, EncryptedEntry>>>,
95}
96
97impl Store {
98 pub fn open(dir: impl AsRef<Path>) -> Result<Self, KeyError> {
100 let dir = dir.as_ref().to_path_buf();
101 fs::create_dir_all(&dir)?;
102
103 let machine_key = derive_machine_key(&dir)?;
104
105 Ok(Self {
106 dir,
107 machine_key,
108 cache: Arc::new(RwLock::new(HashMap::new())),
109 })
110 }
111
112 pub fn generate(&self, set_default: bool) -> Result<KeyInfo, KeyError> {
116 let key_id = new_key_id();
117
118 let signer = Ed25519Signer::generate(&key_id)
119 .map_err(|e| KeyError::Crypto(e.to_string()))?;
120
121 let secret = signer.secret_bytes();
122 let pub_key = signer.public_key_bytes();
123
124 let (enc, nonce) = aes_gcm_encrypt(&self.machine_key, &secret)
125 .map_err(|e| KeyError::Crypto(e))?;
126
127 let entry = EncryptedEntry {
128 id: key_id.clone(),
129 algorithm: "ed25519".into(),
130 created_at: crate::statements::unix_to_rfc3339(unix_now()),
131 public_key: pub_key.clone(),
132 enc_priv_key: enc,
133 nonce,
134 };
135
136 self.write_entry(&entry)?;
137
138 let mut manifest = self.read_manifest()?;
140 manifest.key_ids.push(key_id.clone());
141 if set_default || manifest.default_key_id.is_none() {
142 manifest.default_key_id = Some(key_id.clone());
143 }
144 self.write_manifest(&manifest)?;
145
146 self.cache.write().unwrap().insert(key_id.clone(), entry);
148
149 Ok(KeyInfo {
150 id: key_id,
151 algorithm: "ed25519".into(),
152 is_default: manifest.default_key_id.as_deref() == Some(&manifest.key_ids.last().unwrap_or(&String::new())),
153 created_at: crate::statements::unix_to_rfc3339(unix_now()),
154 fingerprint: fingerprint(&pub_key),
155 public_key: pub_key,
156 })
157 }
158
159 pub fn default_signer(&self) -> Result<Box<dyn Signer>, KeyError> {
161 let manifest = self.read_manifest()?;
162 let id = manifest.default_key_id.ok_or(KeyError::NoDefaultKey)?;
163 self.signer(&id)
164 }
165
166 pub fn signer(&self, id: &str) -> Result<Box<dyn Signer>, KeyError> {
168 let entry = self.load_entry(id)?;
169
170 let secret = aes_gcm_decrypt(&self.machine_key, &entry.enc_priv_key, &entry.nonce)
171 .map_err(|e| KeyError::Crypto(e))?;
172
173 let secret_arr: [u8; 32] = secret.try_into()
174 .map_err(|_| KeyError::Crypto("decrypted key is wrong length".into()))?;
175
176 let signer = Ed25519Signer::from_bytes(&entry.id, &secret_arr)
177 .map_err(|e| KeyError::Crypto(e.to_string()))?;
178
179 Ok(Box::new(signer))
180 }
181
182 pub fn default_key_id(&self) -> Result<KeyId, KeyError> {
184 self.read_manifest()?
185 .default_key_id
186 .ok_or(KeyError::NoDefaultKey)
187 }
188
189 pub fn list(&self) -> Result<Vec<KeyInfo>, KeyError> {
191 let manifest = self.read_manifest()?;
192 let default = manifest.default_key_id.as_deref().unwrap_or("");
193
194 manifest.key_ids.iter().map(|id| {
195 let entry = self.load_entry(id)?;
196 Ok(KeyInfo {
197 id: entry.id.clone(),
198 algorithm: entry.algorithm.clone(),
199 is_default: entry.id == default,
200 created_at: entry.created_at.clone(),
201 fingerprint: fingerprint(&entry.public_key),
202 public_key: entry.public_key.clone(),
203 })
204 }).collect()
205 }
206
207 pub fn set_default(&self, id: &str) -> Result<(), KeyError> {
209 self.load_entry(id)?;
211 let mut manifest = self.read_manifest()?;
212 manifest.default_key_id = Some(id.to_string());
213 self.write_manifest(&manifest)
214 }
215
216 pub fn public_key(&self, id: &str) -> Result<Vec<u8>, KeyError> {
218 Ok(self.load_entry(id)?.public_key)
219 }
220
221 fn load_entry(&self, id: &str) -> Result<EncryptedEntry, KeyError> {
224 if let Ok(cache) = self.cache.read() {
226 if let Some(entry) = cache.get(id) {
227 return Ok(EncryptedEntry {
229 id: entry.id.clone(),
230 algorithm: entry.algorithm.clone(),
231 created_at: entry.created_at.clone(),
232 public_key: entry.public_key.clone(),
233 enc_priv_key: entry.enc_priv_key.clone(),
234 nonce: entry.nonce.clone(),
235 });
236 }
237 }
238 self.read_entry(id)
239 }
240
241 fn entry_path(&self, id: &str) -> PathBuf {
242 self.dir.join(format!("{}.json", id))
243 }
244
245 fn write_entry(&self, entry: &EncryptedEntry) -> Result<(), KeyError> {
246 let path = self.entry_path(&entry.id);
247 let json = serde_json::to_vec_pretty(entry)?;
248 write_file_600(&path, &json)?;
249 Ok(())
250 }
251
252 fn read_entry(&self, id: &str) -> Result<EncryptedEntry, KeyError> {
253 let path = self.entry_path(id);
254 if !path.exists() {
255 return Err(KeyError::NotFound(id.to_string()));
256 }
257 let bytes = fs::read(&path)?;
258 let entry: EncryptedEntry = serde_json::from_slice(&bytes)?;
259 Ok(entry)
260 }
261
262 fn manifest_path(&self) -> PathBuf {
263 self.dir.join("manifest.json")
264 }
265
266 fn read_manifest(&self) -> Result<Manifest, KeyError> {
267 let path = self.manifest_path();
268 if !path.exists() {
269 return Ok(Manifest::default());
270 }
271 let bytes = fs::read(&path)?;
272 Ok(serde_json::from_slice(&bytes)?)
273 }
274
275 fn write_manifest(&self, m: &Manifest) -> Result<(), KeyError> {
276 let json = serde_json::to_vec_pretty(m)?;
277 write_file_600(&self.manifest_path(), &json)?;
278 Ok(())
279 }
280}
281
282pub fn aes_gcm_encrypt(key: &[u8; 32], plaintext: &[u8]) -> Result<(Vec<u8>, Vec<u8>), String> {
287 use sha2::Sha256;
298
299 let mut nonce = [0u8; 12];
300 rand::thread_rng().fill_bytes(&mut nonce);
301
302 let mut enc_key_input = key.to_vec();
304 enc_key_input.extend_from_slice(&nonce);
305 enc_key_input.extend_from_slice(b"enc");
306 let enc_key = Sha256::digest(&enc_key_input);
307
308 let mut mac_key_input = key.to_vec();
309 mac_key_input.extend_from_slice(&nonce);
310 mac_key_input.extend_from_slice(b"mac");
311 let mac_key = Sha256::digest(&mac_key_input);
312
313 let ciphertext: Vec<u8> = plaintext.iter().enumerate().map(|(i, &b)| {
315 let mut block_input = enc_key.to_vec();
316 block_input.extend_from_slice(&(i as u64).to_le_bytes());
317 let block = Sha256::digest(&block_input);
318 b ^ block[i % 32]
319 }).collect();
320
321 let mut mac_input = mac_key.to_vec();
323 mac_input.extend_from_slice(&nonce);
324 mac_input.extend_from_slice(&ciphertext);
325 let mac = Sha256::digest(&mac_input);
326
327 let mut out = Vec::with_capacity(12 + 32 + ciphertext.len());
329 out.extend_from_slice(&nonce);
330 out.extend_from_slice(&mac);
331 out.extend_from_slice(&ciphertext);
332
333 Ok((out, nonce.to_vec()))
334}
335
336pub fn aes_gcm_decrypt(key: &[u8; 32], enc_data: &[u8], _nonce_unused: &[u8]) -> Result<Vec<u8>, String> {
337 if enc_data.len() < 44 {
338 return Err("ciphertext too short".into());
339 }
340 use sha2::Sha256;
341
342 let nonce = &enc_data[..12];
343 let stored_mac = &enc_data[12..44];
344 let ciphertext = &enc_data[44..];
345
346 let nonce_arr: [u8; 12] = nonce.try_into().unwrap();
347
348 let mut enc_key_input = key.to_vec();
349 enc_key_input.extend_from_slice(&nonce_arr);
350 enc_key_input.extend_from_slice(b"enc");
351 let enc_key = Sha256::digest(&enc_key_input);
352
353 let mut mac_key_input = key.to_vec();
354 mac_key_input.extend_from_slice(&nonce_arr);
355 mac_key_input.extend_from_slice(b"mac");
356 let mac_key = Sha256::digest(&mac_key_input);
357
358 let mut mac_input = mac_key.to_vec();
360 mac_input.extend_from_slice(&nonce_arr);
361 mac_input.extend_from_slice(ciphertext);
362 let computed_mac = Sha256::digest(&mac_input);
363
364 let mac_ok = stored_mac.iter().zip(computed_mac.iter())
366 .fold(0u8, |acc, (a, b)| acc | (a ^ b)) == 0;
367
368 if !mac_ok {
369 return Err("MAC verification failed — key file may be corrupt or wrong machine".into());
370 }
371
372 let plaintext: Vec<u8> = ciphertext.iter().enumerate().map(|(i, &b)| {
373 let mut block_input = enc_key.to_vec();
374 block_input.extend_from_slice(&(i as u64).to_le_bytes());
375 let block = Sha256::digest(&block_input);
376 b ^ block[i % 32]
377 }).collect();
378
379 Ok(plaintext)
380}
381
382pub fn derive_machine_key(store_dir: &Path) -> Result<[u8; 32], KeyError> {
385 if let Ok(id) = fs::read_to_string("/etc/machine-id") {
387 let trimmed = id.trim();
388 if !trimmed.is_empty() {
389 let mut h = Sha256::new();
390 h.update(trimmed.as_bytes());
391 h.update(store_dir.to_string_lossy().as_bytes());
392 return Ok(h.finalize().into());
393 }
394 }
395
396 #[cfg(target_os = "macos")]
408 {
409 let hostname = std::process::Command::new("hostname")
410 .output()
411 .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string())
412 .unwrap_or_default();
413 let username = std::env::var("USER").unwrap_or_default();
414 if !hostname.is_empty() && !username.is_empty() {
415 let mut h = Sha256::new();
416 h.update(b"treeship-machine-key:");
417 h.update(hostname.as_bytes());
418 h.update(b":");
419 h.update(username.as_bytes());
420 h.update(b":");
421 h.update(store_dir.to_string_lossy().as_bytes());
422 return Ok(h.finalize().into());
423 }
424 }
425
426 let home = std::env::var("HOME")
429 .map(std::path::PathBuf::from)
430 .map_err(|_| KeyError::Crypto("HOME not set".to_string()))?;
431 let seed_path = home.join(".treeship").join("machine_seed");
432 let seed = if seed_path.exists() {
433 fs::read_to_string(&seed_path).map_err(KeyError::Io)?
434 } else {
435 let mut bytes = [0u8; 32];
436 rand::thread_rng().fill_bytes(&mut bytes);
437 let seed_hex = hex_encode(&bytes);
438 let _ = fs::create_dir_all(seed_path.parent().unwrap_or(Path::new(".")));
439 fs::write(&seed_path, &seed_hex).map_err(KeyError::Io)?;
440 #[cfg(unix)]
441 {
442 use std::os::unix::fs::PermissionsExt;
443 let _ = fs::set_permissions(&seed_path, fs::Permissions::from_mode(0o600));
444 }
445 seed_hex
446 };
447
448 let mut h = Sha256::new();
449 h.update(b"treeship-machine-key-fallback:");
450 h.update(seed.trim().as_bytes());
451 h.update(b":");
452 h.update(store_dir.to_string_lossy().as_bytes());
453 Ok(h.finalize().into())
454}
455
456pub fn derive_machine_key_stable(store_dir: &Path) -> Result<[u8; 32], KeyError> {
460 if let Ok(id) = fs::read_to_string("/etc/machine-id") {
462 let trimmed = id.trim();
463 if !trimmed.is_empty() {
464 let mut h = Sha256::new();
465 h.update(b"treeship-machine-key-v2:");
466 h.update(trimmed.as_bytes());
467 h.update(b":");
468 h.update(store_dir.to_string_lossy().as_bytes());
469 return Ok(h.finalize().into());
470 }
471 }
472
473 #[cfg(target_os = "macos")]
476 {
477 if let Ok(output) = std::process::Command::new("ioreg")
478 .args(["-rd1", "-c", "IOPlatformExpertDevice"])
479 .output()
480 {
481 let stdout = String::from_utf8_lossy(&output.stdout);
482 for line in stdout.lines() {
483 if line.contains("IOPlatformSerialNumber") {
484 if let Some(serial) = line.split('"').nth(3) {
485 if !serial.is_empty() {
486 let mut h = Sha256::new();
487 h.update(b"treeship-machine-key-v2:");
488 h.update(serial.as_bytes());
489 h.update(b":");
490 h.update(store_dir.to_string_lossy().as_bytes());
491 return Ok(h.finalize().into());
492 }
493 }
494 }
495 }
496 }
497 }
498
499 let home = std::env::var("HOME")
502 .map(std::path::PathBuf::from)
503 .map_err(|_| KeyError::Crypto("HOME not set".to_string()))?;
504 let seed_dir = home.join(".treeship").join(".internal");
505 let _ = fs::create_dir_all(&seed_dir);
506 #[cfg(unix)]
507 {
508 use std::os::unix::fs::PermissionsExt;
509 let _ = fs::set_permissions(&seed_dir, fs::Permissions::from_mode(0o700));
510 }
511
512 let seed_path = seed_dir.join("machine_seed_v2");
513 let seed = if seed_path.exists() {
514 fs::read_to_string(&seed_path).map_err(KeyError::Io)?
515 } else {
516 let mut bytes = [0u8; 32];
517 rand::thread_rng().fill_bytes(&mut bytes);
518 let seed_hex = hex_encode(&bytes);
519 fs::write(&seed_path, &seed_hex).map_err(KeyError::Io)?;
520 #[cfg(unix)]
521 {
522 use std::os::unix::fs::PermissionsExt;
523 let _ = fs::set_permissions(&seed_path, fs::Permissions::from_mode(0o600));
524 }
525 seed_hex
526 };
527
528 let mut h = Sha256::new();
529 h.update(b"treeship-machine-key-v2-fallback:");
530 h.update(seed.trim().as_bytes());
531 h.update(b":");
532 h.update(store_dir.to_string_lossy().as_bytes());
533 Ok(h.finalize().into())
534}
535
536fn new_key_id() -> KeyId {
539 let mut b = [0u8; 8];
540 rand::thread_rng().fill_bytes(&mut b);
541 format!("key_{}", hex_encode(&b))
542}
543
544fn fingerprint(pub_key: &[u8]) -> String {
545 let h = Sha256::digest(pub_key);
546 hex_encode(&h[..8])
547}
548
549fn hex_encode(b: &[u8]) -> String {
550 b.iter().fold(String::new(), |mut s, byte| {
551 s.push_str(&format!("{:02x}", byte));
552 s
553 })
554}
555
556fn write_file_600(path: &Path, data: &[u8]) -> Result<(), KeyError> {
557 let mut f = fs::OpenOptions::new()
558 .write(true)
559 .create(true)
560 .truncate(true)
561 .open(path)?;
562 f.write_all(data)?;
563 #[cfg(unix)]
565 {
566 use std::os::unix::fs::PermissionsExt;
567 fs::set_permissions(path, fs::Permissions::from_mode(0o600))?;
568 }
569 Ok(())
570}
571
572fn unix_now() -> u64 {
573 use std::time::{SystemTime, UNIX_EPOCH};
574 SystemTime::now()
575 .duration_since(UNIX_EPOCH)
576 .unwrap_or_default()
577 .as_secs()
578}
579
580#[cfg(test)]
581mod tests {
582 use super::*;
583
584 fn temp_dir_path() -> PathBuf {
585 let mut p = std::env::temp_dir();
586 p.push(format!("treeship-test-{}", {
587 let mut b = [0u8; 4];
588 rand::thread_rng().fill_bytes(&mut b);
589 hex_encode(&b)
590 }));
591 p
592 }
593
594 fn make_store() -> (Store, PathBuf) {
595 let dir = temp_dir_path();
596 let store = Store::open(&dir).unwrap();
597 (store, dir)
598 }
599
600 fn cleanup(dir: PathBuf) {
601 let _ = fs::remove_dir_all(dir);
602 }
603
604 #[test]
605 fn generate_key() {
606 let (store, dir) = make_store();
607 let info = store.generate(true).unwrap();
608 assert!(info.id.starts_with("key_"));
609 assert_eq!(info.algorithm, "ed25519");
610 assert!(!info.fingerprint.is_empty());
611 assert_eq!(info.public_key.len(), 32);
612 cleanup(dir);
613 }
614
615 #[test]
616 fn default_signer_works() {
617 let (store, dir) = make_store();
618 store.generate(true).unwrap();
619 let signer = store.default_signer().unwrap();
620 assert!(!signer.key_id().is_empty());
621 let pae = crate::attestation::pae("text/plain", b"test");
622 let sig = signer.sign(&pae).unwrap();
623 assert_eq!(sig.len(), 64);
624 cleanup(dir);
625 }
626
627 #[test]
628 fn encrypt_decrypt_roundtrip() {
629 let key = [42u8; 32];
630 let plaintext = b"super secret private key material here!";
631 let (enc, nonce) = aes_gcm_encrypt(&key, plaintext).unwrap();
632 let dec = aes_gcm_decrypt(&key, &enc, &nonce).unwrap();
633 assert_eq!(dec, plaintext);
634 }
635
636 #[test]
637 fn decrypt_wrong_key_fails() {
638 let key = [42u8; 32];
639 let wrong = [99u8; 32];
640 let (enc, nonce) = aes_gcm_encrypt(&key, b"secret").unwrap();
641 assert!(aes_gcm_decrypt(&wrong, &enc, &nonce).is_err());
642 }
643
644 #[test]
645 fn persist_and_reload() {
646 let (store, dir) = make_store();
647 let info = store.generate(true).unwrap();
648
649 let store2 = Store::open(&dir).unwrap();
651 let signer = store2.signer(&info.id).unwrap();
652 assert_eq!(signer.key_id(), info.id);
653
654 let verifier = {
657 use crate::attestation::Verifier;
658 use ed25519_dalek::VerifyingKey;
659 let pk_bytes: [u8; 32] = info.public_key.try_into().unwrap();
660 let vk = VerifyingKey::from_bytes(&pk_bytes).unwrap();
661 let mut v = Verifier::new(std::collections::HashMap::new());
662 v.add_key(info.id.clone(), vk);
663 v
664 };
665
666 use crate::attestation::sign;
667 use crate::statements::ActionStatement;
668 let stmt = ActionStatement::new("agent://test", "tool.call");
669 let pt = crate::statements::payload_type("action");
670 let signed = sign(&pt, &stmt, signer.as_ref()).unwrap();
671 verifier.verify(&signed.envelope).unwrap();
672
673 cleanup(dir);
674 }
675
676 #[test]
677 fn list_keys() {
678 let (store, dir) = make_store();
679 store.generate(true).unwrap();
680 store.generate(false).unwrap();
681
682 let keys = store.list().unwrap();
683 assert_eq!(keys.len(), 2);
684 assert_eq!(keys.iter().filter(|k| k.is_default).count(), 1);
685 cleanup(dir);
686 }
687
688 #[test]
689 fn no_default_key_errors() {
690 let (store, dir) = make_store();
691 assert!(store.default_signer().is_err());
692 cleanup(dir);
693 }
694}