Skip to main content

wae_authentication/totp/
secret.rs

1//! TOTP 密钥管理
2
3use crate::totp::{TotpError, TotpResult};
4use base64::{Engine, engine::general_purpose::STANDARD as BASE64_STANDARD};
5use rand::Rng;
6
7/// Base32 字符集
8const BASE32_CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";
9
10/// TOTP 密钥
11#[derive(Debug, Clone)]
12pub struct TotpSecret {
13    /// 原始字节
14    bytes: Vec<u8>,
15
16    /// Base32 编码
17    base32: String,
18}
19
20impl TotpSecret {
21    /// 生成新的随机密钥
22    ///
23    /// # Arguments
24    /// * `length` - 密钥长度 (字节),推荐 20 字节 (160 位)
25    pub fn generate(length: usize) -> TotpResult<Self> {
26        let mut bytes = vec![0u8; length];
27        rand::rng().fill_bytes(&mut bytes);
28        Self::from_bytes(&bytes)
29    }
30
31    /// 生成推荐长度的密钥 (20 字节)
32    pub fn generate_default() -> TotpResult<Self> {
33        Self::generate(20)
34    }
35
36    /// 从字节创建密钥
37    pub fn from_bytes(bytes: &[u8]) -> TotpResult<Self> {
38        let base32 = Self::encode_base32(bytes);
39        Ok(Self { bytes: bytes.to_vec(), base32 })
40    }
41
42    /// 从 Base32 字符串创建密钥
43    pub fn from_base32(s: &str) -> TotpResult<Self> {
44        let bytes = Self::decode_base32(s)?;
45        Ok(Self { bytes, base32: s.to_uppercase().replace(" ", "") })
46    }
47
48    /// 获取原始字节
49    pub fn as_bytes(&self) -> &[u8] {
50        &self.bytes
51    }
52
53    /// 获取 Base32 编码
54    pub fn as_base32(&self) -> &str {
55        &self.base32
56    }
57
58    /// 获取密钥长度 (字节)
59    pub fn len(&self) -> usize {
60        self.bytes.len()
61    }
62
63    /// 检查密钥是否为空
64    pub fn is_empty(&self) -> bool {
65        self.bytes.is_empty()
66    }
67
68    /// 编码为 Base32
69    fn encode_base32(data: &[u8]) -> String {
70        let mut result = String::new();
71        let mut i = 0;
72        let n = data.len();
73
74        while i < n {
75            let mut word: u64 = 0;
76            let mut bits = 0;
77
78            for j in 0..5 {
79                if i + j < n {
80                    word = (word << 8) | (data[i + j] as u64);
81                    bits += 8;
82                }
83            }
84
85            i += 5;
86
87            while bits >= 5 {
88                bits -= 5;
89                let index = ((word >> bits) & 0x1F) as usize;
90                result.push(BASE32_CHARS[index] as char);
91            }
92
93            if bits > 0 {
94                let index = ((word << (5 - bits)) & 0x1F) as usize;
95                result.push(BASE32_CHARS[index] as char);
96            }
97        }
98
99        result
100    }
101
102    /// 解码 Base32
103    fn decode_base32(s: &str) -> TotpResult<Vec<u8>> {
104        let s = s.to_uppercase().replace(" ", "").replace("-", "");
105        let mut result = Vec::new();
106        let chars: Vec<char> = s.chars().collect();
107
108        let mut i = 0;
109        while i < chars.len() {
110            let mut word: u64 = 0;
111            let mut bits = 0;
112
113            for j in 0..8 {
114                if i + j < chars.len() {
115                    let val = Self::base32_char_to_value(chars[i + j])?;
116                    word = (word << 5) | (val as u64);
117                    bits += 5;
118                }
119            }
120
121            i += 8;
122
123            while bits >= 8 {
124                bits -= 8;
125                result.push(((word >> bits) & 0xFF) as u8);
126            }
127        }
128
129        Ok(result)
130    }
131
132    /// Base32 字符转值
133    fn base32_char_to_value(c: char) -> TotpResult<u8> {
134        match c {
135            'A'..='Z' => Ok((c as u8) - b'A'),
136            '2'..='7' => Ok((c as u8) - b'2' + 26),
137            _ => Err(TotpError::Base32Error(format!("Invalid character: {}", c))),
138        }
139    }
140}
141
142impl std::fmt::Display for TotpSecret {
143    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
144        write!(f, "{}", self.base32)
145    }
146}
147
148/// 密钥格式化选项
149#[derive(Debug, Clone, Copy, PartialEq, Eq)]
150pub enum SecretFormat {
151    /// Base32 格式 (无分隔符)
152    Base32,
153    /// Base32 格式 (每 4 字符空格分隔)
154    Base32Spaced,
155    /// 原始字节
156    Raw,
157    /// Base64 格式
158    Base64,
159}
160
161impl TotpSecret {
162    /// 格式化密钥
163    pub fn format(&self, format: SecretFormat) -> String {
164        match format {
165            SecretFormat::Base32 => self.base32.clone(),
166            SecretFormat::Base32Spaced => self
167                .base32
168                .as_bytes()
169                .chunks(4)
170                .map(|chunk| std::str::from_utf8(chunk).unwrap_or(""))
171                .collect::<Vec<_>>()
172                .join(" "),
173            SecretFormat::Raw => self.bytes.iter().map(|b| format!("{:02x}", b)).collect(),
174            SecretFormat::Base64 => BASE64_STANDARD.encode(&self.bytes),
175        }
176    }
177}