this repo has no description
lewis.moe
1use base32::Alphabet;
2use rand::RngCore;
3use subtle::ConstantTimeEq;
4use totp_rs::{Algorithm, TOTP};
5
6const TOTP_DIGITS: usize = 6;
7const TOTP_STEP: u64 = 30;
8const TOTP_SECRET_LENGTH: usize = 20;
9
10pub fn generate_totp_secret() -> Vec<u8> {
11 let mut secret = vec![0u8; TOTP_SECRET_LENGTH];
12 rand::thread_rng().fill_bytes(&mut secret);
13 secret
14}
15
16pub fn encrypt_totp_secret(secret: &[u8]) -> Result<Vec<u8>, String> {
17 crate::config::encrypt_key(secret)
18}
19
20pub fn decrypt_totp_secret(encrypted: &[u8], version: i32) -> Result<Vec<u8>, String> {
21 crate::config::decrypt_key(encrypted, Some(version))
22}
23
24fn create_totp(
25 secret: Vec<u8>,
26 issuer: Option<String>,
27 account_name: String,
28) -> Result<TOTP, String> {
29 TOTP::new(
30 Algorithm::SHA1,
31 TOTP_DIGITS,
32 1,
33 TOTP_STEP,
34 secret,
35 issuer,
36 account_name,
37 )
38 .map_err(|e| format!("Failed to create TOTP: {}", e))
39}
40
41pub fn verify_totp_code(secret: &[u8], code: &str) -> bool {
42 let code = code.trim();
43 if code.len() != TOTP_DIGITS {
44 return false;
45 }
46
47 let Ok(totp) = create_totp(secret.to_vec(), None, String::new()) else {
48 return false;
49 };
50
51 let now = std::time::SystemTime::now()
52 .duration_since(std::time::UNIX_EPOCH)
53 .map(|d| d.as_secs())
54 .unwrap_or(0);
55
56 for offset in [-1i64, 0, 1] {
57 let time = (now as i64 + offset * TOTP_STEP as i64) as u64;
58 let expected = totp.generate(time);
59 let is_valid: bool = code.as_bytes().ct_eq(expected.as_bytes()).into();
60 if is_valid {
61 return true;
62 }
63 }
64
65 false
66}
67
68pub fn generate_totp_uri(secret: &[u8], account_name: &str, issuer: &str) -> String {
69 let secret_base32 = base32::encode(Alphabet::Rfc4648 { padding: false }, secret);
70 format!(
71 "otpauth://totp/{}:{}?secret={}&issuer={}&algorithm=SHA1&digits={}&period={}",
72 urlencoding::encode(issuer),
73 urlencoding::encode(account_name),
74 secret_base32,
75 urlencoding::encode(issuer),
76 TOTP_DIGITS,
77 TOTP_STEP
78 )
79}
80
81pub fn generate_qr_png_base64(
82 secret: &[u8],
83 account_name: &str,
84 issuer: &str,
85) -> Result<String, String> {
86 use base64::{Engine, engine::general_purpose::STANDARD};
87
88 let totp = create_totp(
89 secret.to_vec(),
90 Some(issuer.to_string()),
91 account_name.to_string(),
92 )?;
93
94 let qr_png = totp
95 .get_qr_png()
96 .map_err(|e| format!("Failed to generate QR code: {}", e))?;
97
98 Ok(STANDARD.encode(qr_png))
99}
100
101const BACKUP_CODE_ALPHABET: &[u8] = b"23456789ABCDEFGHJKMNPQRSTUVWXYZ";
102const BACKUP_CODE_LENGTH: usize = 8;
103const BACKUP_CODE_COUNT: usize = 10;
104const BACKUP_CODE_BCRYPT_COST: u32 = 10;
105
106pub fn generate_backup_codes() -> Vec<String> {
107 let mut codes = Vec::with_capacity(BACKUP_CODE_COUNT);
108 let mut rng = rand::thread_rng();
109
110 for _ in 0..BACKUP_CODE_COUNT {
111 let mut code = String::with_capacity(BACKUP_CODE_LENGTH);
112 for _ in 0..BACKUP_CODE_LENGTH {
113 let idx = (rng.next_u32() as usize) % BACKUP_CODE_ALPHABET.len();
114 code.push(BACKUP_CODE_ALPHABET[idx] as char);
115 }
116 codes.push(code);
117 }
118
119 codes
120}
121
122pub fn hash_backup_code(code: &str) -> Result<String, String> {
123 bcrypt::hash(code, BACKUP_CODE_BCRYPT_COST).map_err(|e| format!("Failed to hash code: {}", e))
124}
125
126pub fn verify_backup_code(code: &str, hash: &str) -> bool {
127 bcrypt::verify(code, hash).unwrap_or(false)
128}
129
130pub fn is_backup_code_format(code: &str) -> bool {
131 let code = code.trim().to_uppercase();
132 code.len() == BACKUP_CODE_LENGTH
133 && code
134 .chars()
135 .all(|c| BACKUP_CODE_ALPHABET.contains(&(c as u8)))
136}
137
138#[cfg(test)]
139mod tests {
140 use super::*;
141
142 #[test]
143 fn test_generate_totp_secret() {
144 let secret = generate_totp_secret();
145 assert_eq!(secret.len(), TOTP_SECRET_LENGTH);
146 }
147
148 #[test]
149 fn test_verify_totp_code() {
150 let secret = generate_totp_secret();
151 let totp = create_totp(secret.clone(), None, String::new()).unwrap();
152 let code = totp.generate_current().unwrap();
153 assert!(verify_totp_code(&secret, &code));
154 assert!(!verify_totp_code(&secret, "000000"));
155 }
156
157 #[test]
158 fn test_generate_totp_uri() {
159 let secret = vec![0u8; 20];
160 let uri = generate_totp_uri(&secret, "test@example.com", "TestPDS");
161 assert!(uri.starts_with("otpauth://totp/"));
162 assert!(uri.contains("secret="));
163 assert!(uri.contains("issuer=TestPDS"));
164 }
165
166 #[test]
167 fn test_backup_codes() {
168 let codes = generate_backup_codes();
169 assert_eq!(codes.len(), BACKUP_CODE_COUNT);
170 for code in &codes {
171 assert_eq!(code.len(), BACKUP_CODE_LENGTH);
172 assert!(is_backup_code_format(code));
173 }
174 }
175
176 #[test]
177 fn test_backup_code_hash_verify() {
178 let codes = generate_backup_codes();
179 let code = &codes[0];
180 let hash = hash_backup_code(code).unwrap();
181 assert!(verify_backup_code(code, &hash));
182 assert!(!verify_backup_code("WRONGCOD", &hash));
183 }
184
185 #[test]
186 fn test_is_backup_code_format() {
187 assert!(is_backup_code_format("ABCD2345"));
188 assert!(is_backup_code_format(" abcd2345 "));
189 assert!(!is_backup_code_format("ABCD234"));
190 assert!(!is_backup_code_format("ABCD23456"));
191 assert!(!is_backup_code_format("ABCD234O"));
192 assert!(!is_backup_code_format("ABCD2341"));
193 }
194}