pub use humantime::FormattedDuration;
use sha1::{Digest, Sha1};
use totp_rs::{Algorithm, Secret, TOTP as TotpStd};
#[derive(Debug, Clone)]
pub struct TOTP {
algorithm: Algorithm,
digits: usize,
skew: u8,
step: u64,
}
impl Default for TOTP {
fn default() -> Self {
Self {
algorithm: Algorithm::SHA1,
digits: 6,
skew: 1,
step: Self::OTP_DEFAULT_STEP,
}
}
}
impl TOTP {
const OTP_DEFAULT_STEP: u64 = 30;
pub fn default_otp() -> Self {
Self::default()
}
const EMAIL_DEFAULT_STEP: u64 = 10 * 60;
pub fn default_email() -> Self {
Self::default().with_step(Self::EMAIL_DEFAULT_STEP)
}
const SMS_DEFAULT_STEP: u64 = 5 * 60;
pub fn default_sms() -> Self {
Self::default().with_step(Self::SMS_DEFAULT_STEP)
}
pub const fn with_step(mut self, step: u64) -> Self {
self.step = step;
self
}
pub const fn with_digits(mut self, digits: usize) -> Self {
self.digits = digits;
self
}
pub fn secret_from_root(root: impl AsRef<[u8]>) -> String {
let mut hasher = Sha1::new();
hasher.update(root);
let result = hasher.finalize();
let Secret::Encoded(secret_base32) = Secret::Raw(result.to_vec()).to_encoded() else {
unreachable!()
};
secret_base32
}
pub fn gen_secret_base32() -> String {
let Secret::Encoded(secret_base32) = Secret::generate_secret().to_encoded() else {
unreachable!()
};
secret_base32
}
pub fn qr_base64(&self, secret_base32: String, app_name: String, username: String) -> anyhow::Result<String> {
let secret = Secret::Encoded(secret_base32);
let totp = TotpStd::new(
self.algorithm,
self.digits,
self.skew,
self.step,
secret.to_bytes()?,
Some(app_name),
username,
)?;
let qr_base64 = totp.get_qr_base64().map_err(|err| anyhow::anyhow!(err))?;
Ok(format!("data:image/png;base64,{qr_base64}"))
}
pub fn qr_png(&self, secret_base32: String, app_name: String, username: String) -> anyhow::Result<Vec<u8>> {
let secret = Secret::Encoded(secret_base32);
let totp = TotpStd::new(
self.algorithm,
self.digits,
self.skew,
self.step,
secret.to_bytes()?,
Some(app_name),
username,
)?;
let qr_png = totp.get_qr_png().map_err(|err| anyhow::anyhow!(err))?;
Ok(qr_png)
}
pub fn check_current(&self, secret_base32: String, totp_token: &str) -> anyhow::Result<bool> {
let totp = TotpStd::new(
self.algorithm,
self.digits,
self.skew,
self.step,
Secret::Encoded(secret_base32).to_bytes()?,
None,
"".to_owned(),
)?;
Ok(totp.check_current(totp_token)?)
}
pub fn gen_current(&self, secret_base32: String) -> anyhow::Result<String> {
let totp = TotpStd::new(
self.algorithm,
self.digits,
self.skew,
self.step,
Secret::Encoded(secret_base32).to_bytes()?,
None,
"".to_owned(),
)?;
Ok(totp.generate_current()?)
}
pub fn gen_current2(&self, secret_base32: String) -> anyhow::Result<(String, FormattedDuration)> {
let code = self.gen_current(secret_base32)?;
let duration = humantime::format_duration(core::time::Duration::from_secs(self.step));
Ok((code, duration))
}
}
#[cfg(test)]
mod tests {
use super::TOTP;
const SECRET_BASE32: &'static str = "BQ3OBNNHYUZJZFTF6R2XBTKY2OJLISPW";
#[test]
fn gen_secret_base32() {
let secret_base32 = TOTP::gen_secret_base32();
println!("secret_base32={secret_base32}");
}
#[test]
fn qr_base64() {
let totp = TOTP::default();
let qr_base64 = totp
.qr_base64(
SECRET_BASE32.to_owned(),
env!("CARGO_PKG_NAME").to_owned(),
"andeya".to_owned(),
)
.unwrap();
println!("{qr_base64}");
}
#[test]
fn check_current() {
let totp_token = "379244";
let totp = TOTP::default();
let r = totp.check_current(SECRET_BASE32.to_owned(), totp_token);
println!("{r:?}");
}
#[test]
fn gen_current() {
let totp = TOTP::default();
let totp_token = totp.gen_current(SECRET_BASE32.to_owned());
println!("{totp_token:?}");
}
#[test]
fn gen_current2() {
let totp = TOTP::default();
let (token, fmt_duration) = totp.gen_current2(SECRET_BASE32.to_owned()).unwrap();
println!("The TOTP code {token} is valid within {fmt_duration}.");
}
}