use std::io::ErrorKind; use base32::Alphabet; use rand::Rng; use totp_rfc6238::{HashAlgorithm, TotpGenerator}; use crate::data::app_config::AppConfig; use crate::data::user::User; use crate::utils::err::Res; use crate::utils::time::time; const BASE32_ALPHABET: Alphabet = Alphabet::RFC4648 { padding: true }; const NUM_DIGITS: usize = 6; const PERIOD: u64 = 30; #[derive(Clone, serde::Serialize, serde::Deserialize, Debug)] pub struct TotpKey { encoded: String, } impl TotpKey { /// Generate a new TOTP key pub fn new_random() -> Self { let random_bytes = rand::thread_rng().gen::<[u8; 20]>(); Self { encoded: base32::encode(BASE32_ALPHABET, &random_bytes), } } /// Get a key from an encoded secret pub fn from_encoded_secret(s: &str) -> Self { Self { encoded: s.to_string(), } } /// Get QrCode URL for user /// /// Based on pub fn url_for_user(&self, u: &User, conf: &AppConfig) -> String { format!( "otpauth://totp/{}:{}?secret={}&issuer={}&algorithm=SHA1&digits={}&period={}", urlencoding::encode(conf.domain_name_without_port()), urlencoding::encode(&u.username), self.encoded, urlencoding::encode(conf.domain_name_without_port()), NUM_DIGITS, PERIOD, ) } /// Get account name pub fn account_name(&self, u: &User, conf: &AppConfig) -> String { format!( "{}:{}", urlencoding::encode(conf.domain_name_without_port()), urlencoding::encode(&u.username) ) } /// Get current secret in base32 format pub fn get_secret(&self) -> String { self.encoded.to_string() } /// Get current code pub fn current_code(&self) -> Res { self.get_code_at(time) } /// Get previous code pub fn previous_code(&self) -> Res { self.get_code_at(|| time() - PERIOD) } /// Get following code pub fn following_code(&self) -> Res { self.get_code_at(|| time() + PERIOD) } /// Get the code at a specific time fn get_code_at u64>(&self, get_time: F) -> Res { let gen = TotpGenerator::new() .set_digit(NUM_DIGITS) .unwrap() .set_step(PERIOD) .unwrap() .set_hash_algorithm(HashAlgorithm::SHA1) .build(); let key = match base32::decode(BASE32_ALPHABET, &self.encoded) { None => { return Err(Box::new(std::io::Error::new( ErrorKind::Other, "Failed to decode base32 secret!", ))); } Some(k) => k, }; Ok(gen.get_code_with(&key, get_time)) } /// Check a code's validity pub fn check_code(&self, code: &str) -> Res { Ok(self.previous_code()?.eq(code) || self.current_code()?.eq(code) || self.following_code()?.eq(code)) } } #[cfg(test)] mod test { use crate::data::totp_key::TotpKey; #[test] fn test_timing() { let key = TotpKey::new_random(); let code = key.current_code().unwrap(); let old_code = key.previous_code().unwrap(); let following_code = key.following_code().unwrap(); assert_ne!(code, old_code); assert_ne!(code, following_code); assert_ne!(old_code, following_code); } #[test] fn test_generation() { let key = TotpKey::from_encoded_secret("JBSWY3DPEHPK3PXP"); assert_eq!("124851", key.get_code_at(|| 1650470683).unwrap()); } }