unifly-api 0.8.2

// Session API authentication
//
// Cookie-based session login/logout and controller platform detection.
// The login endpoint sets a session cookie in the client's jar;
// subsequent requests use that cookie automatically.
//
// MFA flow: when the controller requires 2FA, login returns HTTP 499
// with `{"errors":["ubic_2fa_token required"]}` and a `token` field
// containing a Set-Cookie value for `UBIC_2FA`. We inject that cookie,
// then retry the login POST with `ubic_2fa_token` in the body.

use secrecy::{ExposeSecret, SecretString};
use serde_json::json;
use tracing::debug;
use url::Url;

use crate::auth::ControllerPlatform;
use crate::error::Error;
use crate::session::client::SessionClient;

/// Response body from a 499 MFA challenge.
#[derive(serde::Deserialize)]
struct MfaChallengeResponse {
    /// Set-Cookie value for `UBIC_2FA` (e.g. `UBIC_2FA=eyJ...`).
    token: Option<String>,
}

impl SessionClient {
    /// Authenticate with the controller using username/password.
    ///
    /// If the controller has MFA enabled it returns HTTP 499. When a
    /// `totp_token` is provided, we complete the two-step challenge
    /// automatically. Without one, `Error::TwoFactorRequired` is returned
    /// so the caller can prompt the user.
    pub async fn login(
        &self,
        username: &str,
        password: &SecretString,
        totp_token: Option<&SecretString>,
    ) -> Result<(), Error> {
        let login_path = self
            .platform()
            .login_path()
            .ok_or_else(|| Error::Authentication {
                message: "login not supported on cloud platform".into(),
            })?;

        let url = self
            .base_url()
            .join(login_path)
            .map_err(Error::InvalidUrl)?;

        debug!("logging in at {}", url);

        let body = json!({
            "username": username,
            "password": password.expose_secret(),
        });

        let resp = self
            .http()
            .post(url.clone())
            .json(&body)
            .send()
            .await
            .map_err(Error::Transport)?;

        let status = resp.status();

        // HTTP 499 = MFA challenge from UniFi SSO
        if status.as_u16() == 499 {
            return self
                .handle_mfa_challenge(resp, &url, username, password, totp_token)
                .await;
        }

        if !status.is_success() {
            let body = resp.text().await.unwrap_or_default();
            return Err(Error::Authentication {
                message: format!("login failed (HTTP {status}): {body}"),
            });
        }

        self.capture_csrf(&resp);
        debug!("login successful");
        Ok(())
    }

    /// Handle a 499 MFA challenge: inject the UBIC_2FA cookie and retry
    /// with the TOTP token in the login body.
    async fn handle_mfa_challenge(
        &self,
        resp: reqwest::Response,
        login_url: &Url,
        username: &str,
        password: &SecretString,
        totp_token: Option<&SecretString>,
    ) -> Result<(), Error> {
        let totp = totp_token.ok_or(Error::TwoFactorRequired)?;
        let code = totp.expose_secret();

        // Validate TOTP format before sending
        if code.len() != 6 || !code.chars().all(|c| c.is_ascii_digit()) {
            return Err(Error::Authentication {
                message: "TOTP token must be exactly 6 digits".into(),
            });
        }

        debug!("MFA challenge received — completing 2FA handshake");

        // Extract the MFA cookie from the challenge response
        let challenge: MfaChallengeResponse =
            resp.json().await.map_err(|_| Error::Authentication {
                message: "failed to parse MFA challenge response".into(),
            })?;

        if let Some(cookie_value) = challenge.token {
            if !cookie_value.starts_with("UBIC_2FA=") {
                return Err(Error::Authentication {
                    message: format!(
                        "unexpected MFA cookie format (expected UBIC_2FA=...): {cookie_value}"
                    ),
                });
            }
            self.add_cookie(&cookie_value, login_url)?;
        }

        // Retry login with TOTP token
        let body = json!({
            "username": username,
            "password": password.expose_secret(),
            "ubic_2fa_token": code,
        });

        let resp = self
            .http()
            .post(login_url.clone())
            .json(&body)
            .send()
            .await
            .map_err(Error::Transport)?;

        let status = resp.status();
        if !status.is_success() {
            let body = resp.text().await.unwrap_or_default();
            return Err(Error::Authentication {
                message: format!("MFA login failed (HTTP {status}): {body}"),
            });
        }

        self.capture_csrf(&resp);
        debug!("MFA login successful");
        Ok(())
    }

