//! The hub as an OpenID Connect issuer for jobs (hub spec §13.5,

//! oxygenhub#103): a job that declares `identity "AUDIENCE" { mount-path … }`

//! gets a signed JWT naming the repository, bookmark, run and job, so a

//! cloud or a vault that trusts this hub can let the job in without a

//! long-lived secret. The key is RSA 2048 in `keys/oidc.pem`, made at

//! first use; `/.well-known/openid-configuration` and

//! `/.well-known/jwks.json` publish what verifiers need.

use rsa::pkcs1v15::{SigningKey, VerifyingKey};

use rsa::pkcs8::{DecodePrivateKey, EncodePrivateKey};

use rsa::signature::{SignatureEncoding, Signer, Verifier};

use rsa::traits::PublicKeyParts;

use rsa::RsaPrivateKey;

use sha2::Sha256;

use std::path::Path;

use crate::api::ApiError;

use crate::Shared;

use axum::extract::State;

use axum::http::StatusCode;

use axum::response::{IntoResponse, Response};

use axum::Json;

/// The longest a job token lives, whatever the job's limit.

pub const MAX_TTL_S: u64 = 6 * 3600;

pub struct Issuer {

    signing: SigningKey<Sha256>,

    verifying: VerifyingKey<Sha256>,

    /// The key id: the first 16 hex of BLAKE3 over the public key's

    /// modulus, as the JWKS names it.

    pub kid: String,

    n: String,

    e: String,

}

/// What a token says about its job.

pub struct JobClaims<'a> {

    pub issuer: &'a str,

    pub audience: &'a str,

    pub owner: &'a str,

    pub name: &'a str,

    pub visibility: &'a str,

    pub bookmark: &'a str,

    pub run_number: i64,

    pub job: &'a str,

    pub state: &'a str,

    pub runner: &'a str,

    pub actor: &'a str,

    /// The environment the job deploys to, only once it was approved

    /// (oxygenflow#59).

    pub environment: Option<&'a str>,

    pub now_ms: i64,

    pub ttl_s: u64,

}

impl Issuer {

    /// The key from `<keys>/oidc.pem`, or a new 2048-bit one written

    /// there (mode 0600).

