koprogo_api/domain/entities/
two_factor_secret.rs

1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use uuid::Uuid;
4
5/// Two-factor authentication secret for TOTP (Time-based One-Time Password)
6/// Stores encrypted TOTP secret and backup codes for account recovery
7#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
8pub struct TwoFactorSecret {
9    pub id: Uuid,
10    pub user_id: Uuid,
11    pub secret_encrypted: String, // Base32-encoded TOTP secret (AES-256 encrypted at application level)
12    pub backup_codes_encrypted: Vec<String>, // 10 backup codes (bcrypt hashed)
13    pub is_enabled: bool,
14    pub verified_at: Option<DateTime<Utc>>, // First successful verification
15    pub last_used_at: Option<DateTime<Utc>>, // Last successful TOTP verification
16    pub created_at: DateTime<Utc>,
17    pub updated_at: DateTime<Utc>,
18}
19
20impl TwoFactorSecret {
21    /// Create a new 2FA secret (not yet enabled)
22    pub fn new(user_id: Uuid, secret_encrypted: String) -> Result<Self, String> {
23        // Validate secret is non-empty
24        if secret_encrypted.trim().is_empty() {
25            return Err("Secret cannot be empty".to_string());
26        }
27
28        // Generate 10 backup codes (will be encrypted by caller)
29        let backup_codes = Self::generate_backup_codes_placeholders();
30
31        Ok(Self {
32            id: Uuid::new_v4(),
33            user_id,
34            secret_encrypted,
35            backup_codes_encrypted: backup_codes,
36            is_enabled: false,
37            verified_at: None,
38            last_used_at: None,
39            created_at: Utc::now(),
40            updated_at: Utc::now(),
41        })
42    }
43
44    /// Set backup codes (encrypted/hashed by caller)
45    pub fn with_backup_codes(
46        mut self,
47        backup_codes_encrypted: Vec<String>,
48    ) -> Result<Self, String> {
49        if backup_codes_encrypted.len() != 10 {
50            return Err("Must provide exactly 10 backup codes".to_string());
51        }
52
53        self.backup_codes_encrypted = backup_codes_encrypted;
54        Ok(self)
55    }
56
57    /// Enable 2FA after successful verification
58    pub fn enable(&mut self) -> Result<(), String> {
59        if self.is_enabled {
60            return Err("2FA is already enabled".to_string());
61        }
62
63        self.is_enabled = true;
64        self.verified_at = Some(Utc::now());
65        self.updated_at = Utc::now();
66
67        Ok(())
68    }
69
70    /// Disable 2FA
71    pub fn disable(&mut self) {
72        self.is_enabled = false;
73        self.updated_at = Utc::now();
74    }
75
76    /// Mark TOTP as used (update last_used_at)
77    pub fn mark_used(&mut self) {
78        self.last_used_at = Some(Utc::now());
79        self.updated_at = Utc::now();
80    }
81
82    /// Regenerate backup codes (caller must hash/encrypt new codes)
83    pub fn regenerate_backup_codes(
84        &mut self,
85        new_backup_codes_encrypted: Vec<String>,
86    ) -> Result<(), String> {
87        if new_backup_codes_encrypted.len() != 10 {
88            return Err("Must provide exactly 10 backup codes".to_string());
89        }
90
91        self.backup_codes_encrypted = new_backup_codes_encrypted;
92        self.updated_at = Utc::now();
93
94        Ok(())
95    }
96
97    /// Remove a used backup code (caller must identify which code was used)
98    pub fn remove_backup_code(&mut self, code_index: usize) -> Result<(), String> {
99        if code_index >= self.backup_codes_encrypted.len() {
100            return Err("Invalid backup code index".to_string());
101        }
102
103        self.backup_codes_encrypted.remove(code_index);
104        self.updated_at = Utc::now();
105
106        Ok(())
107    }
108
109    /// Check if backup codes are exhausted (< 3 remaining)
110    pub fn backup_codes_low(&self) -> bool {
111        self.backup_codes_encrypted.len() < 3
112    }
113
114    /// Check if 2FA needs re-verification (not used in 90 days)
115    pub fn needs_reverification(&self) -> bool {
116        if !self.is_enabled {
117            return false;
118        }
119
120        match self.last_used_at {
121            Some(last_used) => {
122                let days_since_use = (Utc::now() - last_used).num_days();
123                days_since_use > 90
124            }
125            None => false, // Never used yet
126        }
127    }
128
129    /// Generate placeholder backup codes (will be replaced with real codes)
130    fn generate_backup_codes_placeholders() -> Vec<String> {
131        vec!["placeholder".to_string(); 10]
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    fn sample_user_id() -> Uuid {
140        Uuid::parse_str("550e8400-e29b-41d4-a716-446655440000").unwrap()
141    }
142
143    #[test]
144    fn test_create_two_factor_secret_success() {
145        let secret = TwoFactorSecret::new(
146            sample_user_id(),
147            "JBSWY3DPEHPK3PXP".to_string(), // Base32 encoded secret
148        );
149
150        assert!(secret.is_ok());
151        let s = secret.unwrap();
152        assert_eq!(s.user_id, sample_user_id());
153        assert_eq!(s.secret_encrypted, "JBSWY3DPEHPK3PXP");
154        assert!(!s.is_enabled);
155        assert!(s.verified_at.is_none());
156        assert_eq!(s.backup_codes_encrypted.len(), 10);
157    }
158
159    #[test]
160    fn test_create_two_factor_secret_empty() {
161        let secret = TwoFactorSecret::new(sample_user_id(), "".to_string());
162
163        assert!(secret.is_err());
164        assert!(secret.unwrap_err().contains("Secret cannot be empty"));
165    }
166
167    #[test]
168    fn test_with_backup_codes() {
169        let codes: Vec<String> = (0..10).map(|i| format!("CODE{}", i)).collect();
170
171        let secret = TwoFactorSecret::new(sample_user_id(), "SECRET123".to_string())
172            .unwrap()
173            .with_backup_codes(codes.clone());
174
175        assert!(secret.is_ok());
176        let s = secret.unwrap();
177        assert_eq!(s.backup_codes_encrypted.len(), 10);
178        assert_eq!(s.backup_codes_encrypted[0], "CODE0");
179    }
180
181    #[test]
182    fn test_with_backup_codes_invalid_count() {
183        let codes: Vec<String> = vec!["CODE1".to_string(), "CODE2".to_string()]; // Only 2 codes
184
185        let secret = TwoFactorSecret::new(sample_user_id(), "SECRET123".to_string())
186            .unwrap()
187            .with_backup_codes(codes);
188
189        assert!(secret.is_err());
190        assert!(secret
191            .unwrap_err()
192            .contains("Must provide exactly 10 backup codes"));
193    }
194
195    #[test]
196    fn test_enable_2fa() {
197        let mut secret = TwoFactorSecret::new(sample_user_id(), "SECRET123".to_string()).unwrap();
198
199        assert!(!secret.is_enabled);
200        assert!(secret.verified_at.is_none());
201
202        let result = secret.enable();
203        assert!(result.is_ok());
204        assert!(secret.is_enabled);
205        assert!(secret.verified_at.is_some());
206        assert!(secret.verified_at.unwrap() <= Utc::now());
207    }
208
209    #[test]
210    fn test_enable_2fa_already_enabled() {
211        let mut secret = TwoFactorSecret::new(sample_user_id(), "SECRET123".to_string()).unwrap();
212
213        secret.enable().unwrap();
214
215        let result = secret.enable();
216        assert!(result.is_err());
217        assert!(result.unwrap_err().contains("already enabled"));
218    }
219
220    #[test]
221    fn test_disable_2fa() {
222        let mut secret = TwoFactorSecret::new(sample_user_id(), "SECRET123".to_string()).unwrap();
223
224        secret.enable().unwrap();
225        assert!(secret.is_enabled);
226
227        secret.disable();
228        assert!(!secret.is_enabled);
229    }
230
231    #[test]
232    fn test_mark_used() {
233        let mut secret = TwoFactorSecret::new(sample_user_id(), "SECRET123".to_string()).unwrap();
234
235        assert!(secret.last_used_at.is_none());
236
237        secret.mark_used();
238        assert!(secret.last_used_at.is_some());
239        assert!(secret.last_used_at.unwrap() <= Utc::now());
240    }
241
242    #[test]
243    fn test_regenerate_backup_codes() {
244        let mut secret = TwoFactorSecret::new(sample_user_id(), "SECRET123".to_string()).unwrap();
245
246        let new_codes: Vec<String> = (0..10).map(|i| format!("NEWCODE{}", i)).collect();
247
248        let result = secret.regenerate_backup_codes(new_codes.clone());
249        assert!(result.is_ok());
250        assert_eq!(secret.backup_codes_encrypted.len(), 10);
251        assert_eq!(secret.backup_codes_encrypted[0], "NEWCODE0");
252    }
253
254    #[test]
255    fn test_regenerate_backup_codes_invalid_count() {
256        let mut secret = TwoFactorSecret::new(sample_user_id(), "SECRET123".to_string()).unwrap();
257
258        let new_codes: Vec<String> = vec!["CODE1".to_string()]; // Only 1 code
259
260        let result = secret.regenerate_backup_codes(new_codes);
261        assert!(result.is_err());
262    }
263
264    #[test]
265    fn test_remove_backup_code() {
266        let codes: Vec<String> = (0..10).map(|i| format!("CODE{}", i)).collect();
267        let mut secret = TwoFactorSecret::new(sample_user_id(), "SECRET123".to_string())
268            .unwrap()
269            .with_backup_codes(codes)
270            .unwrap();
271
272        assert_eq!(secret.backup_codes_encrypted.len(), 10);
273
274        let result = secret.remove_backup_code(0);
275        assert!(result.is_ok());
276        assert_eq!(secret.backup_codes_encrypted.len(), 9);
277        assert_eq!(secret.backup_codes_encrypted[0], "CODE1"); // CODE0 removed, CODE1 is now first
278    }
279
280    #[test]
281    fn test_remove_backup_code_invalid_index() {
282        let mut secret = TwoFactorSecret::new(sample_user_id(), "SECRET123".to_string()).unwrap();
283
284        let result = secret.remove_backup_code(20); // Invalid index
285        assert!(result.is_err());
286        assert!(result.unwrap_err().contains("Invalid backup code index"));
287    }
288
289    #[test]
290    fn test_backup_codes_low() {
291        let codes: Vec<String> = (0..10).map(|i| format!("CODE{}", i)).collect();
292        let mut secret = TwoFactorSecret::new(sample_user_id(), "SECRET123".to_string())
293            .unwrap()
294            .with_backup_codes(codes)
295            .unwrap();
296
297        assert!(!secret.backup_codes_low());
298
299        // Remove 8 codes (leaving 2)
300        for _ in 0..8 {
301            secret.remove_backup_code(0).unwrap();
302        }
303
304        assert_eq!(secret.backup_codes_encrypted.len(), 2);
305        assert!(secret.backup_codes_low());
306    }
307
308    #[test]
309    fn test_needs_reverification() {
310        let mut secret = TwoFactorSecret::new(sample_user_id(), "SECRET123".to_string()).unwrap();
311
312        // Not enabled
313        assert!(!secret.needs_reverification());
314
315        // Enable
316        secret.enable().unwrap();
317        assert!(!secret.needs_reverification()); // Never used yet
318
319        // Mark as used recently
320        secret.mark_used();
321        assert!(!secret.needs_reverification());
322
323        // Simulate 91 days ago (needs reverification)
324        secret.last_used_at = Some(Utc::now() - chrono::Duration::days(91));
325        assert!(secret.needs_reverification());
326    }
327}