reinhardt_testkit/auth/
secondary.rs1use async_trait::async_trait;
2
3use super::error::TestAuthError;
4use super::identity::SessionIdentity;
5use crate::client::APIClient;
6
7#[async_trait]
12pub trait SecondaryAuth: Send + Sync {
13 async fn apply_to_client(
17 &self,
18 client: &APIClient,
19 primary: &SessionIdentity,
20 ) -> Result<(), TestAuthError>;
21}
22
23#[cfg(native)]
28pub struct TotpSecondaryAuth {
29 manager: reinhardt_auth::mfa::MFAAuthentication,
30 code_override: Option<String>,
31}
32
33#[cfg(native)]
34impl TotpSecondaryAuth {
35 pub fn new(manager: reinhardt_auth::mfa::MFAAuthentication) -> Self {
37 Self {
38 manager,
39 code_override: None,
40 }
41 }
42
43 pub fn with_code_only(code: impl Into<String>) -> Self {
48 Self {
49 manager: reinhardt_auth::mfa::MFAAuthentication::new("test"),
50 code_override: Some(code.into()),
51 }
52 }
53
54 pub fn with_code(mut self, code: impl Into<String>) -> Self {
58 self.code_override = Some(code.into());
59 self
60 }
61}
62
63#[cfg(native)]
64#[async_trait]
65impl SecondaryAuth for TotpSecondaryAuth {
66 async fn apply_to_client(
67 &self,
68 client: &APIClient,
69 primary: &SessionIdentity,
70 ) -> Result<(), TestAuthError> {
71 let code = match &self.code_override {
72 Some(c) => c.clone(),
73 None => {
74 let secret = self
75 .manager
76 .get_secret(&primary.user_id)
77 .await
78 .ok_or_else(|| TestAuthError::MfaUserNotRegistered(primary.user_id.clone()))?;
79 generate_totp_code(&secret)
80 .map_err(|e| TestAuthError::SecondaryAuthError(e.to_string()))?
81 }
82 };
83 client
84 .set_header("X-MFA-Code", &code)
85 .await
86 .map_err(|e| TestAuthError::ClientError(e.to_string()))?;
87 Ok(())
88 }
89}
90
91#[cfg(native)]
94fn generate_totp_code(secret: &str) -> Result<String, String> {
95 use std::time::{SystemTime, UNIX_EPOCH};
96
97 let secret_bytes = base32_decode(secret).ok_or("invalid base32 secret")?;
98 let time = SystemTime::now()
99 .duration_since(UNIX_EPOCH)
100 .map_err(|e| e.to_string())?
101 .as_secs();
102 let step = time / 30;
103 Ok(totp_lite::totp_custom::<totp_lite::Sha256>(
104 30,
105 6,
106 &secret_bytes,
107 step,
108 ))
109}
110
111#[cfg(native)]
113fn base32_decode(input: &str) -> Option<Vec<u8>> {
114 const ALPHABET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";
115 let mut bits = 0u64;
116 let mut bit_count = 0u32;
117 let mut result = Vec::new();
118
119 for &byte in input.as_bytes() {
120 if byte == b'=' {
121 break;
122 }
123 let val = ALPHABET
124 .iter()
125 .position(|&c| c == byte.to_ascii_uppercase())? as u64;
126 bits = (bits << 5) | val;
127 bit_count += 5;
128 if bit_count >= 8 {
129 bit_count -= 8;
130 result.push((bits >> bit_count) as u8);
131 bits &= (1 << bit_count) - 1;
132 }
133 }
134 Some(result)
135}
136
137#[cfg(test)]
138mod tests {
139 #[cfg(native)]
140 mod totp_tests {
141 use super::super::*;
142 use rstest::*;
143
144 #[rstest]
145 fn base32_decode_known_value() {
146 let decoded = base32_decode("JBSWY3DPEHPK3PXP");
147 assert!(decoded.is_some());
148 let bytes = decoded.unwrap();
149 assert!(!bytes.is_empty());
150 }
151
152 #[rstest]
153 fn generate_totp_produces_six_digit_code() {
154 let code = generate_totp_code("JBSWY3DPEHPK3PXP").unwrap();
155 assert_eq!(code.len(), 6);
156 assert!(code.chars().all(|c| c.is_ascii_digit()));
157 }
158 }
159}