secrets/crates/secrets-mcp/src/web.rs

use askama::Template;
use chrono::SecondsFormat;
use std::net::SocketAddr;

use axum::{
    Json, Router,
    body::Body,
    extract::{ConnectInfo, Path, Query, State},
    http::{HeaderMap, StatusCode, header},
    response::{Html, IntoResponse, Redirect, Response},
    routing::{get, post},
};
use serde::{Deserialize, Serialize};
use tower_sessions::Session;
use uuid::Uuid;

use secrets_core::audit::log_login;
use secrets_core::crypto::hex;
use secrets_core::service::{
    api_key::{ensure_api_key, regenerate_api_key},
    audit_log::list_for_user,
    user::{
        OAuthProfile, bind_oauth_account, find_or_create_user, get_user_by_id,
        unbind_oauth_account, update_user_key_setup,
    },
};

use crate::AppState;
use crate::oauth::{OAuthConfig, OAuthUserInfo, google_auth_url, random_state};

const SESSION_USER_ID: &str = "user_id";
const SESSION_OAUTH_STATE: &str = "oauth_state";
const SESSION_OAUTH_BIND_MODE: &str = "oauth_bind_mode";
const SESSION_LOGIN_PROVIDER: &str = "login_provider";

// ── Template types ────────────────────────────────────────────────────────────

#[derive(Template)]
#[template(path = "login.html")]
struct LoginTemplate {
    has_google: bool,
    base_url: String,
    version: &'static str,
}

#[derive(Template)]
#[template(path = "home.html")]
struct HomeTemplate {
    is_logged_in: bool,
    base_url: String,
    version: &'static str,
}

#[derive(Template)]
#[template(path = "dashboard.html")]
struct DashboardTemplate {
    user_name: String,
    user_email: String,
    has_passphrase: bool,
    base_url: String,
    version: &'static str,
}

#[derive(Template)]
#[template(path = "audit.html")]
struct AuditPageTemplate {
    user_name: String,
    user_email: String,
    entries: Vec<AuditEntryView>,
    version: &'static str,
}

struct AuditEntryView {
    /// RFC3339 UTC for `<time datetime>`; rendered as browser-local in audit.html.
    created_at_iso: String,
    action: String,
    target: String,
    detail: String,
}

// ── App state helpers ─────────────────────────────────────────────────────────

fn google_cfg(state: &AppState) -> Option<&OAuthConfig> {
    state.google_config.as_ref()
}

async fn current_user_id(session: &Session) -> Option<Uuid> {
    match session.get::<String>(SESSION_USER_ID).await {
        Ok(opt) => match opt {
            Some(s) => match Uuid::parse_str(&s) {
                Ok(id) => Some(id),
                Err(e) => {
                    tracing::warn!(error = %e, user_id_str = %s, "invalid user_id UUID in session");
                    None
                }
            },
            None => None,
        },
        Err(e) => {
            tracing::warn!(error = %e, "failed to read user_id from session");
            None
        }
    }
}

fn request_client_ip(headers: &HeaderMap, connect_info: ConnectInfo<SocketAddr>) -> Option<String> {
    if let Some(first) = headers
        .get("x-forwarded-for")
        .and_then(|v| v.to_str().ok())
        .and_then(|s| s.split(',').next())
    {
        let value = first.trim();
        if !value.is_empty() {
            return Some(value.to_string());
        }
    }

    Some(connect_info.ip().to_string())
}

fn request_user_agent(headers: &HeaderMap) -> Option<String> {
    headers
        .get(header::USER_AGENT)
        .and_then(|value| value.to_str().ok())
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(ToOwned::to_owned)
}

// ── Routes ────────────────────────────────────────────────────────────────────

