nintypes 0.2.11

use jsonwebtoken::{decode, encode, Algorithm, DecodingKey, EncodingKey, Header, TokenData, Validation};
use serde::{Deserialize, Serialize};
use std::{
    fs,
    time::{SystemTime, UNIX_EPOCH},
};
use tower_cookies::{cookie::time::OffsetDateTime, Cookie};

pub fn load_encoding_key() -> EncodingKey {
    let private_key_pem = fs::read_to_string("jwt_private_key.pem").expect("Unable to read jwt private key");
    EncodingKey::from_rsa_pem(private_key_pem.as_bytes()).expect("Invalid jwt private key")
}

pub fn load_decoding_key() -> DecodingKey {
    let public_key_pem = fs::read_to_string("jwt_public_key.pem").expect("Unable to read jwt public key");
    DecodingKey::from_rsa_pem(public_key_pem.as_bytes()).expect("Invalid jwt public key")
}

pub struct ClaimsController {
    encode_key: EncodingKey,
    decode_key: DecodingKey,
    expiration_secs: u64,
    algorithm: Algorithm,
    frontend_domain: String,
}

impl ClaimsController {
    pub fn new(encode_key: EncodingKey, decode_key: DecodingKey, expiration_secs: u64, frontend_domain: String) -> Self {
        Self {
            encode_key,
            decode_key,
            expiration_secs,
            frontend_domain,
            algorithm: Algorithm::RS256,
        }
    }

    pub fn encode(&self, claims: &Claims) -> Result<String, Error> {
        encode(&Header::new(self.algorithm), claims, &self.encode_key).map_err(Error::FailedEncode)
    }

    pub fn now_secs() -> u64 {
        SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs()
    }

    pub fn decode(raw_claims: &str, algorithm: Algorithm, decode_key: &DecodingKey) -> Result<TokenData<Claims>, Error> {
        let claims = decode::<Claims>(raw_claims, decode_key, &Validation::new(algorithm)).map_err(Error::FailedDecode)?;
        let exp = claims.claims.exp as u64;

        if exp < Self::now_secs() {
            return Err(Error::Expired);
        }

        Ok(claims)
    }

    pub fn verify(&self, raw_claims: &str) -> Result<TokenData<Claims>, Error> {
        Self::decode(raw_claims, self.algorithm, &self.decode_key)
    }

    pub fn expiration(&self) -> u64 {
        Self::now_secs() + self.expiration_secs
    }

    pub fn generate_cookie(&self, claims: Claims) -> Result<Cookie<'static>, Error> {
        let raw_token = self.encode(&claims)?;
        let cookie = Cookie::build(("access_token", raw_token))
            .path("/")
            .domain(self.frontend_domain.clone())
            .same_site(tower_cookies::cookie::SameSite::None)
            .secure(true)
            .expires(OffsetDateTime::from_unix_timestamp(claims.exp as i64)?)
            .build();

        Ok(cookie)
    }

    pub fn generate_empty_cookie(&self) -> Result<Cookie<'static>, Error> {
        let cookie = Cookie::build(("access_token", ""))
            .path("/")
            .domain(self.frontend_domain.clone())
            .same_site(tower_cookies::cookie::SameSite::None)
            .secure(true)
            .expires(OffsetDateTime::from_unix_timestamp(0)?)
            .build();

        Ok(cookie)
    }
}

#[derive(Debug, Serialize, Deserialize, Clone)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
pub struct Claims {
    pub address: String,
    pub exp: usize,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub admin: Option<bool>,
}

impl Claims {
    pub fn is_admin(&self) -> bool {
        self.admin.unwrap_or_default()
    }

    #[allow(clippy::result_large_err)]
    pub fn check_admin(&self) -> Result<(), http::Response<String>> {
        if self.is_admin() {
            return Ok(());
        }
        Err(http::Response::builder()
            .status(http::StatusCode::UNAUTHORIZED)
            .body(serde_json::to_string("user is not admin").unwrap())
            .unwrap())
    }
}

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("Token has expired")]
    Expired,
    #[error("Failed to encode token: {0}")]
    FailedEncode(jsonwebtoken::errors::Error),
    #[error("Failed to decode token: {0}")]
    FailedDecode(jsonwebtoken::errors::Error),
    #[error("Failed to create offset date time: {0}")]
    OffsetDateTime(#[from] tower_cookies::cookie::time::error::ComponentRange),
}