#[cfg(feature = "serde_support")]
use serde::{Deserialize, Serialize};
use byteorder::{BigEndian, ReadBytesExt};
use std::io::Cursor;
#[cfg(feature = "qr")]
use {base64, image::Luma, qrcode::QrCode};
use hmac::{Hmac, Mac, NewMac};
use sha1::Sha1;
use sha2::{Sha256, Sha512};
type HmacSha1 = Hmac<Sha1>;
type HmacSha256 = Hmac<Sha256>;
type HmacSha512 = Hmac<Sha512>;
#[derive(Debug, Copy, Clone)]
#[cfg_attr(feature = "serde_support", derive(Serialize, Deserialize))]
pub enum Algorithm {
SHA1,
SHA256,
SHA512,
}
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde_support", derive(Serialize, Deserialize))]
pub struct TOTP<T = Vec<u8>> {
pub algorithm: Algorithm,
pub digits: usize,
pub skew: u8,
pub step: u64,
pub secret: T,
}
impl<T: AsRef<[u8]>> TOTP<T> {
pub fn new(algorithm: Algorithm, digits: usize, skew: u8, step: u64, secret: T) -> TOTP<T> {
TOTP {
algorithm,
digits,
skew,
step,
secret,
}
}
pub fn sign(&self, time: u64) -> Vec<u8> {
let ctr = (time / self.step).to_be_bytes().to_vec();
match self.algorithm {
Algorithm::SHA1 => {
let mut mac = HmacSha1::new_varkey(self.secret.as_ref()).expect("no key");
mac.update(&ctr);
mac.finalize().into_bytes().to_vec()
}
Algorithm::SHA256 => {
let mut mac = HmacSha256::new_varkey(self.secret.as_ref()).expect("no key");
mac.update(&ctr);
mac.finalize().into_bytes().to_vec()
}
Algorithm::SHA512 => {
let mut mac = HmacSha512::new_varkey(self.secret.as_ref()).expect("no key");
mac.update(&ctr);
mac.finalize().into_bytes().to_vec()
}
}
}
pub fn generate(&self, time: u64) -> String {
let result: &[u8] = &self.sign(time);
let offset = (result[19] & 15) as usize;
let mut rdr = Cursor::new(result[offset..offset + 4].to_vec());
let result = rdr.read_u32::<BigEndian>().unwrap() & 0x7fff_ffff;
format!(
"{1:00$}",
self.digits,
result % (10 as u32).pow(self.digits as u32)
)
}
pub fn check(&self, token: &str, time: u64) -> bool {
let basestep = time / self.step - (self.skew as u64);
for i in 0..self.skew * 2 + 1 {
let step_time = (basestep + (i as u64)) * (self.step as u64);
if self.generate(step_time) == token {
return true;
}
}
false
}
pub fn get_secret_base32(&self) -> String {
base32::encode(
base32::Alphabet::RFC4648 { padding: false },
self.secret.as_ref(),
)
}
pub fn get_url(&self, label: &str, issuer: &str) -> String {
let algorithm = match self.algorithm {
Algorithm::SHA1 => "SHA1",
Algorithm::SHA256 => "SHA256",
Algorithm::SHA512 => "SHA512",
};
format!(
"otpauth://totp/{}?secret={}&issuer={}&digits={}&algorithm={}",
label,
self.get_secret_base32(),
issuer,
self.digits,
algorithm,
)
}
#[cfg(feature = "qr")]
pub fn get_qr(&self, label: &str, issuer: &str) -> Result<String, Box<dyn std::error::Error>> {
let url = self.get_url(label, issuer);
let code = QrCode::new(&url)?;
let mut vec = Vec::new();
let size: u32 = ((code.width() + 8) * 8) as u32;
let encoder = image::png::PNGEncoder::new(&mut vec);
encoder.encode(
&code.render::<Luma<u8>>().build().to_vec(),
size,
size,
image::ColorType::L8,
)?;
Ok(base64::encode(vec))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn url_for_secret_matches() {
let totp = TOTP::new(Algorithm::SHA1, 6, 1, 1, "TestSecret");
let url = totp.get_url("test_url", "totp-rs");
assert_eq!(url.as_str(), "otpauth://totp/test_url?secret=KRSXG5CTMVRXEZLU&issuer=totp-rs&digits=6&algorithm=SHA1");
}
#[test]
fn returns_base32() {
let totp = TOTP::new(Algorithm::SHA1, 6, 1, 1, "TestSecret");
assert_eq!(totp.get_secret_base32().as_str(), "KRSXG5CTMVRXEZLU");
}
#[test]
fn generates_token() {
let totp = TOTP::new(Algorithm::SHA1, 6, 1, 1, "TestSecret");
assert_eq!(totp.generate(1000).as_str(), "718996");
}
#[test]
fn generates_token_sha256() {
let totp = TOTP::new(Algorithm::SHA256, 6, 1, 1, "TestSecret");
assert_eq!(totp.generate(1000).as_str(), "423657");
}
#[test]
fn generates_token_sha512() {
let totp = TOTP::new(Algorithm::SHA512, 6, 1, 1, "TestSecret");
assert_eq!(totp.generate(1000).as_str(), "416767");
}
#[test]
fn checks_token() {
let totp = TOTP::new(Algorithm::SHA1, 6, 1, 1, "TestSecret");
assert!(totp.check("718996", 1000));
}
#[test]
fn checks_token_with_skew() {
let totp = TOTP::new(Algorithm::SHA1, 6, 1, 1, "TestSecret");
assert!(
totp.check("527544", 2000) && totp.check("712039", 2000) && totp.check("714250", 2000)
);
}
#[test]
#[cfg(feature = "qr")]
fn generates_qr() {
use sha1::{Digest, Sha1};
let totp = TOTP::new(Algorithm::SHA1, 6, 1, 1, "TestSecret");
let qr = totp.get_qr("test_url", "totp-rs").unwrap();
let hash_digest = Sha1::digest(qr.as_bytes());
assert_eq!(
format!("{:x}", hash_digest).as_str(),
"3abc0127e7a2b1013fb25c97ef14422c1fe9e878"
);
}
}