pub fn web_router() -> Router<AppState> {
    Router::new()
        .route("/robots.txt", get(robots_txt))
        .route("/llms.txt", get(llms_txt))
        .route("/ai.txt", get(ai_txt))
        .route("/favicon.svg", get(favicon_svg))
        .route(
            "/favicon.ico",
            get(|| async { Redirect::permanent("/favicon.svg") }),
        )
        .route(
            "/.well-known/oauth-protected-resource",
            get(oauth_protected_resource_metadata),
        )
        .route("/", get(home_page))
        .route("/login", get(login_page))
        .route("/auth/google", get(auth_google))
        .route("/auth/google/callback", get(auth_google_callback))
        .route("/auth/logout", post(auth_logout))
        .route("/dashboard", get(dashboard))
        .route("/audit", get(audit_page))
        .route("/account/bind/google", get(account_bind_google))
        .route(
            "/account/bind/google/callback",
            get(account_bind_google_callback),
        )
        .route("/account/unbind/{provider}", post(account_unbind))
        .route("/api/key-salt", get(api_key_salt))
        .route("/api/key-setup", post(api_key_setup))
        .route("/api/apikey", get(api_apikey_get))
        .route("/api/apikey/regenerate", post(api_apikey_regenerate))
}

fn text_asset_response(content: &'static str, content_type: &'static str) -> Response {
    Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, content_type)
        .header(header::CACHE_CONTROL, "public, max-age=86400")
        .body(Body::from(content))
        .expect("text asset response")
}

async fn robots_txt() -> Response {
    text_asset_response(
        include_str!("../static/robots.txt"),
        "text/plain; charset=utf-8",
    )
}

async fn llms_txt() -> Response {
    text_asset_response(
        include_str!("../static/llms.txt"),
        "text/markdown; charset=utf-8",
    )
}

async fn ai_txt() -> Response {
    llms_txt().await
}

async fn favicon_svg() -> Response {
    Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "image/svg+xml")
        .header(header::CACHE_CONTROL, "public, max-age=86400")
        .body(Body::from(include_str!("../static/favicon.svg")))
        .expect("favicon response")
}

// ── Home page (public) ───────────────────────────────────────────────────────

async fn home_page(
    State(state): State<AppState>,
    session: Session,
) -> Result<Response, StatusCode> {
    let is_logged_in = current_user_id(&session).await.is_some();
    let tmpl = HomeTemplate {
        is_logged_in,
        base_url: state.base_url.clone(),
        version: env!("CARGO_PKG_VERSION"),
    };
    render_template(tmpl)
}

// ── Login page ────────────────────────────────────────────────────────────────

async fn login_page(
    State(state): State<AppState>,
    session: Session,
) -> Result<Response, StatusCode> {
    if let Some(_uid) = current_user_id(&session).await {
        return Ok(Redirect::to("/dashboard").into_response());
    }

    let tmpl = LoginTemplate {
        has_google: state.google_config.is_some(),
        base_url: state.base_url.clone(),
        version: env!("CARGO_PKG_VERSION"),
    };
    render_template(tmpl)
}

// ── Google OAuth ──────────────────────────────────────────────────────────────

