tuimail/proton-bridge/src/auth.rs

use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::io::{self, Write};
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::{error, info, warn};
use crate::config::ProtonConfig;
use crate::srp;

const API_BASE: &str = "https://mail.proton.me/api";
const SESSION_KEY: &str = "proton_session";
const SERVICE: &str = "tuimail";
/// Refresh 5 minutes before the token actually expires.
const REFRESH_MARGIN_SECS: u64 = 300;

// ── API types ───────────────────────────────────────────────────────────────

#[derive(Serialize)]
struct AuthInfoRequest<'a> {
    #[serde(rename = "Username")]
    username: &'a str,
}

#[derive(Deserialize)]
struct AuthInfoResponse {
    #[serde(rename = "Modulus")]
    modulus: String,
    #[serde(rename = "ServerEphemeral")]
    server_ephemeral: String,
    #[serde(rename = "Version")]
    version: u32,
    #[serde(rename = "Salt")]
    salt: String,
    #[serde(rename = "SRPSession")]
    srp_session: String,
}

#[derive(Serialize)]
struct AuthRequest<'a> {
    #[serde(rename = "Username")]
    username: &'a str,
    #[serde(rename = "ClientEphemeral")]
    client_ephemeral: &'a str,
    #[serde(rename = "ClientProof")]
    client_proof: &'a str,
    #[serde(rename = "SRPSession")]
    srp_session: &'a str,
}

#[derive(Deserialize)]
struct AuthResponse {
    #[serde(rename = "UID")]
    uid: String,
    #[serde(rename = "AccessToken")]
    access_token: String,
    #[serde(rename = "RefreshToken")]
    refresh_token: String,
    #[serde(rename = "ServerProof")]
    server_proof: String,
    #[serde(rename = "ExpiresIn")]
    expires_in: u64,
    /// Integer bitmask: 0 = none, 1 = TOTP, 2 = FIDO2.
    #[serde(rename = "TwoFactor")]
    two_factor: u32,
    /// e.g. "full" or "locked". "locked" means key-sensitive endpoints
    /// (like /core/v4/keys/salts) return 403 until we call the unlock endpoint.
    #[serde(rename = "Scope", default)]
    scope: String,
}

#[derive(Serialize)]
struct RefreshRequest<'a> {
    #[serde(rename = "UID")]
    uid: &'a str,
    #[serde(rename = "RefreshToken")]
    refresh_token: &'a str,
    #[serde(rename = "ResponseType")]
    response_type: &'a str,
    #[serde(rename = "GrantType")]
    grant_type: &'a str,
    #[serde(rename = "RedirectURI")]
    redirect_uri: &'a str,
}

#[derive(Deserialize)]
struct RefreshResponse {
    #[serde(rename = "AccessToken")]
    access_token: String,
    #[serde(rename = "RefreshToken")]
    refresh_token: String,
    #[serde(rename = "ExpiresIn")]
    expires_in: u64,
}

#[derive(Serialize)]
struct TotpRequest<'a> {
    #[serde(rename = "TwoFactorCode")]
    code: &'a str,
}

// ── persisted session ────────────────────────────────────────────────────────

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Session {
    pub uid: String,
    pub access_token: String,
    pub refresh_token: String,
    /// Unix timestamp at which the access token expires.
    pub expires_at: u64,
    /// Space-separated OAuth scopes granted by the server.
    /// The `locked` scope is required for /core/v4/keys/salts.
    /// Refreshed tokens lose `locked`; only fresh SRP+TOTP logins grant it.
    #[serde(default)]
    pub scope: String,
}

impl Session {
    pub fn is_expired(&self) -> bool {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();
        now + REFRESH_MARGIN_SECS >= self.expires_at
    }

    /// Returns true if this session has the `locked` scope needed for key salts.
    pub fn has_locked_scope(&self) -> bool {
        self.scope.split_whitespace().any(|s| s == "locked")
    }

    pub fn save(&self) -> Result<(), String> {
        let json = serde_json::to_string(self).map_err(|e| e.to_string())?;
        keyring::Entry::new(SERVICE, SESSION_KEY)
            .map_err(|e| e.to_string())?
            .set_password(&json)
            .map_err(|e| e.to_string())
    }

    pub fn load() -> Option<Self> {
        let entry = keyring::Entry::new(SERVICE, SESSION_KEY).ok()?;
        let json = entry.get_password().ok()?;
        serde_json::from_str(&json).ok()
    }

    pub fn delete() {
        if let Ok(entry) = keyring::Entry::new(SERVICE, SESSION_KEY) {
            let _ = entry.delete_credential();
        }
    }
}

fn expires_at_from_now(expires_in: u64) -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
        + expires_in
}

// ── auth flow ────────────────────────────────────────────────────────────────

/// Build a reqwest client with the headers Proton's API requires.
pub fn build_client() -> Result<Client, String> {
    Client::builder()
        .user_agent("ProtonMail-Bridge-Rust/0.1")
        .default_headers({
            let mut headers = reqwest::header::HeaderMap::new();
            headers.insert("x-pm-appversion", "Other".parse().unwrap());
            headers
        })
        .build()
        .map_err(|e| e.to_string())
}

