1use hmac::{Hmac, Mac};
2
3pub enum Algorithm {
4 Sha1,
5 Sha256,
6 Sha512,
7 Steam,
8}
9
10const STEAM_CHARS: &[u8] = b"23456789BCDFGHJKMNPQRTVWXY";
11
12pub fn decode_base32(input: &str) -> Option<Vec<u8>> {
13 let cleaned: Vec<u8> = input
14 .trim()
15 .bytes()
16 .filter(|b| !b.is_ascii_whitespace())
17 .collect();
18 let trimmed: &[u8] = {
19 let mut end = cleaned.len();
20 while end > 0 && cleaned[end - 1] == b'=' {
21 end -= 1;
22 }
23 &cleaned[..end]
24 };
25 let mut out = Vec::with_capacity(trimmed.len() * 5 / 8);
26 let mut buffer: u16 = 0;
27 let mut bits: u8 = 0;
28 for &b in trimmed {
29 let v: u16 = match b {
30 b'A'..=b'Z' => u16::from(b - b'A'),
31 b'a'..=b'z' => u16::from(b - b'a'),
32 b'2'..=b'7' => u16::from(b - b'2') + 26,
33 _ => return None,
34 };
35 buffer = (buffer << 5) | v;
36 bits += 5;
37 if bits >= 8 {
38 bits -= 8;
39 let byte = u8::try_from((buffer >> bits) & 0xff).ok()?;
40 out.push(byte);
41 buffer &= (1_u16 << bits).wrapping_sub(1);
42 }
43 }
44 Some(out)
45}
46
47pub fn generate(
48 secret: &[u8],
49 unix_secs: u64,
50 step: u64,
51 digits: u32,
52 algorithm: &Algorithm,
53) -> crate::error::Result<String> {
54 let totp_err = |msg: &str| crate::error::Error::Totp {
55 msg: msg.to_string(),
56 };
57 let counter = (unix_secs / step).to_be_bytes();
58 let mac: Vec<u8> =
59 match algorithm {
60 Algorithm::Sha1 | Algorithm::Steam => {
61 let mut m = Hmac::<sha1::Sha1>::new_from_slice(secret)
62 .map_err(|e| crate::error::Error::Totp {
63 msg: format!("invalid hmac key: {e}"),
64 })?;
65 m.update(&counter);
66 m.finalize().into_bytes().to_vec()
67 }
68 Algorithm::Sha256 => {
69 let mut m = Hmac::<sha2::Sha256>::new_from_slice(secret)
70 .map_err(|e| crate::error::Error::Totp {
71 msg: format!("invalid hmac key: {e}"),
72 })?;
73 m.update(&counter);
74 m.finalize().into_bytes().to_vec()
75 }
76 Algorithm::Sha512 => {
77 let mut m = Hmac::<sha2::Sha512>::new_from_slice(secret)
78 .map_err(|e| crate::error::Error::Totp {
79 msg: format!("invalid hmac key: {e}"),
80 })?;
81 m.update(&counter);
82 m.finalize().into_bytes().to_vec()
83 }
84 };
85
86 let offset = usize::from(
87 *mac.last().ok_or_else(|| totp_err("empty hmac output"))? & 0x0f,
88 );
89 let mut truncated = u32::from_be_bytes(
90 mac[offset..offset + 4]
91 .try_into()
92 .map_err(|_| totp_err("totp truncation failed"))?,
93 ) & 0x7fff_ffff;
94
95 let digits_usize = usize::try_from(digits)
96 .map_err(|_| totp_err("digits out of range"))?;
97
98 match algorithm {
99 Algorithm::Sha1 | Algorithm::Sha256 | Algorithm::Sha512 => {
100 let modulus = 10_u32
101 .checked_pow(digits)
102 .ok_or_else(|| totp_err("digits too large"))?;
103 Ok(format!(
104 "{:0width$}",
105 truncated % modulus,
106 width = digits_usize
107 ))
108 }
109 Algorithm::Steam => {
110 let len = u32::try_from(STEAM_CHARS.len())
111 .map_err(|_| totp_err("steam alphabet too large"))?;
112 let mut s = String::with_capacity(digits_usize);
113 for _ in 0..digits {
114 let idx = usize::try_from(truncated % len)
115 .map_err(|_| totp_err("steam index error"))?;
116 s.push(char::from(STEAM_CHARS[idx]));
117 truncated /= len;
118 }
119 Ok(s)
120 }
121 }
122}
123
124#[cfg(test)]
125#[allow(clippy::as_conversions, clippy::cast_possible_truncation)]
126mod test {
127 use super::{decode_base32, generate, Algorithm, STEAM_CHARS};
128
129 #[test]
130 fn test_decode_base32_basic() {
131 assert_eq!(decode_base32("MZXW6===").unwrap(), b"foo");
132 assert_eq!(decode_base32("MZXW6").unwrap(), b"foo");
133 assert_eq!(decode_base32("mzxw6").unwrap(), b"foo");
134 assert_eq!(decode_base32("MZ XW 6").unwrap(), b"foo");
135 assert_eq!(decode_base32("MZXW6YQ=").unwrap(), b"foob");
136 assert_eq!(decode_base32("MZXW6YTB").unwrap(), b"fooba");
137 assert_eq!(decode_base32("MZXW6YTBOI======").unwrap(), b"foobar");
138 assert!(decode_base32("!!!").is_none());
139 }
140
141 #[test]
142 fn test_rfc6238_sha1() {
143 let secret = b"12345678901234567890";
144 let cases = [
145 (59_u64, "94287082"),
146 (1_111_111_109, "07081804"),
147 (1_111_111_111, "14050471"),
148 (1_234_567_890, "89005924"),
149 (2_000_000_000, "69279037"),
150 (20_000_000_000, "65353130"),
151 ];
152 for (t, expected) in cases {
153 let code = generate(secret, t, 30, 8, &Algorithm::Sha1).unwrap();
154 assert_eq!(code, expected, "time={t}");
155 }
156 }
157
158 #[test]
159 fn test_rfc6238_sha256() {
160 let secret = b"12345678901234567890123456789012";
161 let code = generate(secret, 59, 30, 8, &Algorithm::Sha256).unwrap();
162 assert_eq!(code, "46119246");
163 }
164
165 #[test]
166 fn test_rfc6238_sha512() {
167 let secret =
168 b"1234567890123456789012345678901234567890123456789012345678901234";
169 let code = generate(secret, 59, 30, 8, &Algorithm::Sha512).unwrap();
170 assert_eq!(code, "90693936");
171 }
172
173 #[test]
174 fn test_digits_six() {
175 let secret = b"12345678901234567890";
176 let code = generate(secret, 59, 30, 6, &Algorithm::Sha1).unwrap();
177 assert_eq!(code, "287082");
178 }
179
180 #[test]
181 fn test_steam() {
182 let secret = decode_base32("STEAMKEY234567").unwrap();
183 let code = generate(&secret, 1_000_000_000, 30, 5, &Algorithm::Steam)
184 .unwrap();
185 assert_eq!(code.len(), 5);
186 for c in code.chars() {
187 assert!(STEAM_CHARS.contains(&(c as u8)));
188 }
189 let pinned =
190 generate(b"12345678901234567890", 59, 30, 5, &Algorithm::Steam)
191 .unwrap();
192 assert_eq!(pinned, "PV9M4");
193 }
194}