async fn auth_google(
    State(state): State<AppState>,
    session: Session,
) -> Result<Response, StatusCode> {
    let config = google_cfg(&state).ok_or(StatusCode::SERVICE_UNAVAILABLE)?;

    let oauth_state = random_state();
    session
        .insert(SESSION_OAUTH_STATE, &oauth_state)
        .await
        .map_err(|e| {
            tracing::error!(error = %e, "failed to insert oauth_state into session");
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    let url = google_auth_url(config, &oauth_state);
    Ok(Redirect::to(&url).into_response())
}

#[derive(Deserialize)]
struct OAuthCallbackQuery {
    code: Option<String>,
    state: Option<String>,
    error: Option<String>,
}

async fn auth_google_callback(
    State(state): State<AppState>,
    connect_info: ConnectInfo<SocketAddr>,
    headers: HeaderMap,
    session: Session,
    Query(params): Query<OAuthCallbackQuery>,
) -> Result<Response, StatusCode> {
    let client_ip = request_client_ip(&headers, connect_info);
    let user_agent = request_user_agent(&headers);
    handle_oauth_callback(
        &state,
        &session,
        params,
        "google",
        client_ip.as_deref(),
        user_agent.as_deref(),
        |s, cfg, code| {
            Box::pin(crate::oauth::google::exchange_code(
                &s.http_client,
                cfg,
                code,
            ))
        },
    )
    .await
}

// ── Shared OAuth callback handler ─────────────────────────────────────────────

async fn handle_oauth_callback<F>(
    state: &AppState,
    session: &Session,
    params: OAuthCallbackQuery,
    provider: &str,
    client_ip: Option<&str>,
    user_agent: Option<&str>,
    exchange_fn: F,
) -> Result<Response, StatusCode>
where
    F: for<'a> Fn(
        &'a AppState,
        &'a OAuthConfig,
        &'a str,
    ) -> std::pin::Pin<
        Box<dyn std::future::Future<Output = anyhow::Result<OAuthUserInfo>> + Send + 'a>,
    >,
{
    if let Some(err) = params.error {
        tracing::warn!(provider, error = %err, "OAuth error");
        return Ok(Redirect::to("/login?error=oauth_error").into_response());
    }

    let Some(code) = params.code else {
        tracing::warn!(provider, "OAuth callback missing code");
        return Ok(Redirect::to("/login?error=oauth_missing_code").into_response());
    };
    let Some(returned_state) = params.state.as_deref() else {
        tracing::warn!(provider, "OAuth callback missing state");
        return Ok(Redirect::to("/login?error=oauth_missing_state").into_response());
    };

    let expected_state: Option<String> = session.get(SESSION_OAUTH_STATE).await.map_err(|e| {
        tracing::error!(provider, error = %e, "failed to read oauth_state from session");
        StatusCode::INTERNAL_SERVER_ERROR
    })?;
    if expected_state.as_deref() != Some(returned_state) {
        tracing::warn!(
            provider,
            expected_present = expected_state.is_some(),
            "OAuth state mismatch (empty session often means SameSite=Strict or server restart)"
        );
        return Ok(Redirect::to("/login?error=oauth_state").into_response());
    }
    if let Err(e) = session.remove::<String>(SESSION_OAUTH_STATE).await {
        tracing::warn!(provider, error = %e, "failed to remove oauth_state from session");
    }

    let config = match provider {
        "google" => state
            .google_config
            .as_ref()
            .ok_or(StatusCode::SERVICE_UNAVAILABLE)?,
        _ => return Err(StatusCode::BAD_REQUEST),
    };

    let user_info = exchange_fn(state, config, code.as_str())
        .await
        .map_err(|e| {
            tracing::error!(provider, error = %e, "failed to exchange OAuth code");
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    let bind_mode: bool = match session.get::<bool>(SESSION_OAUTH_BIND_MODE).await {
        Ok(v) => v.unwrap_or(false),
        Err(e) => {
            tracing::error!(
                provider,
                error = %e,
                "failed to read oauth_bind_mode from session"
            );
            return Err(StatusCode::INTERNAL_SERVER_ERROR);
        }
    };

    if bind_mode {
        let user_id = current_user_id(session)
            .await
            .ok_or(StatusCode::UNAUTHORIZED)?;
        if let Err(e) = session.remove::<bool>(SESSION_OAUTH_BIND_MODE).await {
            tracing::warn!(provider, error = %e, "failed to remove oauth_bind_mode from session after bind");
        }

        let profile = OAuthProfile {
            provider: user_info.provider,
            provider_id: user_info.provider_id,
            email: user_info.email,
            name: user_info.name,
            avatar_url: user_info.avatar_url,
        };

        bind_oauth_account(&state.pool, user_id, profile)
            .await
            .map_err(|e| {
                tracing::error!(error = %e, "failed to bind OAuth account");
                StatusCode::INTERNAL_SERVER_ERROR
            })?;

        return Ok(Redirect::to("/dashboard?bound=1").into_response());
    }

    let profile = OAuthProfile {
        provider: user_info.provider,
        provider_id: user_info.provider_id,
        email: user_info.email,
        name: user_info.name,
        avatar_url: user_info.avatar_url,
    };

    let (user, _is_new) = find_or_create_user(&state.pool, profile)
        .await
        .map_err(|e| {
            tracing::error!(error = %e, "failed to find or create user");
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    session
        .insert(SESSION_USER_ID, user.id.to_string())
        .await
        .map_err(|e| {
            tracing::error!(
                error = %e,
                user_id = %user.id,
                "failed to insert user_id into session after OAuth"
            );
            StatusCode::INTERNAL_SERVER_ERROR
        })?;
    session
        .insert(SESSION_LOGIN_PROVIDER, &provider)
        .await
        .map_err(|e| {
            tracing::error!(
                provider,
                error = %e,
                "failed to insert login_provider into session after OAuth"
            );
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    log_login(
        &state.pool,
        "oauth",
        provider,
        user.id,
        client_ip,
        user_agent,
    )
    .await;

    Ok(Redirect::to("/dashboard").into_response())
}

// ── Logout ────────────────────────────────────────────────────────────────────

async fn auth_logout(session: Session) -> impl IntoResponse {
    if let Err(e) = session.flush().await {
        tracing::warn!(error = %e, "failed to flush session on logout");
    }
    Redirect::to("/")
}

// ── Dashboard ─────────────────────────────────────────────────────────────────

async fn dashboard(
    State(state): State<AppState>,
    session: Session,
) -> Result<Response, StatusCode> {
    let Some(user_id) = current_user_id(&session).await else {
        return Ok(Redirect::to("/login").into_response());
    };

    let user = match get_user_by_id(&state.pool, user_id).await.map_err(|e| {
        tracing::error!(error = %e, %user_id, "failed to load user for dashboard");
        StatusCode::INTERNAL_SERVER_ERROR
    })? {
        Some(u) => u,
        None => return Ok(Redirect::to("/login").into_response()),
    };

    let tmpl = DashboardTemplate {
        user_name: user.name.clone(),
        user_email: user.email.clone().unwrap_or_default(),
        has_passphrase: user.key_salt.is_some(),
        base_url: state.base_url.clone(),
        version: env!("CARGO_PKG_VERSION"),
    };

    render_template(tmpl)
}

async fn audit_page(
    State(state): State<AppState>,
    session: Session,
) -> Result<Response, StatusCode> {
    let Some(user_id) = current_user_id(&session).await else {
        return Ok(Redirect::to("/login").into_response());
    };

    let user = match get_user_by_id(&state.pool, user_id).await.map_err(|e| {
        tracing::error!(error = %e, %user_id, "failed to load user for audit page");
        StatusCode::INTERNAL_SERVER_ERROR
    })? {
        Some(u) => u,
        None => return Ok(Redirect::to("/login").into_response()),
    };

    let rows = list_for_user(&state.pool, user_id, 100)
        .await
        .map_err(|e| {
            tracing::error!(error = %e, "failed to load audit log for user");
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    let entries = rows
        .into_iter()
        .map(|row| AuditEntryView {
            created_at_iso: row.created_at.to_rfc3339_opts(SecondsFormat::Secs, true),
            action: row.action,
            target: format_audit_target(&row.folder, &row.entry_type, &row.name),
            detail: serde_json::to_string_pretty(&row.detail).unwrap_or_else(|_| "{}".to_string()),
        })
        .collect();

    let tmpl = AuditPageTemplate {
        user_name: user.name.clone(),
        user_email: user.email.clone().unwrap_or_default(),
        entries,
        version: env!("CARGO_PKG_VERSION"),
    };

    render_template(tmpl)
}

// ── Account bind/unbind ───────────────────────────────────────────────────────

async fn account_bind_google(
    State(state): State<AppState>,
    session: Session,
) -> Result<Response, StatusCode> {
    let _ = current_user_id(&session)
        .await
        .ok_or(StatusCode::UNAUTHORIZED)?;

    session
        .insert(SESSION_OAUTH_BIND_MODE, true)
        .await
        .map_err(|e| {
            tracing::error!(error = %e, "failed to insert oauth_bind_mode into session");
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    let redirect_uri = format!("{}/account/bind/google/callback", state.base_url);
    let mut cfg = state
        .google_config
        .clone()
        .ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
    cfg.redirect_uri = redirect_uri;
    let st = random_state();
    if let Err(e) = session.insert(SESSION_OAUTH_STATE, &st).await {
        tracing::error!(error = %e, "failed to insert oauth_state for account bind flow");
        if let Err(rm) = session.remove::<bool>(SESSION_OAUTH_BIND_MODE).await {
            tracing::warn!(error = %rm, "failed to roll back oauth_bind_mode after oauth_state insert failure");
        }
        return Err(StatusCode::INTERNAL_SERVER_ERROR);
    }

    Ok(Redirect::to(&google_auth_url(&cfg, &st)).into_response())
}

async fn account_bind_google_callback(
    State(state): State<AppState>,
    connect_info: ConnectInfo<SocketAddr>,
    headers: HeaderMap,
    session: Session,
    Query(params): Query<OAuthCallbackQuery>,
) -> Result<Response, StatusCode> {
    let client_ip = request_client_ip(&headers, connect_info);
    let user_agent = request_user_agent(&headers);
    handle_oauth_callback(
        &state,
        &session,
        params,
        "google",
        client_ip.as_deref(),
        user_agent.as_deref(),
        |s, cfg, code| {
            Box::pin(crate::oauth::google::exchange_code(
                &s.http_client,
                cfg,
                code,
            ))
        },
    )
    .await
}

async fn account_unbind(
    State(state): State<AppState>,
    Path(provider): Path<String>,
    session: Session,
) -> Result<Response, StatusCode> {
    let user_id = current_user_id(&session)
        .await
        .ok_or(StatusCode::UNAUTHORIZED)?;

    let current_login_provider = session
        .get::<String>(SESSION_LOGIN_PROVIDER)
        .await
        .map_err(|e| {
            tracing::error!(error = %e, "failed to read login_provider from session");
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    unbind_oauth_account(
        &state.pool,
        user_id,
        &provider,
        current_login_provider.as_deref(),
    )
    .await
    .map_err(|e| {
        tracing::warn!(error = %e, "failed to unbind oauth account");
        StatusCode::BAD_REQUEST
    })?;

    Ok(Redirect::to("/dashboard?unbound=1").into_response())
}

// ── Passphrase / Key setup API ────────────────────────────────────────────────

#[derive(Serialize)]
struct KeySaltResponse {
    has_passphrase: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    salt: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    key_check: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    params: Option<serde_json::Value>,
}

async fn api_key_salt(
    State(state): State<AppState>,
    session: Session,
) -> Result<Json<KeySaltResponse>, StatusCode> {
    let user_id = current_user_id(&session)
        .await
        .ok_or(StatusCode::UNAUTHORIZED)?;

    let user = get_user_by_id(&state.pool, user_id)
        .await
        .map_err(|e| {
            tracing::error!(error = %e, %user_id, "failed to load user for key-salt API");
            StatusCode::INTERNAL_SERVER_ERROR
        })?
        .ok_or(StatusCode::UNAUTHORIZED)?;

    if user.key_salt.is_none() {
        return Ok(Json(KeySaltResponse {
            has_passphrase: false,
            salt: None,
            key_check: None,
            params: None,
        }));
    }

    Ok(Json(KeySaltResponse {
        has_passphrase: true,
        salt: user.key_salt.as_deref().map(hex::encode_hex),
        key_check: user.key_check.as_deref().map(hex::encode_hex),
        params: user.key_params,
    }))
}

#[derive(Deserialize)]
struct KeySetupRequest {
    /// Hex-encoded 32-byte random salt
    salt: String,
    /// Hex-encoded AES-256-GCM encryption of "secrets-mcp-key-check" with the derived key
    key_check: String,
    /// Key derivation parameters, e.g. {"alg":"pbkdf2-sha256","iterations":600000}
    params: serde_json::Value,
}

#[derive(Serialize)]
struct KeySetupResponse {
    ok: bool,
}

async fn api_key_setup(
    State(state): State<AppState>,
    session: Session,
    Json(body): Json<KeySetupRequest>,
) -> Result<Json<KeySetupResponse>, StatusCode> {
    let user_id = current_user_id(&session)
        .await
        .ok_or(StatusCode::UNAUTHORIZED)?;

    let salt = hex::decode_hex(&body.salt).map_err(|e| {
        tracing::warn!(error = %e, "invalid hex in key-setup salt");
        StatusCode::BAD_REQUEST
    })?;
    let key_check = hex::decode_hex(&body.key_check).map_err(|e| {
        tracing::warn!(error = %e, "invalid hex in key-setup key_check");
        StatusCode::BAD_REQUEST
    })?;

    if salt.len() != 32 {
        tracing::warn!(salt_len = salt.len(), "key-setup salt must be 32 bytes");
        return Err(StatusCode::BAD_REQUEST);
    }

    update_user_key_setup(&state.pool, user_id, &salt, &key_check, &body.params)
        .await
        .map_err(|e| {
            tracing::error!(error = %e, "failed to update key setup");
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    Ok(Json(KeySetupResponse { ok: true }))
}

// ── API Key management ────────────────────────────────────────────────────────

#[derive(Serialize)]
struct ApiKeyResponse {
    api_key: String,
}

async fn api_apikey_get(
    State(state): State<AppState>,
    session: Session,
) -> Result<Json<ApiKeyResponse>, StatusCode> {
    let user_id = current_user_id(&session)
        .await
        .ok_or(StatusCode::UNAUTHORIZED)?;

    let api_key = ensure_api_key(&state.pool, user_id).await.map_err(|e| {
        tracing::error!(error = %e, %user_id, "ensure_api_key failed");
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    Ok(Json(ApiKeyResponse { api_key }))
}

async fn api_apikey_regenerate(
    State(state): State<AppState>,
    session: Session,
) -> Result<Json<ApiKeyResponse>, StatusCode> {
    let user_id = current_user_id(&session)
        .await
        .ok_or(StatusCode::UNAUTHORIZED)?;

    let api_key = regenerate_api_key(&state.pool, user_id)
        .await
        .map_err(|e| {
            tracing::error!(error = %e, %user_id, "regenerate_api_key failed");
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    Ok(Json(ApiKeyResponse { api_key }))
}

// ── OAuth / Well-known ────────────────────────────────────────────────────────

/// RFC 9728 — OAuth 2.0 Protected Resource Metadata.
///
/// Advertises that this server accepts Bearer tokens in the `Authorization`
/// header.  We deliberately omit `authorization_servers` because this service
/// issues its own API keys (no external OAuth AS is involved).  MCP clients
/// that probe this endpoint will see the resource identifier and stop looking
/// for a delegated OAuth flow.
async fn oauth_protected_resource_metadata(State(state): State<AppState>) -> impl IntoResponse {
    let body = serde_json::json!({
        "resource": state.base_url,
        "bearer_methods_supported": ["header"],
        "resource_documentation": format!("{}/dashboard", state.base_url),
    });
    (
        StatusCode::OK,
        [(header::CONTENT_TYPE, "application/json")],
        axum::Json(body),
    )
}

// ── Helper ────────────────────────────────────────────────────────────────────

fn render_template<T: Template>(tmpl: T) -> Result<Response, StatusCode> {
    let html = tmpl.render().map_err(|e| {
        tracing::error!(error = %e, "template render error");
        StatusCode::INTERNAL_SERVER_ERROR
    })?;
    Ok(Html(html).into_response())
}

fn format_audit_target(folder: &str, entry_type: &str, name: &str) -> String {
    // Auth events (folder="auth") use entry_type/name as provider-scoped target.
    if folder == "auth" {
        format!("{}/{}", entry_type, name)
    } else if !folder.is_empty() && !entry_type.is_empty() {
        format!("[{}/{}] {}", folder, entry_type, name)
    } else if !folder.is_empty() {
        format!("[{}] {}", folder, name)
    } else {
        name.to_string()
    }
}