defter 0.2.0

//! API keys, the browser login session that hands one to the Zotero client,
//! and the expiring tokens used for uploads and downloads.
//!
//! Keys are provisioned out of band (files or environment) and held in
//! memory, never stored in the database. Each is read/write or read-only.
//! The only path that *hands out* a key is the login session, released only
//! after a human approves it at `/login`, which a deployment puts behind SSO
//! or a private network. Adapted from zhost (<https://github.com/mulatta/zhost>,
//! MIT).

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use std::time::{Duration, Instant};

use axum::{
    Form, Json,
    extract::{Path as UrlPath, Query, Request, State},
    http::{HeaderMap, Method, StatusCode},
    middleware::Next,
    response::{IntoResponse, Redirect, Response},
};
use serde_json::{Value, json};

use crate::app::AppState;
use crate::ui::{self, LoginView};

/// `bytes` of OS randomness, hex-encoded.
pub fn random_hex(bytes: usize) -> String {
    let mut buf = vec![0u8; bytes];
    getrandom::fill(&mut buf).expect("OS random number generator");
    buf.iter().map(|b| format!("{b:02x}")).collect()
}

// ── keys ─────────────────────────────────────────────────────────────────────

/// What a key may do. Reads are always allowed; `write` gates mutations.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Access {
    pub write: bool,
}

/// Where keys come from, in order of precedence.
#[derive(Debug, Default, Clone)]
pub struct KeySource {
    /// `rw:<path>,ro:<path>,…` — one single-line token file per entry.
    pub manifest: Option<String>,
    /// A single read/write token file.
    pub file: Option<PathBuf>,
    /// A single read/write token (prefer a file: the environment is visible in
    /// `/proc`).
    pub key: Option<String>,
    /// With none of the above: a token kept in this file, generated (mode
    /// 0600) on first start.
    pub generated: PathBuf,
}

#[derive(Debug)]
pub struct Keys {
    map: HashMap<String, Access>,
    /// The read/write key handed to the Zotero client by the login session.
    login_key: String,
}

impl Keys {
    pub fn new(map: HashMap<String, Access>) -> Result<Self, String> {
        let login_key = map
            .iter()
            .filter(|(_, a)| a.write)
            .map(|(k, _)| k.clone())
            .min()
            .ok_or("at least one read/write key must be configured")?;
        Ok(Keys { map, login_key })
    }

    /// Load the keys. Returns them and, when a key was generated, its file.
    pub fn load(source: &KeySource) -> Result<(Self, Option<PathBuf>), String> {
        let mut map = HashMap::new();
        let mut generated = None;
        if let Some(manifest) = source.manifest.as_deref().filter(|m| !m.trim().is_empty()) {
            for entry in manifest.split(',').map(str::trim).filter(|s| !s.is_empty()) {
                let (role, path) = entry
                    .split_once(':')
                    .ok_or_else(|| format!("key entry is not <rw|ro>:<path>: {entry}"))?;
                let write = match role {
                    "rw" => true,
                    "ro" => false,
                    other => return Err(format!("unknown key role '{other}' (rw or ro)")),
                };
                map.insert(read_token(Path::new(path))?, Access { write });
            }
        } else if let Some(file) = &source.file {
            map.insert(read_token(file)?, Access { write: true });
        } else if let Some(key) = source.key.as_deref().filter(|k| !k.trim().is_empty()) {
            map.insert(key.trim().to_string(), Access { write: true });
        } else {
            let (token, created) = generated_token(&source.generated)?;
            if created {
                generated = Some(source.generated.clone());
            }
            map.insert(token, Access { write: true });
        }
        Ok((Keys::new(map)?, generated))
    }

    pub fn access(&self, token: &str) -> Option<Access> {
        self.map.get(token).copied()
    }

    pub fn login_key(&self) -> &str {
        &self.login_key
    }
}

fn read_token(path: &Path) -> Result<String, String> {
    let token = std::fs::read_to_string(path)
        .map_err(|e| format!("read key file {}: {e}", path.display()))?
        .trim()
        .to_string();
    if token.is_empty() {
        return Err(format!("key file {} is empty", path.display()));
    }
    Ok(token)
}

