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