1use crate::{OpenADPError, Result};
8use crate::keygen::{generate_encryption_key, recover_encryption_key};
9use crate::client::{ServerInfo, get_servers};
10use serde::{Deserialize, Serialize};
11use aes_gcm::{Aes256Gcm, Key, Nonce, KeyInit};
12use aes_gcm::aead::Aead;
13use rand::Rng;
14use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
15use sha2::{Sha256, Digest};
16use hex;
17
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct WrappedSecret {
22 pub nonce: String,
23 pub ciphertext: String,
24 pub tag: String,
25}
26
27#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct OcryptMetadata {
30 pub servers: Vec<String>,
32 pub threshold: usize,
33 pub version: String,
34 pub auth_code: String,
35 pub user_id: String,
36
37 pub wrapped_long_term_secret: WrappedSecret,
39 pub backup_id: String,
40 pub app_id: String,
41 pub max_guesses: i32,
42 pub ocrypt_version: String,
43}
44
45pub async fn register(
86 user_id: &str,
87 app_id: &str,
88 long_term_secret: &[u8],
89 pin: &str,
90 max_guesses: i32,
91 servers_url: &str,
92) -> Result<Vec<u8>> {
93 register_with_bid(user_id, app_id, long_term_secret, pin, max_guesses, "even", servers_url).await
94}
95
96pub async fn recover(
135 metadata_bytes: &[u8],
136 pin: &str,
137 servers_url: &str,
138) -> Result<(Vec<u8>, u32, Vec<u8>)> {
139 println!("📋 Step 1: Recovering with existing backup...");
141 let (secret, remaining) = recover_without_refresh(metadata_bytes, pin, servers_url).await?;
142
143 let updated_metadata = match attempt_backup_refresh(&secret, metadata_bytes, pin, servers_url).await {
145 Ok(new_metadata) => {
146 println!("✅ Backup refresh successful");
147 new_metadata
148 }
149 Err(e) => {
150 println!("⚠️ Backup refresh failed (using original): {}", e);
151 metadata_bytes.to_vec()
152 }
153 };
154
155 Ok((secret, remaining, updated_metadata))
156}
157
158async fn register_with_bid(
160 user_id: &str,
161 app_id: &str,
162 long_term_secret: &[u8],
163 pin: &str,
164 max_guesses: i32,
165 backup_id: &str,
166 servers_url: &str,
167) -> Result<Vec<u8>> {
168 validate_inputs(user_id, app_id, long_term_secret, pin, max_guesses)?;
170
171 println!("🔐 Protecting secret for user: {}", user_id);
172 println!("📱 Application: {}", app_id);
173 println!("🔑 Secret length: {} bytes", long_term_secret.len());
174
175 println!("🌐 Discovering OpenADP servers...");
177 let server_infos = discover_servers(servers_url).await?;
178
179 if server_infos.is_empty() {
180 return Err(OpenADPError::NoServers);
181 }
182
183 let selected_servers = if server_infos.len() > 15 { use rand::seq::SliceRandom;
186 let mut servers = server_infos;
187 servers.shuffle(&mut rand::thread_rng());
188 servers.into_iter().take(15).collect()
189 } else {
190 server_infos
191 };
192
193 println!("📋 Using {} servers for registration", selected_servers.len());
194
195 println!("🔄 Using backup ID: {}", backup_id);
197 println!("🔑 Generating encryption key using OpenADP servers...");
198
199 let identity = crate::keygen::Identity::new(
201 user_id.to_string(), app_id.to_string(), backup_id.to_string() );
205
206 let key_result = generate_encryption_key(&identity, pin, max_guesses, 0, selected_servers).await?;
207
208 if key_result.encryption_key.is_none() {
209 return Err(OpenADPError::Server(
210 key_result.error.unwrap_or_else(|| "Key generation failed".to_string())
211 ));
212 }
213
214 let encryption_key = key_result.encryption_key.unwrap();
215 let auth_codes = key_result.auth_codes.unwrap();
216 let server_infos = key_result.server_infos.unwrap();
217 let threshold = key_result.threshold.unwrap();
218
219 println!("✅ Generated encryption key with {} servers", server_infos.len());
220
221 println!("🔐 Wrapping long-term secret...");
223 let wrapped_secret = wrap_secret(long_term_secret, &encryption_key)?;
224
225 let server_urls: Vec<String> = server_infos.iter().map(|s| s.url.clone()).collect();
227
228 let metadata = OcryptMetadata {
229 servers: server_urls,
230 threshold,
231 version: "1.0".to_string(),
232 auth_code: auth_codes.base_auth_code,
233 user_id: user_id.to_string(),
234 wrapped_long_term_secret: wrapped_secret,
235 backup_id: backup_id.to_string(),
236 app_id: app_id.to_string(),
237 max_guesses,
238 ocrypt_version: "1.0".to_string(),
239 };
240
241 let metadata_bytes = serde_json::to_vec(&metadata)?;
242 println!("📦 Created metadata ({} bytes)", metadata_bytes.len());
243 println!("🎯 Threshold: {}-of-{} recovery", metadata.threshold, metadata.servers.len());
244
245 Ok(metadata_bytes)
246}
247
248async fn recover_without_refresh(
250 metadata_bytes: &[u8],
251 pin: &str,
252 servers_url: &str,
253) -> Result<(Vec<u8>, u32)> {
254 let metadata: OcryptMetadata = serde_json::from_slice(metadata_bytes)?;
256
257 let identity = crate::keygen::Identity::new(
260 metadata.user_id.clone(), metadata.app_id.clone(), metadata.backup_id.clone() );
264
265 let server_infos = if servers_url.is_empty() {
267 metadata.servers.iter().map(|url| ServerInfo {
269 url: url.clone(),
270 public_key: String::new(),
271 country: String::new(),
272 remaining_guesses: None,
273 }).collect()
274 } else {
275 discover_servers(servers_url).await?
277 };
278
279 let mut server_auth_codes = std::collections::HashMap::new();
281
282 for server_url in &metadata.servers {
284 let combined = format!("{}:{}", metadata.auth_code, server_url);
285 let mut hasher = Sha256::new();
286 hasher.update(combined.as_bytes());
287 let hash = hasher.finalize();
288 let server_code = hex::encode(&hash);
289 server_auth_codes.insert(server_url.clone(), server_code);
290 }
291
292 let auth_codes = crate::keygen::AuthCodes {
293 base_auth_code: metadata.auth_code.clone(),
294 server_auth_codes,
295 };
296
297 let recovery_result = recover_encryption_key(
298 &identity,
299 pin,
300 server_infos,
301 metadata.threshold,
302 auth_codes
303 ).await?;
304
305 if recovery_result.encryption_key.is_none() {
306 return Err(OpenADPError::Server(
307 recovery_result.error.unwrap_or_else(|| "Key recovery failed".to_string())
308 ));
309 }
310
311 let encryption_key = recovery_result.encryption_key.unwrap();
312
313 let secret = unwrap_secret(&metadata.wrapped_long_term_secret, &encryption_key)?;
315
316 Ok((secret, 0)) }
318
319async fn attempt_backup_refresh(
321 secret: &[u8],
322 metadata_bytes: &[u8],
323 pin: &str,
324 servers_url: &str,
325) -> Result<Vec<u8>> {
326 let metadata: OcryptMetadata = serde_json::from_slice(metadata_bytes)?;
327
328 let new_backup_id = generate_next_backup_id(&metadata.backup_id);
330
331 println!("🔄 Attempting backup refresh: {} → {}", metadata.backup_id, new_backup_id);
332
333 let new_metadata = register_with_bid(
335 &metadata.user_id,
336 &metadata.app_id,
337 secret,
338 pin,
339 metadata.max_guesses,
340 &new_backup_id,
341 servers_url,
342 ).await?;
343
344 let (recovered_secret, _) = recover_without_refresh(&new_metadata, pin, servers_url).await?;
346
347 if recovered_secret == secret {
348 println!("✅ Two-phase commit verification successful");
349 Ok(new_metadata)
350 } else {
351 Err(OpenADPError::Server("Two-phase commit verification failed".to_string()))
352 }
353}
354
355async fn discover_servers(servers_url: &str) -> Result<Vec<ServerInfo>> {
357 let registry_url = if servers_url.is_empty() {
358 crate::DEFAULT_REGISTRY_URL
359 } else {
360 servers_url
361 };
362
363 println!("🌐 Discovering servers from registry: {}", registry_url);
364
365 let servers = get_servers(registry_url).await?;
366
367 println!(" ✅ Successfully fetched {} servers from registry", servers.len());
368 println!(" 📋 {} servers are live and ready", servers.len());
369
370 Ok(servers)
371}
372
373fn wrap_secret(secret: &[u8], key: &[u8]) -> Result<WrappedSecret> {
375 if key.len() != 32 {
376 return Err(OpenADPError::Crypto("Key must be 32 bytes".to_string()));
377 }
378
379 let key = Key::<Aes256Gcm>::from_slice(key);
380 let cipher = Aes256Gcm::new(key);
381
382 let nonce_bytes: [u8; 12] = rand::thread_rng().gen();
383 let nonce = Nonce::from_slice(&nonce_bytes);
384
385 let ciphertext = cipher.encrypt(nonce, secret)
386 .map_err(|e| OpenADPError::Crypto(format!("AES-GCM encryption failed: {}", e)))?;
387
388 let (encrypted_data, tag) = ciphertext.split_at(ciphertext.len() - 16);
390
391 Ok(WrappedSecret {
392 nonce: BASE64.encode(&nonce_bytes),
393 ciphertext: BASE64.encode(encrypted_data),
394 tag: BASE64.encode(tag),
395 })
396}
397
398fn unwrap_secret(wrapped: &WrappedSecret, key: &[u8]) -> Result<Vec<u8>> {
400 if key.len() != 32 {
401 return Err(OpenADPError::Crypto("Key must be 32 bytes".to_string()));
402 }
403
404 let key = Key::<Aes256Gcm>::from_slice(key);
405 let cipher = Aes256Gcm::new(key);
406
407 let nonce_bytes = BASE64.decode(&wrapped.nonce)
408 .map_err(|e| OpenADPError::Crypto(format!("Invalid nonce: {}", e)))?;
409 let nonce = Nonce::from_slice(&nonce_bytes);
410
411 let encrypted_data = BASE64.decode(&wrapped.ciphertext)
412 .map_err(|e| OpenADPError::Crypto(format!("Invalid ciphertext: {}", e)))?;
413 let tag = BASE64.decode(&wrapped.tag)
414 .map_err(|e| OpenADPError::Crypto(format!("Invalid tag: {}", e)))?;
415
416 let mut ciphertext_with_tag = encrypted_data;
418 ciphertext_with_tag.extend_from_slice(&tag);
419
420 let plaintext = cipher.decrypt(nonce, ciphertext_with_tag.as_slice())
421 .map_err(|_| OpenADPError::Authentication("Invalid PIN or corrupted data".to_string()))?;
422
423 Ok(plaintext)
424}
425
426fn generate_next_backup_id(current_backup_id: &str) -> String {
428 match current_backup_id {
429 "even" => "odd".to_string(),
430 "odd" => "even".to_string(),
431 _ => {
432 if current_backup_id.starts_with('v') {
434 let version_num: u32 = current_backup_id[1..].parse().unwrap_or(1);
435 format!("v{}", version_num + 1)
436 } else {
437 use std::time::{SystemTime, UNIX_EPOCH};
439 let timestamp = SystemTime::now()
440 .duration_since(UNIX_EPOCH)
441 .unwrap()
442 .as_secs();
443 format!("{}_v{}", current_backup_id, timestamp)
444 }
445 }
446 }
447}
448
449fn validate_inputs(
451 user_id: &str,
452 app_id: &str,
453 long_term_secret: &[u8],
454 pin: &str,
455 max_guesses: i32,
456) -> Result<()> {
457 if user_id.is_empty() {
458 return Err(OpenADPError::InvalidInput("user_id cannot be empty".to_string()));
459 }
460
461 if app_id.is_empty() {
462 return Err(OpenADPError::InvalidInput("app_id cannot be empty".to_string()));
463 }
464
465 if long_term_secret.is_empty() {
466 return Err(OpenADPError::InvalidInput("long_term_secret cannot be empty".to_string()));
467 }
468
469 if pin.is_empty() {
470 return Err(OpenADPError::InvalidInput("pin cannot be empty".to_string()));
471 }
472
473 if max_guesses <= 0 {
474 return Err(OpenADPError::InvalidInput("max_guesses must be at least 1".to_string()));
475 }
476
477 Ok(())
478}
479
480#[cfg(test)]
481mod tests {
482 use super::*;
483
484 #[test]
485 fn test_backup_id_generation() {
486 assert_eq!(generate_next_backup_id("even"), "odd");
487 assert_eq!(generate_next_backup_id("odd"), "even");
488 assert_eq!(generate_next_backup_id("v1"), "v2");
489 assert_eq!(generate_next_backup_id("v42"), "v43");
490
491 let timestamped = generate_next_backup_id("production");
492 assert!(timestamped.starts_with("production_v"));
493 }
494
495 #[test]
496 fn test_secret_wrapping() {
497 let secret = b"test_secret";
498 let key = [42u8; 32];
499
500 let wrapped = wrap_secret(secret, &key).unwrap();
501 let unwrapped = unwrap_secret(&wrapped, &key).unwrap();
502
503 assert_eq!(secret, unwrapped.as_slice());
504 }
505
506 #[test]
507 fn test_input_validation() {
508 assert!(validate_inputs("", "app", b"secret", "pin", 10).is_err());
509 assert!(validate_inputs("user", "", b"secret", "pin", 10).is_err());
510 assert!(validate_inputs("user", "app", b"", "pin", 10).is_err());
511 assert!(validate_inputs("user", "app", b"secret", "", 10).is_err());
512 assert!(validate_inputs("user", "app", b"secret", "pin", 0).is_err());
513 assert!(validate_inputs("user", "app", b"secret", "pin", 10).is_ok());
514 }
515
516 #[test]
517 fn test_metadata_serialization() {
518 let wrapped_secret = WrappedSecret {
519 nonce: "test_nonce".to_string(),
520 ciphertext: "test_ciphertext".to_string(),
521 tag: "test_tag".to_string(),
522 };
523
524 let metadata = OcryptMetadata {
525 servers: vec!["server1".to_string(), "server2".to_string()],
526 threshold: 2,
527 version: "1.0".to_string(),
528 auth_code: "auth123".to_string(),
529 user_id: "user@example.com".to_string(),
530 wrapped_long_term_secret: wrapped_secret,
531 backup_id: "even".to_string(),
532 app_id: "test_app".to_string(),
533 max_guesses: 10,
534 ocrypt_version: "1.0".to_string(),
535 };
536
537 let serialized = serde_json::to_vec(&metadata).unwrap();
538 let deserialized: OcryptMetadata = serde_json::from_slice(&serialized).unwrap();
539
540 assert_eq!(metadata.user_id, deserialized.user_id);
541 assert_eq!(metadata.app_id, deserialized.app_id);
542 assert_eq!(metadata.threshold, deserialized.threshold);
543 }
544}