openlatch-provider 0.2.1

//! `login` / `logout` / `whoami` — PKCE-style + browser-loopback flow.
//!
//! Faithful adaptation of openlatch-client's `src/cli/commands/auth.rs` flow,
//! rewritten to use openlatch-provider's types directly. The wire contract
//! with `app.openlatch.ai` is identical:
//!
//! ```text
//! 1. Bind 127.0.0.1:0 (OS-assigned port)
//! 2. Open browser to {app_url}/cli-auth?callback=...&hostname=...
//! 3. Web console reuses the better-auth session, mints a token, redirects
//!    back to http://127.0.0.1:<port>/callback?token=olk_edt_live_...&org_*
//! 4. Store via KeyringCredentialStore (with file fallback for headless Linux)
//! 5. Validate online with GET /api/v1/users/me; surface the org/user identity
//! ```
//!
//! Tokens never come via CLI flag (leaks to ps/shell-history). `--token-file`
//! is supported for CI; bare `OPENLATCH_TOKEN` env var is also honored by
//! [`crate::auth::retrieve_credential`].

use std::net::TcpListener;
use std::time::{Duration, Instant};

use clap::Args;
use secrecy::{ExposeSecret, SecretString};
use serde::Deserialize;

use crate::auth::{CredentialStore, FileCredentialStore, KeyringCredentialStore};
use crate::cli::GlobalArgs;
use crate::config::{self, Config};
use crate::error::{
    OlError, OL_4200_TOKEN_EXPIRED, OL_4202_BROWSER_LAUNCH_FAILED, OL_4203_PKCE_STATE_MISMATCH,
};
use crate::telemetry::{capture_global, Event};
use crate::ui::output::OutputConfig;

const DEFAULT_API_URL: &str = "https://api.openlatch.ai";
const DEFAULT_APP_URL: &str = "https://app.openlatch.ai";
const LOGIN_TIMEOUT: Duration = Duration::from_secs(300);
const REQUEST_BYTE_CAP: usize = 4096;

#[derive(Args, Debug)]
pub struct LoginArgs {
    /// Read the token from a file instead of triggering the browser flow.
    #[arg(long, value_name = "PATH")]
    pub token_file: Option<String>,

    /// Override the auth URL host (testing only).
    #[arg(long, hide = true)]
    pub auth_url: Option<String>,
}

// ---------------------------------------------------------------------------
// login
// ---------------------------------------------------------------------------

