br-crypto 0.4.8

use base64::{Engine};
use base64::engine::general_purpose::{URL_SAFE_NO_PAD};
use chrono::Local;
use json::{JsonValue, object};
use crate::pkey::{pri_key_sign, pub_key_verify};
use crate::sha256::{encrypt_hmac_byte};

#[derive(Clone, Debug)]
pub struct Jwt {
    pub headers: Header,
    pub payload: Payload,
}

impl Jwt {
    /// 创建
    /// * alg 加密方式
    pub fn new(alg: Alg) -> Jwt {
        Self {
            headers: Header {
                alg,
                typ: "JWT".to_string(),
            },
            payload: Payload::default(),
        }
    }
    /// 设置载荷
    ///
    /// * iss 签发人
    /// * sub 主题
    /// * aud 受众
    /// * jti 编号
    pub fn payload(&mut self, iss: &str, sub: &str, aud: &str, jti: &str) -> &mut Self {
        self.payload.iss = iss.to_string();
        self.payload.sub = sub.to_string();
        self.payload.aud = aud.to_string();
        self.payload.jti = jti.to_string();
        self
    }
    /// 设置数据
    /// * custom 自定义内容
    pub fn payload_custom(&mut self, custom: JsonValue) -> &mut Self {
        self.payload.custom = custom;
        self
    }
    /// 设置数据
    ///
    /// * exp 过期时间 秒s 10位
    /// * nbf 生效时间 秒s 10位
    /// * iat 签发时间 秒s 10位
    pub fn payload_time(&mut self, exp: i32, nbf: i32, iat: i32) -> &mut Self {
        self.payload.exp = exp;
        self.payload.nbf = nbf;
        self.payload.iat = iat;
        self
    }
    pub fn token(&mut self, key: &str) -> String {
        let header = self.headers.sign();
        let payload = self.payload.sign();
        let sign_text = format!("{header}.{payload}");
        let signature = match self.headers.alg {
            Alg::HS256 => base64_url_encode(encrypt_hmac_byte(key, sign_text.as_bytes()).as_slice()),
            Alg::RSA => pri_key_sign(key, sign_text.as_str(), "RSA"),
            Alg::RS256 => pri_key_sign(key, sign_text.as_str(), "RS256"),
        };
        format!("{header}.{payload}.{signature}")
    }
    /// 验证
    /// * key rsa 为公匙
    pub fn verify(jwt_token: &str, key: &str) -> Result<Self, String> {
        let data: Vec<&str> = jwt_token.split(".").collect();
        if data.len() != 3 {
            return Err("TOKEN格式不正确".to_string());
        }
        let header = data[0];
        let payload = data[1];
        let signature = data[2];
        let sign_text = format!("{header}.{payload}");

        let header = Header::from(header)?;
        let payload = Payload::decode(payload)?;
        let signature_new = match header.alg {
            Alg::HS256 => base64_url_encode(&encrypt_hmac_byte(key, sign_text.as_bytes())),
            Alg::RS256 => {
                if !pub_key_verify(key.to_string().clone(), sign_text.as_str(), signature.to_string(), "RS256") {
                    return Err("签名不正确1".to_string());
                }
                signature.to_string()
            }
            Alg::RSA => {
                let sign = base64_url_decode(signature)?;
                if !pub_key_verify(key.to_string().clone(), sign_text.as_str(), sign, "RSA") {
                    return Err("签名不正确2".to_string());
                }
                signature.to_string()
            }
        };

        if signature_new != signature {
            return Err("签名不一致".to_string());
        }
        let time = Local::now().timestamp() as i32;

        // 生效时间早于当前时间
        if payload.nbf > time {
            return Err(format!("未生效: {}", payload.nbf - time));
        }
        // 过期时间小于当前时间
        if time > payload.exp {
            return Err(format!("已过期: {}", payload.exp));
        }

        Ok(Self {
            headers: header,
            payload,
        })
    }
    pub fn info(jwt_token: &str) -> Result<Jwt, String> {
        let data: Vec<&str> = jwt_token.split(".").collect();
        if data.len() != 3 {
            return Err("TOKEN格式不正确".to_string());
        }
        let header = data[0];
        let payload = data[1];
        let header = Header::from(header)?;
        let payload = Payload::decode(payload)?;
        Ok(Self {
            headers: header,
            payload,
        })
    }
}

#[derive(Clone, Debug)]
pub enum Alg {
    HS256,
    RSA,
    RS256,
}

impl Alg {
    pub fn str(&mut self) -> String {
        match self {
            Alg::HS256 => "HS256",
            Alg::RSA => "RSA",
            Alg::RS256 => "RS256"
        }.to_string()
    }
    pub fn from(name: &str) -> Self {
        match name {
            "HS256" => Alg::HS256,
            "RSA" => Alg::RSA,
            "RS256" => Alg::RS256,
            _ => Alg::HS256
        }
    }
}