/// The token in `path`, generating it on first use. Returns whether it was
/// created now.
pub fn generated_token(path: &Path) -> Result<(String, bool), String> {
    if path.exists() {
        return read_token(path).map(|t| (t, false));
    }
    if let Some(dir) = path.parent() {
        std::fs::create_dir_all(dir).map_err(|e| format!("create {}: {e}", dir.display()))?;
    }
    let token = random_hex(32);
    std::fs::write(path, format!("{token}\n"))
        .map_err(|e| format!("write {}: {e}", path.display()))?;
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))
            .map_err(|e| format!("chmod {}: {e}", path.display()))?;
    }
    Ok((token, true))
}

/// The key a request presents: `Zotero-API-Key`, or `Authorization: Bearer`.
pub fn request_key(headers: &HeaderMap) -> Option<&str> {
    headers
        .get("zotero-api-key")
        .and_then(|v| v.to_str().ok())
        .or_else(|| {
            headers
                .get("authorization")
                .and_then(|v| v.to_str().ok())
                .and_then(|v| v.strip_prefix("Bearer "))
        })
        .map(str::trim)
}

/// The `<id>` of a `/users/<id>/…` path.
fn user_segment(path: &str) -> Option<&str> {
    let mut segments = path.split('/');
    segments.find(|s| *s == "users")?;
    segments.next()
}

/// Guard the data API: a known key, write access for mutations, and the
/// configured user's library (`users/0`, the local API's "current user",
/// is accepted as an alias).
pub async fn require_key(State(app): State<AppState>, req: Request, next: Next) -> Response {
    let Some(access) = request_key(req.headers()).and_then(|k| app.keys.access(k)) else {
        return (StatusCode::FORBIDDEN, "Invalid key").into_response();
    };
    let mutating = matches!(
        *req.method(),
        Method::POST | Method::PUT | Method::PATCH | Method::DELETE
    );
    // Unlinking the account in Zotero deletes the key; let a read-only key
    // do that too.
    let unlink = req.method() == Method::DELETE && req.uri().path().ends_with("/keys/current");
    if mutating && !access.write && !unlink {
        return (StatusCode::FORBIDDEN, "Write access denied").into_response();
    }
    if let Some(id) = user_segment(req.uri().path())
        && id != "0"
        && id != app.cfg.user_id.to_string()
    {
        return (StatusCode::FORBIDDEN, "Forbidden").into_response();
    }
    next.run(req).await
}

// ── expiring tokens ──────────────────────────────────────────────────────────

/// Values keyed by unguessable tokens that expire after a fixed lifetime.
/// Expired entries are pruned on insert, so an abandoned one cannot linger.
pub struct Tokens<T> {
    ttl: Duration,
    map: Mutex<HashMap<String, (Instant, T)>>,
}

impl<T: Clone> Tokens<T> {
    pub fn new(ttl: Duration) -> Self {
        Tokens {
            ttl,
            map: Mutex::new(HashMap::new()),
        }
    }

    fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<String, (Instant, T)>> {
        self.map.lock().unwrap_or_else(|e| e.into_inner())
    }

    /// Store `value` under a fresh token. Returns the token and the entries
    /// that had expired.
    pub fn insert(&self, value: T) -> (String, Vec<(String, T)>) {
        let token = random_hex(16);
        let mut map = self.lock();
        let expired_keys: Vec<String> = map
            .iter()
            .filter(|(_, (created, _))| created.elapsed() >= self.ttl)
            .map(|(k, _)| k.clone())
            .collect();
        let expired = expired_keys
            .into_iter()
            .filter_map(|k| map.remove(&k).map(|(_, v)| (k, v)))
            .collect();
        map.insert(token.clone(), (Instant::now(), value));
        (token, expired)
    }

    pub fn get(&self, token: &str) -> Option<T> {
        self.lock()
            .get(token)
            .filter(|(created, _)| created.elapsed() < self.ttl)
            .map(|(_, v)| v.clone())
    }

    /// Change a live entry; false if it is unknown or expired.
    pub fn update(&self, token: &str, f: impl FnOnce(&mut T)) -> bool {
        match self.lock().get_mut(token) {
            Some((created, value)) if created.elapsed() < self.ttl => {
                f(value);
                true
            }
            _ => false,
        }
    }

    pub fn remove(&self, token: &str) -> Option<T> {
        self.lock().remove(token).map(|(_, v)| v)
    }
}

