BasicOIDC/src/data/totp_key.rs

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 <https://github.com/google/google-authenticator/wiki/Key-Uri-Format>
    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<String> {
        self.get_code_at(time)
    }

    /// Get previous code
    pub fn previous_code(&self) -> Res<String> {
        self.get_code_at(|| time() - PERIOD)
    }

    /// Get following code
    pub fn following_code(&self) -> Res<String> {
        self.get_code_at(|| time() + PERIOD)
    }

    /// Get the code at a specific time
    fn get_code_at<F: Fn() -> u64>(&self, get_time: F) -> Res<String> {
        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<bool> {
        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());
    }
}