    pub fn load_or_create(keys: &Path) -> anyhow::Result<Issuer> {

        std::fs::create_dir_all(keys)?;

        let path = keys.join("oidc.pem");

        let key = match std::fs::read_to_string(&path) {

            Ok(pem) => RsaPrivateKey::from_pkcs8_pem(&pem).map_err(|e| anyhow::anyhow!("{}: {e}", path.display()))?,

            Err(_) => {

                let key = RsaPrivateKey::new(&mut rand::rngs::OsRng, 2048)?;

                let pem = key.to_pkcs8_pem(rsa::pkcs8::LineEnding::LF)?;

                std::fs::write(&path, pem.as_bytes())?;

                #[cfg(unix)]

                {

                    use std::os::unix::fs::PermissionsExt;

                    std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))?;

                }

                key

            }

        };

        Ok(Issuer::from_key(key))

    }

    fn from_key(key: RsaPrivateKey) -> Issuer {

        let public = key.to_public_key();

        let n_bytes = public.n().to_bytes_be();

        let kid = blake3::hash(&n_bytes).to_hex()[..16].to_string();

        Issuer {

            n: b64url(&n_bytes),

            e: b64url(&public.e().to_bytes_be()),

            verifying: VerifyingKey::new(public),

            signing: SigningKey::new(key),

            kid,

        }

    }

    /// `/.well-known/openid-configuration`.

    pub fn discovery(&self, issuer: &str) -> serde_json::Value {

        serde_json::json!({

            "issuer": issuer,

            "jwks_uri": format!("{issuer}/.well-known/jwks.json"),

            "response_types_supported": ["id_token"],

            "subject_types_supported": ["public"],

            "id_token_signing_alg_values_supported": ["RS256"],

            "claims_supported": ["iss", "sub", "aud", "exp", "iat", "nbf", "jti", "repository", "repository_owner",

                "repository_visibility", "bookmark", "run_number", "job", "state", "runner", "actor", "environment"],

            "scopes_supported": ["openid"],

        })

    }

    /// `/.well-known/jwks.json`.

    pub fn jwks(&self) -> serde_json::Value {

        serde_json::json!({ "keys": [{ "kty": "RSA", "use": "sig", "alg": "RS256", "kid": self.kid, "n": self.n, "e": self.e }] })

    }

    /// A signed token for a job: `sub` is `repo:OWNER/NAME:bookmark:B`,

    /// so a trust policy can match a repository and a bookmark.

    pub fn mint(&self, c: &JobClaims<'_>) -> String {

        let iat = c.now_ms / 1000;

        let exp = iat + c.ttl_s.clamp(60, MAX_TTL_S) as i64;

        let header = serde_json::json!({ "alg": "RS256", "typ": "JWT", "kid": self.kid });

        let mut payload = serde_json::json!({

            "iss": c.issuer,

            "sub": format!("repo:{}/{}:bookmark:{}", c.owner, c.name, c.bookmark),

            "aud": c.audience,

            "iat": iat,

            "nbf": iat - 30,

            "exp": exp,

            "jti": crate::auth::random_hex(16),

            "repository": format!("{}/{}", c.owner, c.name),

            "repository_owner": c.owner,

            "repository_visibility": c.visibility,

            "bookmark": c.bookmark,

            "run_number": c.run_number,

            "job": c.job,

            "state": c.state,

            "runner": c.runner,

            "actor": c.actor,

        });

        if let Some(e) = c.environment {

            payload["environment"] = serde_json::Value::from(e);

        }

        let signing_input = format!("{}.{}", b64url(header.to_string().as_bytes()), b64url(payload.to_string().as_bytes()));

        let sig = self.signing.sign(signing_input.as_bytes());

        format!("{signing_input}.{}", b64url(&sig.to_bytes()))

    }

    /// A signed token for the hub itself acting as `subject` (a backup's

    /// `backup:OWNER`, oxygenhub#169), for `audience`, living `ttl_s`.

    pub fn mint_for(&self, issuer: &str, subject: &str, audience: &str, now_ms: i64, ttl_s: u64) -> String {

        let iat = now_ms / 1000;

        let header = serde_json::json!({ "alg": "RS256", "typ": "JWT", "kid": self.kid });

        let payload = serde_json::json!({

            "iss": issuer,

            "sub": subject,

            "aud": audience,

            "iat": iat,

            "nbf": iat - 30,

            "exp": iat + ttl_s.clamp(60, MAX_TTL_S) as i64,

            "jti": crate::auth::random_hex(16),

        });

        let signing_input = format!("{}.{}", b64url(header.to_string().as_bytes()), b64url(payload.to_string().as_bytes()));

        let sig = self.signing.sign(signing_input.as_bytes());

        format!("{signing_input}.{}", b64url(&sig.to_bytes()))

    }

    /// The claims of a token this hub signed, when its signature holds,

    /// it is for `audience` and it has not expired at `now_ms`.

    pub fn verify(&self, token: &str, audience: &str, now_ms: i64) -> Option<serde_json::Value> {

        let mut parts = token.split('.');

        let (h, p, s) = (parts.next()?, parts.next()?, parts.next()?);

        if parts.next().is_some() {

            return None;

        }

        let header: serde_json::Value = serde_json::from_slice(&b64url_decode(h)?).ok()?;

        if header["alg"] != "RS256" || header["kid"] != self.kid.as_str() {

            return None;

        }

        let sig = rsa::pkcs1v15::Signature::try_from(b64url_decode(s)?.as_slice()).ok()?;

        self.verifying.verify(format!("{h}.{p}").as_bytes(), &sig).ok()?;

        let claims: serde_json::Value = serde_json::from_slice(&b64url_decode(p)?).ok()?;

        let now = now_ms / 1000;

        (claims["aud"] == audience && claims["exp"].as_i64()? > now && claims["nbf"].as_i64()? <= now).then_some(claims)

    }

}

/// `text` with every token this hub signed replaced by `***`: a job

/// that prints its identity does not leave it in the log.

pub fn mask(text: &str, kid: &str) -> String {

    if !text.contains("eyJ") {

        return text.to_string();

    }

    let is_tok = |c: char| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.';

    let mut out = String::with_capacity(text.len());

    let mut rest = text;

    while let Some(i) = rest.find("eyJ") {

        out.push_str(&rest[..i]);

        let end = rest[i..].find(|c: char| !is_tok(c)).map(|e| i + e).unwrap_or(rest.len());

        let cand = &rest[i..end];

        let ours = cand

            .split('.')

            .next()

            .and_then(b64url_decode)

            .and_then(|h| serde_json::from_slice::<serde_json::Value>(&h).ok())

            .is_some_and(|h| h["kid"] == kid);

        if ours && cand.matches('.').count() == 2 {

            out.push_str("***");

        } else {

            out.push_str(cand);

        }

        rest = &rest[end..];

    }

    out.push_str(rest);

    out

}

fn b64url(bytes: &[u8]) -> String {

    const T: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";

    let mut out = String::with_capacity(bytes.len().div_ceil(3) * 4);

    for chunk in bytes.chunks(3) {

        let n = chunk.iter().fold(0u32, |acc, b| (acc << 8) | *b as u32) << (8 * (3 - chunk.len()));

        for i in 0..=chunk.len() {

            out.push(T[((n >> (18 - 6 * i)) & 63) as usize] as char);

        }

    }

    out

}