// ── login session ────────────────────────────────────────────────────────────

/// A login session: created pending, authorized at `/login`.
#[derive(Debug, Clone, Copy, Default)]
pub struct LoginSession {
    pub authorized: bool,
}

/// How long a login session stays valid; enrollment is a prompt, one-off act.
pub const SESSION_TTL: Duration = Duration::from_secs(600);

/// `GET /keys/current`: who the key belongs to and what it may do.
pub async fn key_current(State(app): State<AppState>, headers: HeaderMap) -> Response {
    let key = request_key(&headers).unwrap_or_default();
    let write = app.keys.access(key).is_some_and(|a| a.write);
    Json(json!({
        "key": key,
        "userID": app.cfg.user_id,
        "username": app.cfg.username,
        "displayName": app.cfg.username,
        "access": {
            "user": { "library": true, "files": true, "notes": true, "write": write },
            "groups": {}
        },
    }))
    .into_response()
}

/// `DELETE /keys/current`, sent when an account is unlinked in Zotero. Keys
/// are provisioned out of band, so nothing is revoked.
pub async fn key_delete() -> StatusCode {
    StatusCode::NO_CONTENT
}

/// `POST /keys/sessions`: Zotero's "Log In" mints a pending session, opens
/// `loginURL` in the browser, then polls the session for the key.
pub async fn create_session(State(app): State<AppState>) -> Response {
    let (token, _) = app.sessions.insert(LoginSession::default());
    (
        StatusCode::CREATED,
        Json(json!({
            "sessionToken": token,
            "loginURL": format!("{}/login?session={token}", app.cfg.public_url),
        })),
    )
        .into_response()
}

/// `GET /keys/sessions/<token>`: the key once `/login` approved the session;
/// `pending` otherwise, so an unknown or unapproved token never yields one.
pub async fn check_session(
    State(app): State<AppState>,
    UrlPath(token): UrlPath<String>,
) -> Json<Value> {
    if app.sessions.get(&token).is_some_and(|s| s.authorized) {
        app.sessions.remove(&token);
        Json(json!({
            "status": "completed",
            "apiKey": app.keys.login_key(),
            "userID": app.cfg.user_id,
            "username": app.cfg.username,
        }))
    } else {
        Json(json!({ "status": "pending" }))
    }
}

pub async fn cancel_session(
    State(app): State<AppState>,
    UrlPath(token): UrlPath<String>,
) -> StatusCode {
    app.sessions.remove(&token);
    StatusCode::NO_CONTENT
}

/// `GET /login?session=<token>`: the consent page Zotero opens in the browser.
/// It only renders the Approve form; approving takes the form's POST.
pub async fn login_page(
    State(app): State<AppState>,
    Query(params): Query<HashMap<String, String>>,
) -> Response {
    if params.get("approved").is_some_and(|v| v == "1") {
        return ui::login_page(StatusCode::OK, LoginView::Approved);
    }
    let Some(token) = params.get("session").filter(|t| !t.is_empty()) else {
        return ui::login_page(StatusCode::BAD_REQUEST, LoginView::Missing);
    };
    match app.sessions.get(token) {
        None => ui::login_page(StatusCode::NOT_FOUND, LoginView::Expired),
        Some(session) if session.authorized => ui::login_page(StatusCode::OK, LoginView::Approved),
        Some(_) => ui::login_page(StatusCode::OK, LoginView::Pending(token)),
    }
}