pub async fn login(g: &GlobalArgs, args: LoginArgs) -> Result<(), OlError> {
    let started = Instant::now();
    let out = OutputConfig::resolve(g);

    // Token-file path bypasses the browser entirely (CI / scripts).
    if let Some(path) = args.token_file.as_deref() {
        let token = std::fs::read_to_string(path)
            .map_err(|e| {
                OlError::new(
                    OL_4200_TOKEN_EXPIRED,
                    format!("cannot read token file '{path}': {e}"),
                )
            })?
            .trim()
            .to_string();
        if token.is_empty() {
            return Err(OlError::new(
                OL_4200_TOKEN_EXPIRED,
                format!("token file '{path}' is empty"),
            ));
        }
        let secret = SecretString::from(token);
        store_token(secret).await?;
        out.print_step("Token stored from --token-file");
        capture_global(Event::login_succeeded(
            "token-file",
            started.elapsed().as_millis() as u64,
        ));
        return Ok(());
    }

    let app_url = args
        .auth_url
        .clone()
        .unwrap_or_else(|| effective_app_url(&Config::load().unwrap_or_default()));

    // Bind a loopback listener so the browser can redirect to us.
    let listener = TcpListener::bind("127.0.0.1:0").map_err(|e| {
        OlError::new(
            OL_4203_PKCE_STATE_MISMATCH,
            format!("cannot bind loopback listener: {e}"),
        )
    })?;
    listener
        .set_nonblocking(false)
        .map_err(|e| OlError::new(OL_4203_PKCE_STATE_MISMATCH, e.to_string()))?;
    let port = listener
        .local_addr()
        .map_err(|e| OlError::new(OL_4203_PKCE_STATE_MISMATCH, e.to_string()))?
        .port();

    let hostname = hostname::get()
        .ok()
        .and_then(|h| h.into_string().ok())
        .unwrap_or_else(|| "unknown".into());
    let callback_url = format!("http://127.0.0.1:{port}/callback");
    let auth_url = format!(
        "{app}/cli-auth?callback={cb}&hostname={host}",
        app = app_url,
        cb = url_encode(&callback_url),
        host = url_encode(&hostname),
    );

    out.print_info(&format!("Opening {} in your browser…", &auth_url));
    if try_open_browser(&auth_url).is_err() {
        out.print_info(&format!(
            "Open this URL manually if your browser didn't pop up:\n  {auth_url}"
        ));
    }

    let spinner = out.create_spinner("Waiting for browser sign-in…");

    // Block until exactly one request lands on the loopback. 5-minute cap.
    let token = wait_for_callback(&listener, LOGIN_TIMEOUT)?;
    if let Some(s) = spinner {
        s.finish_and_clear();
    }

    let secret = SecretString::from(token);
    store_token(secret.clone()).await?;
    out.print_step("Token stored in OS keyring");

    // Online validation. Failure here doesn't undo the storage — login
    // succeeded server-side, the token just couldn't be cross-checked.
    match validate_online(&secret).await {
        Ok(identity) => {
            let org = identity.organization_name.as_deref().unwrap_or("");
            let user = identity.user_db_id.as_deref().or(identity.id.as_deref());
            out.print_step(&format!(
                "Authenticated as {} ({})",
                identity.email.as_deref().unwrap_or("?"),
                if org.is_empty() { "no org" } else { org }
            ));
            // Stitch identity for telemetry.
            if let (Some(uid), Some(machine_id)) = (user, config::machine_id_or_init().ok()) {
                if let Some(handle) = crate::telemetry::global() {
                    crate::telemetry::identity::record_auth_success(
                        handle,
                        &config::provider_dir(),
                        &machine_id,
                        uid,
                        identity.organization_id.as_deref(),
                    );
                }
            }
        }
        Err(e) => {
            out.print_info(&format!(
                "Token stored but online validation failed: {} ({})",
                e.message, e.code.code
            ));
        }
    }

    capture_global(Event::login_succeeded(
        "pkce-loopback",
        started.elapsed().as_millis() as u64,
    ));
    Ok(())
}

fn try_open_browser(url: &str) -> Result<(), OlError> {
    // Headless Linux (no DISPLAY/WAYLAND_DISPLAY) → skip silently; the URL
    // print above gives the user a manual path.
    #[cfg(target_os = "linux")]
    {
        if std::env::var_os("DISPLAY").is_none() && std::env::var_os("WAYLAND_DISPLAY").is_none() {
            return Err(OlError::new(
                OL_4202_BROWSER_LAUNCH_FAILED,
                "headless Linux: open the URL manually",
            ));
        }
    }
    open::that(url).map_err(|e| {
        OlError::new(
            OL_4202_BROWSER_LAUNCH_FAILED,
            format!("could not open browser: {e}"),
        )
    })
}

