laron-crypto 0.2.3

// This file is part of the laron-crypto
//
// Copyright 2023 Ade M Ramdani
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use constant_time_eq::constant_time_eq;
use hmac::{Hmac, Mac};
use rand::RngCore;
use std::time::SystemTime;

/// Error type for TOTP.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Error {
    /// System time error.
    SystemTimeError(String),
    /// Base32 decode error.
    Base32DecodeError,
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match &self {
            Error::SystemTimeError(e) => write!(f, "System time error: {}", e),
            Error::Base32DecodeError => write!(f, "Base32 decode error"),
        }
    }
}

impl std::error::Error for Error {}

/// Result type for TOTP.
pub type Result<T> = std::result::Result<T, Error>;

fn system_time() -> Result<u64> {
    SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .map(|d| d.as_secs())
        .map_err(|e| Error::SystemTimeError(e.to_string()))
}

/// Algorithm supported for TOTP.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Algorithm {
    /// HMAC-based One-time Password Algorithm.
    HmacSha1,
    /// HMAC-based One-time Password Algorithm.
    HmacSha256,
    /// HMAC-based One-time Password Algorithm.
    HmacSha512,
}

impl Default for Algorithm {
    fn default() -> Self {
        Algorithm::HmacSha1
    }
}

impl Algorithm {
    fn digest<M: Mac>(mut mac: M, data: &[u8]) -> Vec<u8> {
        mac.update(data);
        mac.finalize().into_bytes().to_vec()
    }

    fn sign(&self, key: &[u8], data: &[u8]) -> Vec<u8> {
        match self {
            Algorithm::HmacSha1 => {
                Algorithm::digest(Hmac::<sha1::Sha1>::new_from_slice(key).unwrap(), data)
            }
            Algorithm::HmacSha256 => {
                Algorithm::digest(Hmac::<sha2::Sha256>::new_from_slice(key).unwrap(), data)
            }
            Algorithm::HmacSha512 => {
                Algorithm::digest(Hmac::<sha2::Sha512>::new_from_slice(key).unwrap(), data)
            }
        }
    }
}

/// Secret is a shared secret key for TOTP.
#[derive(Debug, Clone)]
pub struct Secret {
    /// The shared secret key.
    key: Vec<u8>,
}

impl PartialEq for Secret {
    fn eq(&self, other: &Self) -> bool {
        constant_time_eq(&self.key, &other.key)
    }
}

impl Eq for Secret {}

impl Secret {
    /// Create a new Secret instance.
    pub fn new() -> Self {
        let mut key = vec![0u8; 16];
        rand::thread_rng().fill_bytes(&mut key);
        Self { key }
    }

    /// Create a new Secret instance from a base32 encoded string.
    pub fn from_base32(key: &str) -> Result<Self> {
        let key = base32::decode(base32::Alphabet::RFC4648 { padding: false }, key);
        if key.is_none() {
            return Err(Error::Base32DecodeError);
        }
        Ok(Self { key: key.unwrap() })
    }

    /// Create a new Secret instance from raw secret key.
    pub fn from_raw(key: &str) -> Self {
        Self {
            key: key.as_bytes().to_vec(),
        }
    }

    /// Create a new Secret instance from slice.
    pub fn from_slice(key: &[u8]) -> Self {
        Self { key: key.to_vec() }
    }

    /// Get the raw secret key.
    pub fn raw(&self) -> String {
        String::from_utf8(self.key.clone()).unwrap()
    }

    /// Get the base32 encoded secret key.
    pub fn base32(&self) -> String {
        base32::encode(base32::Alphabet::RFC4648 { padding: false }, &self.key)
    }

    /// Get the slice of the secret key.
    pub fn as_slice(&self) -> &[u8] {
        &self.key
    }
}

/// TOTP is a time-based one-time password algorithm.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TOTP {
    /// Algorithm used for TOTP.
    algorithm: Algorithm,
    /// The number of digits in the generated code.
    digits: usize,
    /// Number of steps allowed as a delay in the counter.
    window: u64,
    /// The time step in seconds.
    time_step: u64,
    /// The shared secret key.
    key: Secret,
}

impl TOTP {
    /// Create a new TOTP instance.
    pub fn new(
        algorithm: Algorithm,
        digits: usize,
        window: u64,
        time_step: u64,
        key: Secret,
    ) -> Self {
        Self::new_unchecked(algorithm, digits, window, time_step, key)
    }

    /// Create a new TOTP instance without checking the validity of the parameters.
    pub fn new_unchecked(
        algorithm: Algorithm,
        digits: usize,
        window: u64,
        time_step: u64,
        key: Secret,
    ) -> Self {
        Self {
            algorithm,
            digits,
            window,
            time_step,
            key,
        }
    }

    /// Generate a TOTP code.
    pub fn generate(&self, time: u64) -> String {
        let hash = self.algorithm.sign(
            self.key.as_slice(),
            (time / self.time_step).to_be_bytes().as_ref(),
        );
        let offset = (hash.last().unwrap() & 0xf) as usize;
        let result = u32::from_be_bytes(hash[offset..offset + 4].try_into().unwrap()) & 0x7fffffff;
        format!(
            "{1:00$}",
            self.digits,
            result % 10_u32.pow(self.digits as u32)
        )
    }

    /// Generate a TOTP code using the current time.
    pub fn generate_now(&self) -> Result<String> {
        let time = system_time()?;
        Ok(self.generate(time))
    }

    /// Returns the timestamp of the first second for the next time step.
    pub fn next_time_step(&self, time: u64) -> u64 {
        let time_step = time / self.time_step;
        (time_step + 1) * self.time_step
    }

    /// Get the time to live for the current token.
    pub fn ttl(&self) -> Result<u64> {
        let time = system_time()?;
        let remain = self.time_step - (time % self.time_step);
        Ok(remain)
    }

    /// Verify a TOTP code.
    pub fn verify(&self, token: &str) -> Result<bool> {
        let time = system_time()?;
        Ok(self.verify_with_time(token, time))
    }

    /// Verify a TOTP code using the given time.
    pub fn verify_with_time(&self, token: &str, time: u64) -> bool {
        let step = time / self.time_step - self.window;
        for i in 0..self.window * 2 + 1 {
            let t = (step + i) * self.time_step;
            let code = self.generate(t);
            if constant_time_eq(token.as_bytes(), code.as_bytes()) {
                return true;
            }
        }
        false
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_totp() {
        let secret =
            Secret::from_base32("KLOKQRZDX46UDXHSOUUM32PKUPPKGI7LM4U3G4GEDVUTSNXLTLNA").unwrap();
        let totp = TOTP::new(Algorithm::HmacSha1, 6, 1, 30, secret);
        let code = totp.generate_now().unwrap();
        assert_eq!(code.len(), 6);
        assert!(totp.verify(&code).unwrap());
    }
}