/// `POST /login` (form `session=<token>`): the consent page's Approve button.
///
/// The page itself (`GET /login`) never authorizes: a prefetch or a cross-site
/// `GET` would otherwise approve a session an attacker created. The POST also
/// rejects a foreign `Origin` and, with `--login-authorized-user`, requires the
/// SSO proxy in front to forward that identity.
pub async fn login_authorize(
    State(app): State<AppState>,
    headers: HeaderMap,
    Form(form): Form<HashMap<String, String>>,
) -> Response {
    if let Some(origin) = headers.get("origin").and_then(|v| v.to_str().ok())
        && origin.trim_end_matches('/') != public_origin(&app.cfg.public_url)
    {
        return ui::login_page(
            StatusCode::FORBIDDEN,
            LoginView::Refused("The approval was sent from another site."),
        );
    }
    if let Some(want) = &app.cfg.login_authorized_user {
        // oauth2-proxy forwards the identity as X-Forwarded-* when it proxies
        // and as X-Auth-Request-* in nginx auth_request mode.
        let identity = [
            "x-auth-request-email",
            "x-forwarded-email",
            "x-auth-request-user",
            "x-forwarded-user",
        ]
        .iter()
        .find_map(|h| headers.get(*h))
        .and_then(|v| v.to_str().ok());
        if !identity.is_some_and(|got| got.eq_ignore_ascii_case(want)) {
            return ui::login_page(
                StatusCode::FORBIDDEN,
                LoginView::Refused("You are not allowed to approve logins to this library."),
            );
        }
    }
    let Some(token) = form.get("session") else {
        return ui::login_page(StatusCode::BAD_REQUEST, LoginView::Missing);
    };
    if app.sessions.update(token, |s| s.authorized = true) {
        Redirect::to("/login?approved=1").into_response()
    } else {
        ui::login_page(StatusCode::NOT_FOUND, LoginView::Expired)
    }
}

/// `scheme://host[:port]` of the public URL, as browsers send it in `Origin`.
fn public_origin(public_url: &str) -> &str {
    let url = public_url.trim_end_matches('/');
    let after_scheme = url.find("://").map(|i| i + 3).unwrap_or(0);
    match url[after_scheme..].find('/') {
        Some(i) => &url[..after_scheme + i],
        None => url,
    }
}

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

    #[test]
    fn keys_load_by_precedence() {
        let dir = tempfile::tempdir().unwrap();
        let rw = dir.path().join("rw");
        let ro = dir.path().join("ro");
        std::fs::write(&rw, "write-token\n").unwrap();
        std::fs::write(&ro, "read-token").unwrap();
        let source = KeySource {
            manifest: Some(format!("rw:{},ro:{}", rw.display(), ro.display())),
            key: Some("ignored".into()),
            generated: dir.path().join("api-key"),
            ..Default::default()
        };
        let (keys, generated) = Keys::load(&source).unwrap();
        assert!(generated.is_none());
        assert_eq!(keys.access("write-token"), Some(Access { write: true }));
        assert_eq!(keys.access("read-token"), Some(Access { write: false }));
        assert_eq!(keys.access("ignored"), None);
        assert_eq!(keys.login_key(), "write-token");

        let only_ro = KeySource {
            manifest: Some(format!("ro:{}", ro.display())),
            ..Default::default()
        };
        assert!(Keys::load(&only_ro).is_err());
    }

    #[test]
    fn a_generated_key_is_kept() {
        let dir = tempfile::tempdir().unwrap();
        let source = KeySource {
            generated: dir.path().join("sub/api-key"),
            ..Default::default()
        };
        let (first, generated) = Keys::load(&source).unwrap();
        assert!(generated.is_some());
        let (second, generated) = Keys::load(&source).unwrap();
        assert!(generated.is_none());
        assert_eq!(first.login_key(), second.login_key());
        assert_eq!(first.login_key().len(), 64);
    }

    #[test]
    fn tokens_expire() {
        let tokens = Tokens::new(Duration::from_millis(0));
        let (token, _) = tokens.insert(1);
        assert_eq!(tokens.get(&token), None);
        let (_, expired) = tokens.insert(2);
        assert_eq!(expired, vec![(token, 1)]);

        let tokens = Tokens::new(Duration::from_secs(60));
        let (token, _) = tokens.insert(1);
        assert!(tokens.update(&token, |v| *v = 2));
        assert_eq!(tokens.remove(&token), Some(2));
        assert!(!tokens.update(&token, |v| *v = 3));
    }

    #[test]
    fn paths_and_origins() {
        assert_eq!(user_segment("/api/users/1/items"), Some("1"));
        assert_eq!(user_segment("/users/0/items/top"), Some("0"));
        assert_eq!(user_segment("/keys/current"), None);
        assert_eq!(
            public_origin("https://z.example.org/"),
            "https://z.example.org"
        );
        assert_eq!(
            public_origin("https://z.example.org/sub"),
            "https://z.example.org"
        );
        assert_eq!(
            public_origin("http://127.0.0.1:8189"),
            "http://127.0.0.1:8189"
        );
    }
}