fn wait_for_callback(listener: &TcpListener, timeout: Duration) -> Result<String, OlError> {
    use std::io::{BufRead, BufReader, Write};
    listener
        .set_nonblocking(true)
        .map_err(|e| OlError::new(OL_4203_PKCE_STATE_MISMATCH, e.to_string()))?;
    let deadline = Instant::now() + timeout;

    loop {
        match listener.accept() {
            Ok((mut stream, _)) => {
                stream.set_read_timeout(Some(Duration::from_secs(10))).ok();
                let mut reader = BufReader::new(
                    stream
                        .try_clone()
                        .map_err(|e| OlError::new(OL_4203_PKCE_STATE_MISMATCH, e.to_string()))?,
                );
                let mut request_line = String::new();
                reader.read_line(&mut request_line).ok();
                if request_line.len() > REQUEST_BYTE_CAP {
                    return Err(OlError::new(
                        OL_4203_PKCE_STATE_MISMATCH,
                        "callback request exceeded 4 KB cap",
                    ));
                }
                // Drain the rest of the headers so the browser doesn't stall.
                let mut buf = String::new();
                while reader.read_line(&mut buf).map(|n| n > 2).unwrap_or(false) {
                    if buf == "\r\n" {
                        break;
                    }
                    buf.clear();
                }

                let token = parse_callback_token(&request_line)?;
                let resp = "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\nConnection: close\r\n\r\n\
                            <html><body><h2>Sign-in complete — you can close this tab.</h2></body></html>";
                let _ = stream.write_all(resp.as_bytes());
                return Ok(token);
            }
            Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                if Instant::now() >= deadline {
                    return Err(OlError::new(
                        OL_4203_PKCE_STATE_MISMATCH,
                        "browser sign-in timed out after 5 minutes",
                    ));
                }
                std::thread::sleep(Duration::from_millis(100));
            }
            Err(e) => {
                return Err(OlError::new(
                    OL_4203_PKCE_STATE_MISMATCH,
                    format!("loopback accept failed: {e}"),
                ));
            }
        }
    }
}

fn parse_callback_token(request_line: &str) -> Result<String, OlError> {
    // GET /callback?key=...&org_name=...&org_id=... HTTP/1.1
    //
    // Param name matches openlatch-client + openlatch-platform's `/cli-auth`
    // redirect — the platform sends `key`, not `token`. Mismatch here is what
    // produced the "callback was missing the `token` query parameter" bug.
    let parts: Vec<&str> = request_line.split_whitespace().collect();
    if parts.len() < 2 || parts[0] != "GET" {
        return Err(OlError::new(
            OL_4203_PKCE_STATE_MISMATCH,
            "callback request was not a GET",
        ));
    }
    let path = parts[1];
    let query = path
        .split_once('?')
        .map(|(_, q)| q)
        .ok_or_else(|| OlError::new(OL_4203_PKCE_STATE_MISMATCH, "callback had no query"))?;

    for kv in query.split('&') {
        if let Some((k, v)) = kv.split_once('=') {
            if k == "key" {
                let decoded = url_decode(v);
                if decoded.is_empty() {
                    return Err(OlError::new(
                        OL_4203_PKCE_STATE_MISMATCH,
                        "callback `key` was empty",
                    ));
                }
                return Ok(decoded);
            }
        }
    }
    Err(OlError::new(
        OL_4203_PKCE_STATE_MISMATCH,
        "callback was missing the `key` query parameter",
    ))
}

