unifly-api 0.9.0

// 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.
    ///
    /// Recent standalone controllers can return a misleading 401 from
    /// `/api/auth/login`, so we need both probes before deciding.
    ///
    /// Observed patterns:
    /// - Classic standalone: `/api/login` responds with 200/400, while
    ///   `/api/auth/login` is usually 404 and may occasionally be 401.
    /// - UniFi OS: both `/api/login` and `/api/auth/login` can return 401.
    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)?;

        let standalone_url = base_url.join("/api/login").map_err(Error::InvalidUrl)?;
        debug!("probing standalone at {}", standalone_url);
        let unifi_os_url = base_url
            .join("/api/auth/login")
            .map_err(Error::InvalidUrl)?;
        debug!("probing UniFi OS at {}", unifi_os_url);
        let standalone_status = http
            .get(standalone_url)
            .send()
            .await
            .map(|resp| resp.status());
        let unifi_os_status = http
            .get(unifi_os_url)
            .send()
            .await
            .map(|resp| resp.status());

        if matches!(
            standalone_status.as_ref(),
            Ok(&reqwest::StatusCode::OK | &reqwest::StatusCode::BAD_REQUEST)
        ) {
            debug!("detected standalone (classic) controller");
            return Ok(ControllerPlatform::ClassicController);
        }

        if matches!(
            (standalone_status.as_ref(), unifi_os_status.as_ref()),
            (
                Ok(&reqwest::StatusCode::UNAUTHORIZED),
                Ok(&reqwest::StatusCode::UNAUTHORIZED)
            )
        ) {
            debug!("detected UniFi OS controller from dual-401 login probes");
            return Ok(ControllerPlatform::UnifiOs);
        }

        if let Ok(status) = unifi_os_status.as_ref()
            && *status != reqwest::StatusCode::NOT_FOUND
        {
            debug!(?status, "detected UniFi OS controller");
            return Ok(ControllerPlatform::UnifiOs);
        }

        if let Ok(status) = standalone_status.as_ref()
            && *status != reqwest::StatusCode::NOT_FOUND
        {
            debug!(?status, "detected standalone (classic) controller");
            return Ok(ControllerPlatform::ClassicController);
        }

        if let Err(e) = standalone_status {
            return Err(Error::Transport(e));
        }

        if let Err(e) = unifi_os_status {
            return Err(Error::Transport(e));
        }

        Err(Error::Authentication {
            message: "could not detect controller platform: both login probes returned 404".into(),
        })
    }
}