Skip to content

Commit f0d792b

Browse files
Josh-XTJosh-XT
committed
Update login workflow
1 parent 01dba91 commit f0d792b

File tree

1 file changed

+263
-21
lines changed

1 file changed

+263
-21
lines changed

‎agixt-sdk/src/client/mod.rs‎

Lines changed: 263 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -50,11 +50,69 @@ impl AGiXTSDK {
5050

5151
// ==================== Authentication ====================
5252

53-
/// Login to the AGiXT server.
54-
pub async fn login(&self, email: &str, otp: &str) -> Result<Option<String>> {
53+
/// Login with username/password authentication.
54+
///
55+
/// # Arguments
56+
/// * `username` - Username or email address
57+
/// * `password` - User's password
58+
/// * `mfa_token` - Optional TOTP code if MFA is enabled
59+
///
60+
/// # Returns
61+
/// JWT token on success, or None on failure
62+
pub async fn login(&self, username: &str, password: &str, mfa_token: Option<&str>) -> Result<Option<String>> {
63+
let mut payload = serde_json::json!({
64+
"username": username,
65+
"password": password,
66+
});
67+
if let Some(token) = mfa_token {
68+
payload["mfa_token"] = serde_json::json!(token);
69+
}
70+
5571
let response = self
5672
.client
5773
.post(&format!("{}/v1/login", self.base_uri))
74+
.json(&payload)
75+
.send()
76+
.await?;
77+
78+
let status = response.status();
79+
let text = response.text().await?;
80+
81+
if self.verbose {
82+
self.parse_response(status, &text).await?;
83+
}
84+
85+
let json: serde_json::Value = serde_json::from_str(&text)?;
86+
87+
// Check for token in response (new auth flow)
88+
if status.is_success() {
89+
if let Some(token) = json.get("token").and_then(|t| t.as_str()) {
90+
let mut headers = self.headers.lock().await;
91+
if let Ok(value) = HeaderValue::from_str(token) {
92+
headers.insert(AUTHORIZATION, value);
93+
}
94+
if self.verbose {
95+
println!("Logged in successfully");
96+
}
97+
return Ok(Some(token.to_string()));
98+
}
99+
}
100+
Ok(None)
101+
}
102+
103+
/// Legacy login with magic link (email + OTP token).
104+
/// Maintained for backward compatibility.
105+
///
106+
/// # Arguments
107+
/// * `email` - User's email address
108+
/// * `otp` - TOTP code from authenticator app
109+
///
110+
/// # Returns
111+
/// JWT token on success, or None on failure
112+
pub async fn login_magic_link(&self, email: &str, otp: &str) -> Result<Option<String>> {
113+
let response = self
114+
.client
115+
.post(&format!("{}/v1/login/magic-link", self.base_uri))
58116
.json(&serde_json::json!({
59117
"email": email,
60118
"token": otp,
@@ -70,10 +128,7 @@ impl AGiXTSDK {
70128
}
71129

72130
let json: serde_json::Value = serde_json::from_str(&text)?;
73-
self.process_login_response(json).await
74-
}
75-
76-
async fn process_login_response(&self, json: serde_json::Value) -> Result<Option<String>> {
131+
77132
if let Some(detail) = json.get("detail").and_then(|d| d.as_str()) {
78133
if detail.contains("?token=") {
79134
let token = detail.split("token=").nth(1).unwrap_or_default();
@@ -88,15 +143,184 @@ impl AGiXTSDK {
88143
Ok(None)
89144
}
90145

91-
/// Register a new user.
92-
pub async fn register_user(&self, email: &str, first_name: &str, last_name: &str) -> Result<String> {
146+
/// Register a new user with username/password authentication.
147+
///
148+
/// # Arguments
149+
/// * `email` - User's email address
150+
/// * `password` - User's password
151+
/// * `confirm_password` - Password confirmation
152+
/// * `first_name` - User's first name (optional)
153+
/// * `last_name` - User's last name (optional)
154+
/// * `username` - Desired username (optional)
155+
/// * `organization_name` - Company/organization name (optional)
156+
///
157+
/// # Returns
158+
/// Response JSON with user_id, username, token on success
159+
pub async fn register_user(
160+
&self,
161+
email: &str,
162+
password: &str,
163+
confirm_password: &str,
164+
first_name: Option<&str>,
165+
last_name: Option<&str>,
166+
username: Option<&str>,
167+
organization_name: Option<&str>,
168+
) -> Result<serde_json::Value> {
169+
let mut payload = serde_json::json!({
170+
"email": email,
171+
"password": password,
172+
"confirm_password": confirm_password,
173+
"first_name": first_name.unwrap_or(""),
174+
"last_name": last_name.unwrap_or(""),
175+
});
176+
if let Some(u) = username {
177+
payload["username"] = serde_json::json!(u);
178+
}
179+
if let Some(org) = organization_name {
180+
payload["organization_name"] = serde_json::json!(org);
181+
}
182+
93183
let response = self
94184
.client
95185
.post(&format!("{}/v1/user", self.base_uri))
186+
.json(&payload)
187+
.send()
188+
.await?;
189+
190+
let status = response.status();
191+
let text = response.text().await?;
192+
193+
if self.verbose {
194+
self.parse_response(status, &text).await?;
195+
}
196+
197+
let json: serde_json::Value = serde_json::from_str(&text)?;
198+
199+
// Auto-login if token is returned
200+
if status.is_success() {
201+
if let Some(token) = json.get("token").and_then(|t| t.as_str()) {
202+
let mut headers = self.headers.lock().await;
203+
if let Ok(value) = HeaderValue::from_str(token) {
204+
headers.insert(AUTHORIZATION, value);
205+
}
206+
if self.verbose {
207+
println!("Registered and logged in as {}", json.get("username").and_then(|u| u.as_str()).unwrap_or(""));
208+
}
209+
}
210+
}
211+
212+
Ok(json)
213+
}
214+
215+
/// Get MFA setup information including QR code URI.
216+
///
217+
/// # Returns
218+
/// JSON with provisioning_uri, secret, and mfa_enabled status
219+
pub async fn get_mfa_setup(&self) -> Result<serde_json::Value> {
220+
let response = self
221+
.client
222+
.get(&format!("{}/v1/user/mfa/setup", self.base_uri))
223+
.headers(self.headers.lock().await.clone())
224+
.send()
225+
.await?;
226+
227+
let status = response.status();
228+
let text = response.text().await?;
229+
230+
if self.verbose {
231+
self.parse_response(status, &text).await?;
232+
}
233+
234+
let json: serde_json::Value = serde_json::from_str(&text)?;
235+
Ok(json)
236+
}
237+
238+
/// Enable MFA for the current user.
239+
///
240+
/// # Arguments
241+
/// * `mfa_token` - TOTP code from authenticator app to verify setup
242+
///
243+
/// # Returns
244+
/// Response JSON with success message
245+
pub async fn enable_mfa(&self, mfa_token: &str) -> Result<serde_json::Value> {
246+
let response = self
247+
.client
248+
.post(&format!("{}/v1/user/mfa/enable", self.base_uri))
249+
.headers(self.headers.lock().await.clone())
250+
.json(&serde_json::json!({ "mfa_token": mfa_token }))
251+
.send()
252+
.await?;
253+
254+
let status = response.status();
255+
let text = response.text().await?;
256+
257+
if self.verbose {
258+
self.parse_response(status, &text).await?;
259+
}
260+
261+
let json: serde_json::Value = serde_json::from_str(&text)?;
262+
Ok(json)
263+
}
264+
265+
/// Disable MFA for the current user.
266+
///
267+
/// # Arguments
268+
/// * `password` - User's password (optional)
269+
/// * `mfa_token` - Current TOTP code (optional)
270+
///
271+
/// # Returns
272+
/// Response JSON with success message
273+
pub async fn disable_mfa(&self, password: Option<&str>, mfa_token: Option<&str>) -> Result<serde_json::Value> {
274+
let mut payload = serde_json::json!({});
275+
if let Some(p) = password {
276+
payload["password"] = serde_json::json!(p);
277+
}
278+
if let Some(t) = mfa_token {
279+
payload["mfa_token"] = serde_json::json!(t);
280+
}
281+
282+
let response = self
283+
.client
284+
.post(&format!("{}/v1/user/mfa/disable", self.base_uri))
285+
.headers(self.headers.lock().await.clone())
286+
.json(&payload)
287+
.send()
288+
.await?;
289+
290+
let status = response.status();
291+
let text = response.text().await?;
292+
293+
if self.verbose {
294+
self.parse_response(status, &text).await?;
295+
}
296+
297+
let json: serde_json::Value = serde_json::from_str(&text)?;
298+
Ok(json)
299+
}
300+
301+
/// Change the current user's password.
302+
///
303+
/// # Arguments
304+
/// * `current_password` - Current password
305+
/// * `new_password` - New password
306+
/// * `confirm_password` - New password confirmation
307+
///
308+
/// # Returns
309+
/// Response JSON with success message
310+
pub async fn change_password(
311+
&self,
312+
current_password: &str,
313+
new_password: &str,
314+
confirm_password: &str,
315+
) -> Result<serde_json::Value> {
316+
let response = self
317+
.client
318+
.post(&format!("{}/v1/user/password/change", self.base_uri))
319+
.headers(self.headers.lock().await.clone())
96320
.json(&serde_json::json!({
97-
"email": email,
98-
"first_name": first_name,
99-
"last_name": last_name,
321+
"current_password": current_password,
322+
"new_password": new_password,
323+
"confirm_password": confirm_password,
100324
}))
101325
.send()
102326
.await?;
@@ -109,20 +333,38 @@ impl AGiXTSDK {
109333
}
110334

111335
let json: serde_json::Value = serde_json::from_str(&text)?;
336+
Ok(json)
337+
}
112338

113-
if let Some(otp_uri) = json.get("otp_uri").and_then(|u| u.as_str()) {
114-
let mfa_token = otp_uri
115-
.split("secret=")
116-
.nth(1)
117-
.and_then(|s| s.split('&').next())
118-
.ok_or_else(|| crate::Error::Other("Invalid OTP URI format".to_string()))?;
339+
/// Set a password for users who don't have one (e.g., social login users).
340+
///
341+
/// # Arguments
342+
/// * `new_password` - New password
343+
/// * `confirm_password` - New password confirmation
344+
///
345+
/// # Returns
346+
/// Response JSON with success message
347+
pub async fn set_password(&self, new_password: &str, confirm_password: &str) -> Result<serde_json::Value> {
348+
let response = self
349+
.client
350+
.post(&format!("{}/v1/user/password/set", self.base_uri))
351+
.headers(self.headers.lock().await.clone())
352+
.json(&serde_json::json!({
353+
"new_password": new_password,
354+
"confirm_password": confirm_password,
355+
}))
356+
.send()
357+
.await?;
119358

120-
self.login(email, mfa_token).await?;
359+
let status = response.status();
360+
let text = response.text().await?;
121361

122-
Ok(otp_uri.to_string())
123-
} else {
124-
Ok(text)
362+
if self.verbose {
363+
self.parse_response(status, &text).await?;
125364
}
365+
366+
let json: serde_json::Value = serde_json::from_str(&text)?;
367+
Ok(json)
126368
}
127369

128370
/// Check if a user exists.

0 commit comments

Comments
 (0)