async fn store_token(secret: SecretString) -> Result<(), OlError> {
    let store = KeyringCredentialStore::new();
    if let Err(_e) = store.store_async(secret.clone()).await {
        // Fallback to encrypted file when keyring is unavailable.
        let path = config::provider_dir().join("credentials.enc");
        let machine_id = config::machine_id_or_init().unwrap_or_else(|_| "unknown".into());
        let file_store = FileCredentialStore::new(path, machine_id);
        file_store.store(secret)?;
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// logout
// ---------------------------------------------------------------------------

pub async fn logout(g: &GlobalArgs) -> Result<(), OlError> {
    let out = OutputConfig::resolve(g);

    // Best-effort server-side revocation first (fail open — local delete still proceeds).
    if let Ok(secret) = retrieve_token().await {
        let api_url = effective_api_url(&Config::load().unwrap_or_default());
        let _ = revoke_remote(&api_url, &secret).await;
    }

    let store = KeyringCredentialStore::new();
    let keyring_result = store.delete_async().await;
    let path = config::provider_dir().join("credentials.enc");
    if path.exists() {
        let _ = std::fs::remove_file(&path);
    }

    if keyring_result.is_ok() || path.exists() {
        out.print_step("Local credentials cleared");
    } else {
        out.print_step("No stored credentials to clear");
    }
    Ok(())
}

async fn retrieve_token() -> Result<SecretString, OlError> {
    let store = KeyringCredentialStore::new();
    if let Ok(s) = store.retrieve_async().await {
        return Ok(s);
    }
    if let Ok(val) = std::env::var("OPENLATCH_TOKEN") {
        if !val.is_empty() {
            return Ok(SecretString::from(val));
        }
    }
    let path = config::provider_dir().join("credentials.enc");
    let machine_id = config::machine_id_or_init().unwrap_or_else(|_| "unknown".into());
    FileCredentialStore::new(path, machine_id).retrieve()
}

async fn revoke_remote(api_url: &str, token: &SecretString) -> Result<(), OlError> {
    let client = reqwest::Client::builder()
        .use_rustls_tls()
        .timeout(Duration::from_secs(10))
        .build()
        .map_err(|e| {
            OlError::new(
                OL_4203_PKCE_STATE_MISMATCH,
                format!("reqwest client build: {e}"),
            )
        })?;
    let url = format!("{}/api/v1/api-keys/self", api_url.trim_end_matches('/'));
    let _ = client
        .delete(&url)
        .bearer_auth(token.expose_secret())
        .send()
        .await;
    Ok(())
}

// ---------------------------------------------------------------------------
// whoami
// ---------------------------------------------------------------------------

pub async fn whoami(g: &GlobalArgs) -> Result<(), OlError> {
    let out = OutputConfig::resolve(g);
    let secret = retrieve_token().await?;

    match validate_online(&secret).await {
        Ok(identity) => {
            if out.is_machine() {
                out.print_json(&serde_json::json!({
                    "online": true,
                    "email": identity.email,
                    "user_db_id": identity.user_db_id.as_deref().or(identity.id.as_deref()),
                    "organization_id": identity.organization_id,
                    "organization_name": identity.organization_name,
                }));
            } else {
                out.print_step(&format!(
                    "Authenticated as {} (org: {})",
                    identity.email.as_deref().unwrap_or("?"),
                    identity.organization_name.as_deref().unwrap_or("(none)")
                ));
            }
            Ok(())
        }
        Err(e) => {
            if out.is_machine() {
                out.print_json(&serde_json::json!({ "online": false, "error": e.code.code }));
            }
            Err(e)
        }
    }
}

#[derive(Debug, Deserialize, Default)]
pub(crate) struct AuthMeResponse {
    #[serde(default)]
    pub email: Option<String>,
    #[serde(default)]
    pub id: Option<String>,
    #[serde(default)]
    pub user_db_id: Option<String>,
    #[serde(default)]
    pub organization_id: Option<String>,
    #[serde(default)]
    pub organization_name: Option<String>,
}

/// Try to retrieve a stored token and validate it online. Returns `None` if
/// no token exists or validation fails — callers (e.g. `init`) use this to
/// detect "already logged in" without forcing a re-auth.
pub(crate) async fn current_identity() -> Option<AuthMeResponse> {
    let secret = retrieve_token().await.ok()?;
    validate_online(&secret).await.ok()
}

async fn validate_online(token: &SecretString) -> Result<AuthMeResponse, OlError> {
    let api_url = effective_api_url(&Config::load().unwrap_or_default());
    let client = reqwest::Client::builder()
        .use_rustls_tls()
        .timeout(Duration::from_secs(10))
        .build()
        .map_err(|e| OlError::new(OL_4200_TOKEN_EXPIRED, format!("reqwest build: {e}")))?;
    let url = format!("{}/api/v1/users/me", api_url.trim_end_matches('/'));
    let resp = client
        .get(&url)
        .bearer_auth(token.expose_secret())
        .send()
        .await
        .map_err(|e| OlError::new(OL_4200_TOKEN_EXPIRED, format!("network error: {e}")))?;
    let status = resp.status();
    if !status.is_success() {
        return Err(OlError::new(
            OL_4200_TOKEN_EXPIRED,
            format!("backend returned {} on /users/me", status.as_u16()),
        ));
    }
    resp.json::<AuthMeResponse>()
        .await
        .map_err(|e| OlError::new(OL_4200_TOKEN_EXPIRED, format!("parse /users/me: {e}")))
}

fn effective_api_url(cfg: &Config) -> String {
    if let Ok(env) = std::env::var("OPENLATCH_API_URL") {
        if !env.is_empty() {
            return env;
        }
    }
    if let Some(profile) = cfg.profiles.get("default") {
        if let Some(ref url) = profile.api_url {
            return url.clone();
        }
    }
    DEFAULT_API_URL.to_string()
}

/// Resolve the web-app origin used by the browser auth step.
///
/// Precedence (mirrors openlatch-client::cli::commands::auth::resolve_app_url):
///   1. `OPENLATCH_APP_URL` — explicit override (production sets this when
///      app and API live on different hosts).
///   2. `[profiles.default] app_url` — explicit per-profile override.
///   3. `OPENLATCH_API_URL` with a trailing `/api` stripped — lets dev setups
///      (Vite at `http://localhost:5173`) point both at one origin via the
///      single `OPENLATCH_API_URL` env var.
///   4. `[profiles.default] api_url` with a trailing `/api` stripped.
///   5. `https://app.openlatch.ai` default.
fn effective_app_url(cfg: &Config) -> String {
    if let Ok(env) = std::env::var("OPENLATCH_APP_URL") {
        if !env.is_empty() {
            return env;
        }
    }
    if let Some(profile) = cfg.profiles.get("default") {
        if let Some(ref url) = profile.app_url {
            return url.clone();
        }
    }
    if let Ok(env) = std::env::var("OPENLATCH_API_URL") {
        if !env.is_empty() {
            return derive_app_url_from_api(&env);
        }
    }
    if let Some(profile) = cfg.profiles.get("default") {
        if let Some(ref url) = profile.api_url {
            return derive_app_url_from_api(url);
        }
    }
    DEFAULT_APP_URL.to_string()
}

/// Derive the web-app origin from an API URL by stripping a trailing `/api`.
///
/// - `https://api.openlatch.ai/api`     → `https://api.openlatch.ai`
/// - `http://localhost:5173`            → `http://localhost:5173`
/// - `https://example.com/api/v2`       → `https://example.com/api/v2` (mid-path `/api` is preserved)
fn derive_app_url_from_api(api_url: &str) -> String {
    let trimmed = api_url.trim_end_matches('/');
    trimmed.strip_suffix("/api").unwrap_or(trimmed).to_string()
}

// ---------------------------------------------------------------------------
// URL encode/decode helpers
// ---------------------------------------------------------------------------

fn url_encode(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for &b in s.as_bytes() {
        match b {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
                out.push(b as char)
            }
            _ => out.push_str(&format!("%{b:02X}")),
        }
    }
    out
}

fn url_decode(s: &str) -> String {
    let mut out = Vec::with_capacity(s.len());
    let bytes = s.as_bytes();
    let mut i = 0;
    while i < bytes.len() {
        if bytes[i] == b'+' {
            out.push(b' ');
            i += 1;
        } else if bytes[i] == b'%' && i + 2 < bytes.len() {
            let hex = std::str::from_utf8(&bytes[i + 1..i + 3]).unwrap_or("");
            if let Ok(n) = u8::from_str_radix(hex, 16) {
                out.push(n);
                i += 3;
            } else {
                out.push(bytes[i]);
                i += 1;
            }
        } else {
            out.push(bytes[i]);
            i += 1;
        }
    }
    String::from_utf8(out).unwrap_or_default()
}

// hostname is provided by the `hostname` crate. We avoid it to keep deps lean.
mod hostname {
    pub fn get() -> std::io::Result<std::ffi::OsString> {
        // Cross-platform hostname via env or system call. Best-effort: failures
        // map to "unknown" upstream.
        if let Ok(s) = std::env::var("HOSTNAME") {
            if !s.is_empty() {
                return Ok(s.into());
            }
        }
        if let Ok(s) = std::env::var("COMPUTERNAME") {
            if !s.is_empty() {
                return Ok(s.into());
            }
        }
        Err(std::io::Error::other("no hostname env var set"))
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn url_encode_round_trips_simple_ascii() {
        assert_eq!(url_encode("hello"), "hello");
    }

    #[test]
    fn url_encode_escapes_special_chars() {
        let encoded = url_encode("https://x.io/path?a=b");
        assert!(encoded.contains("%3A%2F%2F"));
        assert!(encoded.contains("%3D"));
    }

    #[test]
    fn url_decode_reverses_encode() {
        let original = "olk_edt_live_AbCd123";
        assert_eq!(url_decode(&url_encode(original)), original);
    }

    #[test]
    fn parse_callback_token_extracts_key() {
        let line = "GET /callback?key=olk_edt_live_test123&org_name=Acme HTTP/1.1\r\n";
        assert_eq!(parse_callback_token(line).unwrap(), "olk_edt_live_test123");
    }

    #[test]
    fn parse_callback_token_rejects_empty_key() {
        let line = "GET /callback?key= HTTP/1.1\r\n";
        let err = parse_callback_token(line).unwrap_err();
        assert_eq!(err.code.code, "OL-4203");
    }

    #[test]
    fn parse_callback_token_rejects_missing_key() {
        let line = "GET /callback?other=value HTTP/1.1\r\n";
        let err = parse_callback_token(line).unwrap_err();
        assert_eq!(err.code.code, "OL-4203");
    }

    #[test]
    fn parse_callback_token_rejects_post() {
        let line = "POST /callback?key=foo HTTP/1.1\r\n";
        let err = parse_callback_token(line).unwrap_err();
        assert_eq!(err.code.code, "OL-4203");
    }

    #[test]
    fn effective_api_url_falls_back_to_default() {
        let cfg = Config::default();
        // OPENLATCH_API_URL may be set in dev env; just check shape.
        let url = effective_api_url(&cfg);
        assert!(url.starts_with("http"));
    }

    #[test]
    fn effective_app_url_falls_back_to_default() {
        let cfg = Config::default();
        let url = effective_app_url(&cfg);
        assert!(url.starts_with("http"));
    }

    #[test]
    fn derive_app_url_strips_trailing_api_suffix() {
        assert_eq!(
            derive_app_url_from_api("https://api.openlatch.ai/api"),
            "https://api.openlatch.ai"
        );
    }

    #[test]
    fn derive_app_url_strips_trailing_slash_before_api() {
        assert_eq!(
            derive_app_url_from_api("https://api.openlatch.ai/api/"),
            "https://api.openlatch.ai"
        );
    }

    #[test]
    fn derive_app_url_passes_through_bare_origin() {
        assert_eq!(
            derive_app_url_from_api("http://localhost:5173"),
            "http://localhost:5173"
        );
    }

    #[test]
    fn derive_app_url_passes_through_bare_origin_with_trailing_slash() {
        assert_eq!(
            derive_app_url_from_api("http://localhost:5173/"),
            "http://localhost:5173"
        );
    }

    #[test]
    fn derive_app_url_does_not_strip_mid_path_api_segment() {
        assert_eq!(
            derive_app_url_from_api("https://example.com/api/v2"),
            "https://example.com/api/v2"
        );
    }
}