    /// Extract and store CSRF token from a login response.
    fn capture_csrf(&self, resp: &reqwest::Response) {
        if let Some(token) = resp
            .headers()
            .get("X-CSRF-Token")
            .or_else(|| resp.headers().get("x-csrf-token"))
            .and_then(|v| v.to_str().ok())
        {
            self.set_csrf_token(token.to_owned());
        }
    }

    /// Attempt to restore a session from cache, falling back to fresh login.
    ///
    /// If a cached session is available and the validation probe succeeds,
    /// the cookie and CSRF token are restored without hitting the login
    /// endpoint. Otherwise, performs a normal login and caches the result.
    pub async fn login_with_cache(
        &self,
        username: &str,
        password: &secrecy::SecretString,
        totp_token: Option<&secrecy::SecretString>,
        cache: &super::session_cache::SessionCache,
    ) -> Result<(), Error> {
        use super::session_cache::{EXPIRY_MARGIN_SECS, fallback_expiry, jwt_expiry};

        // Try cached session first
        if let Some((cookie, csrf)) = cache.load() {
            self.add_cookie(&cookie, self.base_url())?;
            if let Some(token) = csrf {
                self.set_csrf_token(token);
            }

            if self.validate_session().await {
                debug!("restored session from cache");
                return Ok(());
            }
            debug!("cached session invalid, performing fresh login");
        }

        // Fresh login
        self.login(username, password, totp_token).await?;

        // Cache the new session
        if let Some(cookie) = self.cookie_header() {
            let csrf = self.csrf_token_value();
            let expires_at = jwt_expiry(&cookie).map_or_else(fallback_expiry, |exp| {
                exp.saturating_sub(EXPIRY_MARGIN_SECS)
            });
            cache.save(&cookie, csrf.as_deref(), expires_at);
        }

        Ok(())
    }

    /// Validate the current session by probing a lightweight endpoint.
    ///
    /// Returns `true` if the session is still alive.
    async fn validate_session(&self) -> bool {
        let prefix = self.platform().session_prefix().unwrap_or("");
        let base = self.base_url().as_str().trim_end_matches('/');
        let prefix = prefix.trim_end_matches('/');
        let probe_url = format!("{base}{prefix}/api/s/{}/self", self.site());

        let Ok(url) = url::Url::parse(&probe_url) else {
            return false;
        };

        match self.http().get(url).send().await {
            Ok(resp) => resp.status().is_success(),
            Err(_) => false,
        }
    }

    /// End the current session.
    ///
    /// Platform-specific logout endpoint:
    /// - UniFi OS: `POST /api/auth/logout`
    /// - Standalone: `POST /api/logout`
    pub async fn logout(&self) -> Result<(), Error> {
        let logout_path = self
            .platform()
            .logout_path()
            .ok_or_else(|| Error::Authentication {
                message: "logout not supported on cloud platform".into(),
            })?;

        let url = self
            .base_url()
            .join(logout_path)
            .map_err(Error::InvalidUrl)?;

        debug!("logging out at {}", url);

        let _resp = self
            .http()
            .post(url)
            .send()
            .await
            .map_err(Error::Transport)?;

        debug!("logout complete");
        Ok(())
    }

    /// Auto-detect the controller platform by probing login endpoints.
    ///
    /// Tries the UniFi OS endpoint first (`/api/auth/login`). If it
    /// responds (even with an error), we're on UniFi OS. If the connection
    /// fails or returns 404, falls back to standalone detection.
    pub async fn detect_platform(base_url: &Url) -> Result<ControllerPlatform, Error> {
        let http = reqwest::Client::builder()
            .danger_accept_invalid_certs(true)
            .build()
            .map_err(Error::Transport)?;

        // Probe UniFi OS endpoint
        let unifi_os_url = base_url
            .join("/api/auth/login")
            .map_err(Error::InvalidUrl)?;

        debug!("probing UniFi OS at {}", unifi_os_url);

        if let Ok(resp) = http.get(unifi_os_url).send().await {
            // UniFi OS returns a response (even 401/405) at this path.
            // Standalone controllers don't have this endpoint at all.
            if resp.status() != reqwest::StatusCode::NOT_FOUND {
                debug!("detected UniFi OS platform");
                return Ok(ControllerPlatform::UnifiOs);
            }
        }
        // Connection error -- might be standalone on a different port

        // Probe standalone endpoint
        let standalone_url = base_url.join("/api/login").map_err(Error::InvalidUrl)?;

        debug!("probing standalone at {}", standalone_url);

        match http.get(standalone_url).send().await {
            Ok(_) => {
                debug!("detected standalone (classic) controller");
                Ok(ControllerPlatform::ClassicController)
            }
            Err(e) => Err(Error::Transport(e)),
        }
    }
}