/// Return a valid session, refreshing or re-logging in as needed.
///
/// Requires the `locked` scope (needed for /core/v4/keys/salts).
/// Refreshed tokens lose that scope, so we always fall back to a full
/// SRP+TOTP login when it is absent.
pub async fn authenticate(client: &Client, config: &ProtonConfig) -> Result<Session, String> {
    if let Some(session) = Session::load() {
        if !session.has_locked_scope() {
            error!("Cached session missing 'locked' scope — re-authenticating…");
        } else if !session.is_expired() {
            info!("Using cached session for {}", config.username);
            return Ok(session);
        } else {
            info!("Access token expired, refreshing…");
            match refresh(client, &session).await {
                Ok(refreshed) => {
                    // Refreshed tokens don't carry the locked scope; save the
                    // new tokens but only use this session if it has the scope.
                    if refreshed.has_locked_scope() {
                        refreshed.save()?;
                        info!("Session refreshed for {}", config.username);
                        return Ok(refreshed);
                    }
                    warn!("Refreshed token lacks 'locked' scope — falling back to full login");
                }
                Err(e) => {
                    error!("Refresh failed ({}), falling back to full login", e);
                }
            }
        }
    }
    login(client, config).await
}

async fn refresh(client: &Client, session: &Session) -> Result<Session, String> {
    let resp: RefreshResponse = client
        .post(format!("{}/auth/v4/refresh", API_BASE))
        .header("x-pm-uid", &session.uid)
        .json(&RefreshRequest {
            uid: &session.uid,
            refresh_token: &session.refresh_token,
            response_type: "token",
            grant_type: "refresh_token",
            redirect_uri: "https://proton.me",
        })
        .send()
        .await
        .map_err(|e| e.to_string())?
        .error_for_status()
        .map_err(|e| e.to_string())?
        .json()
        .await
        .map_err(|e| e.to_string())?;

    Ok(Session {
        uid: session.uid.clone(),
        access_token: resp.access_token,
        refresh_token: resp.refresh_token,
        expires_at: expires_at_from_now(resp.expires_in),
        scope: String::new(), // refreshed tokens don't carry the locked scope
    })
}

async fn login(client: &Client, config: &ProtonConfig) -> Result<Session, String> {
    // Step 1: request SRP challenge
    let info: AuthInfoResponse = client
        .post(format!("{}/auth/info", API_BASE))
        .json(&AuthInfoRequest { username: &config.username })
        .send()
        .await
        .map_err(|e| e.to_string())?
        .error_for_status()
        .map_err(|e| e.to_string())?
        .json()
        .await
        .map_err(|e| e.to_string())?;

    // The modulus arrives PGP-armored — strip the armor to get the raw base64.
    let modulus_b64 = strip_pgp_armor(&info.modulus)?;

    // Step 2: compute SRP proofs
    let proofs = srp::generate_proofs(
        info.version,
        &config.password,
        &info.salt,
        &info.server_ephemeral,
        &modulus_b64,
    )?;

    // Step 3: submit proof
    let auth_resp = client
        .post(format!("{}/auth/v4", API_BASE))
        .json(&AuthRequest {
            username: &config.username,
            client_ephemeral: &proofs.client_ephemeral,
            client_proof: &proofs.client_proof,
            srp_session: &info.srp_session,
        })
        .send()
        .await
        .map_err(|e| e.to_string())?;

    if !auth_resp.status().is_success() {
        let status = auth_resp.status();
        let body = auth_resp.text().await.unwrap_or_default();
        return Err(format!("auth/v4 failed ({status}): {body}"));
    }

    let auth: AuthResponse = auth_resp.json().await.map_err(|e| e.to_string())?;

    // Verify the server proved it knows the password too
    if auth.server_proof != proofs.expected_server_proof {
        return Err("Server proof verification failed — possible MITM".to_string());
    }

    let session = Session {
        uid: auth.uid.clone(),
        access_token: auth.access_token.clone(),
        refresh_token: auth.refresh_token.clone(),
        expires_at: expires_at_from_now(auth.expires_in),
        scope: auth.scope.clone(),
    };

    // Step 4: handle TOTP if enabled
    if auth.two_factor & 1 != 0 {
        submit_totp(client, &session).await?;
    }

    let has_full = auth.scope.split_whitespace().any(|s| s == "full");
    info!("Auth scope tokens: {}", auth.scope);
    if !has_full {
        warn!("Token does not have 'full' scope — some endpoints may return 403.");
    }

    session.save()?;
    info!("Authenticated as {}", config.username);
    Ok(session)
}

async fn submit_totp(client: &Client, session: &Session) -> Result<(), String> {
    info!("TOTP code: ");
    io::stdout().flush().unwrap();
    let mut code = String::new();
    io::stdin().read_line(&mut code).map_err(|e| e.to_string())?;
    let code = code.trim();

    client
        .post(format!("{}/auth/v4/2fa", API_BASE))
        .bearer_auth(&session.access_token)
        .header("x-pm-uid", &session.uid)
        .json(&TotpRequest { code })
        .send()
        .await
        .map_err(|e| e.to_string())?
        .error_for_status()
        .map_err(|e| e.to_string())?;

    Ok(())
}

// ── helpers ──────────────────────────────────────────────────────────────────

/// Remove PGP armor lines and return only the base64 body.
/// The modulus arrives as:
///   -----BEGIN PGP SIGNED MESSAGE-----
///   ...base64...
///   -----BEGIN PGP SIGNATURE-----
///   ...
fn strip_pgp_armor(armored: &str) -> Result<String, String> {
    let mut in_body = false;
    let mut body = String::new();
    for line in armored.lines() {
        if line.starts_with("-----BEGIN PGP SIGNED MESSAGE") {
            in_body = false;
            continue;
        }
        if line.is_empty() && !in_body {
            in_body = true;
            continue;
        }
        if line.starts_with("-----BEGIN PGP SIGNATURE") {
            break;
        }
        if in_body {
            body.push_str(line.trim());
        }
    }
    if body.is_empty() {
        Err("failed to extract modulus from PGP armor".to_string())
    } else {
        Ok(body)
    }
}