#[derive(Clone, Debug)]
pub struct Header {
    /// 签名方式
    alg: Alg,
    /// 令牌类型
    typ: String,
}

impl Header {
    pub fn sign(&mut self) -> String {
        let header = self.json();
        base64_url_encode(header.to_string().as_bytes())
    }
    pub fn set_alg(&mut self, alg: Alg) {
        self.alg = alg;
    }
    pub fn from(base64: &str) -> Result<Header, String> {
        let mut header = Header::default();
        let header_text = base64_url_decode(base64)?;
        header.alg = match json::parse(header_text.as_str()) {
            Ok(e) => {
                Alg::from(e["alg"].as_str().unwrap_or(""))
            }
            Err(e) => {
                return Err(e.to_string());
            }
        };
        Ok(header)
    }
    pub fn json(&mut self) -> JsonValue {
        object! {
            alg:self.alg.str(),
            typ:self.typ.clone()
        }
    }
}

impl Default for Header {
    fn default() -> Self {
        Self {
            alg: Alg::HS256,
            typ: "JWT".to_string(),
        }
    }
}
/// 载荷
#[derive(Clone)]
#[derive(Debug)]
pub struct Payload {
    /// 签发人
    pub iss: String,
    /// 签发时间
    pub iat: i32,
    /// 签发TOKEN
    pub jti: String,

    /// 主题
    pub sub: String,
    /// 受众
    pub aud: String,
    /// 数据
    pub custom: JsonValue,
    /// 生效时间 秒s
    pub nbf: i32,
    /// 过期时间 秒s
    pub exp: i32,

}
impl Default for Payload {
    fn default() -> Self {
        Self {
            iss: "".to_string(),
            //iat: Local::now().timestamp(),
            iat: 0,
            jti: "".to_string(),
            sub: "".to_string(),
            aud: "".to_string(),
            custom: JsonValue::Null,
            nbf: 0,
            exp: 0,
        }
    }
}

impl Payload {
    pub fn sign(&mut self) -> String {
        let payload = self.json();
        base64_url_encode(payload.to_string().as_bytes())
    }

    pub fn decode(payload: &str) -> Result<Payload, String> {
        let payload_text = base64_url_decode(payload)?;
        match json::parse(payload_text.as_str()) {
            Ok(payload) => {
                Ok(Self {
                    iss: payload["iss"].to_string(),
                    iat: payload["iat"].as_i32().unwrap_or(0),
                    jti: payload["jti"].to_string(),
                    sub: payload["sub"].to_string(),
                    aud: payload["aud"].to_string(),
                    custom: payload["custom"].clone(),
                    nbf: payload["nbf"].as_i32().unwrap_or(0),
                    exp: payload["exp"].as_i32().unwrap_or(0),
                })
            }
            Err(_) => Err("载荷错误".to_string())
        }
    }
    pub fn json(&mut self) -> JsonValue {
        object! {
            iss: self.iss.clone(),
            iat:  self.iat,
            jti:  self.jti.clone(),
            sub: self.sub.clone(),
            aud:  self.aud.clone(),
            custom:  self.custom.clone(),
            nbf: self.nbf,
            exp:  self.exp
        }
    }
}


/// 解码
pub fn base64_url_decode(base64: &str) -> Result<String, String> {
    unsafe {
        match URL_SAFE_NO_PAD.decode(base64) {
            Ok(e) => Ok(String::from_utf8_unchecked(e)),
            Err(e) => Err(e.to_string())
        }
    }
}
/// 加密
pub fn base64_url_encode(data: &[u8]) -> String {
    URL_SAFE_NO_PAD.encode(data)
}

#[test]
fn tests() {
    let mut jwt = Jwt::new(Alg::HS256);
    jwt.payload("1", "1", "1", "1");
    jwt.payload_custom(object! {"name" : "John","admin" : true});
    jwt.payload_time(1751572000, 1751472000, 1751472000);
    let data = jwt.token("1");
    println!("{data}");

    let jwt_new = Jwt::verify(data.as_str(), "1");
    println!("{jwt_new:#?}");

    let jwt_new = Jwt::info(data.as_str());
    println!("{jwt_new:#?}");
}

#[cfg(test)]
mod test {
    use crate::pkey::create_pkey_pkcs8;
    use crate::jwt::Alg;
    use crate::jwt::Jwt;
    use json::object;

    #[test]
    fn test_rsa() {
        let mut jwt = Jwt::new(Alg::RS256);
        jwt.payload("1", "1", "1", "1");
        let (pub_key, pri_key) = create_pkey_pkcs8();

        jwt.payload_custom(object! {"name" : "John","admin" : true});
        jwt.payload_time(1751572000, 1751472000, 1751472000);
        let data = jwt.token(pri_key.as_str());
        println!("{data}");

        let jwt_new = Jwt::verify(data.as_str(), pub_key.as_str());
        println!("{jwt_new:#?}");

        let jwt_new = Jwt::info(data.as_str());
        println!("{jwt_new:#?}");
    }
}