fn b64url_decode(s: &str) -> Option<Vec<u8>> {

    let val = |c: u8| -> Option<u32> {

        Some(match c {

            b'A'..=b'Z' => (c - b'A') as u32,

            b'a'..=b'z' => (c - b'a') as u32 + 26,

            b'0'..=b'9' => (c - b'0') as u32 + 52,

            b'-' => 62,

            b'_' => 63,

            _ => return None,

        })

    };

    let bytes = s.trim_end_matches('=').as_bytes();

    let mut out = Vec::with_capacity(bytes.len() * 3 / 4);

    for chunk in bytes.chunks(4) {

        if chunk.len() == 1 {

            return None;

        }

        let mut n = 0u32;

        for (i, c) in chunk.iter().enumerate() {

            n |= val(*c)? << (18 - 6 * i);

        }

        out.push((n >> 16) as u8);

        if chunk.len() > 2 {

            out.push((n >> 8) as u8);

        }

        if chunk.len() > 3 {

            out.push(n as u8);

        }

    }

    Some(out)

}

/// `GET /.well-known/openid-configuration`.

pub async fn discovery(State(app): State<Shared>) -> Result<Response, ApiError> {

    let issuer = app.oidc().await.map_err(|e| ApiError(StatusCode::INTERNAL_SERVER_ERROR, format!("oidc: {e}")))?;

    Ok(([(axum::http::header::CACHE_CONTROL, "public, max-age=3600")], Json(issuer.discovery(app.config.base_url()))).into_response())

}

/// `GET /.well-known/jwks.json`.

pub async fn jwks(State(app): State<Shared>) -> Result<Response, ApiError> {

    let issuer = app.oidc().await.map_err(|e| ApiError(StatusCode::INTERNAL_SERVER_ERROR, format!("oidc: {e}")))?;

    Ok(([(axum::http::header::CACHE_CONTROL, "public, max-age=3600")], Json(issuer.jwks())).into_response())

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn tokens_verify_against_the_published_key_and_are_masked() {

        let dir = tempfile::tempdir().unwrap();

        let issuer = Issuer::load_or_create(dir.path()).unwrap();

        let again = Issuer::load_or_create(dir.path()).unwrap();

        assert_eq!(issuer.kid, again.kid, "the key is kept");

        let claims = JobClaims {

            issuer: "https://hub.example",

            audience: "sts.amazonaws.com",

            owner: "andy",

            name: "demo",

            visibility: "public",

            bookmark: "main",

            run_number: 7,

            job: "deploy",

            state: "ab".repeat(32).leak(),

            runner: "earth-1",

            actor: "andy",

            environment: Some("prod"),

            now_ms: 1_800_000_000_000,

            ttl_s: 900,

        };

        let token = issuer.mint(&claims);

        assert_eq!(token.matches('.').count(), 2);

        let got = again.verify(&token, "sts.amazonaws.com", 1_800_000_500_000).expect("verifies");

        assert_eq!(got["sub"], "repo:andy/demo:bookmark:main");

        assert_eq!(got["run_number"], 7);

        assert_eq!(got["environment"], "prod");

        let plain = issuer.mint(&JobClaims { environment: None, ..claims });

        assert!(again.verify(&plain, "sts.amazonaws.com", 1_800_000_500_000).unwrap().get("environment").is_none(), "only when approved");

        assert_eq!(got["exp"].as_i64().unwrap() - got["iat"].as_i64().unwrap(), 900);

        assert!(again.verify(&token, "other", 1_800_000_500_000).is_none(), "another audience");

        assert!(again.verify(&token, "sts.amazonaws.com", 1_800_001_000_000).is_none(), "expired");

        let other = Issuer::load_or_create(tempfile::tempdir().unwrap().path()).unwrap();

        assert!(other.verify(&token, "sts.amazonaws.com", 1_800_000_500_000).is_none(), "another key");

        let tampered = format!("{}x", token);

        assert!(again.verify(&tampered, "sts.amazonaws.com", 1_800_000_500_000).is_none());

        let jwks = issuer.jwks();

        assert_eq!(jwks["keys"][0]["kid"], issuer.kid.as_str());

        assert_eq!(b64url_decode(jwks["keys"][0]["e"].as_str().unwrap()).unwrap(), vec![1, 0, 1]);

        assert_eq!(mask(&format!("token={token} done eyJnot.a.token"), &issuer.kid), "token=*** done eyJnot.a.token");

        assert_eq!(b64url_decode(&b64url(b"hello?")).unwrap(), b"hello?");

    }

}