json-proof-token 0.4.1

// Copyright 2023 Fondazione LINKS

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::errors::CustomError;
use serde::{Deserialize, Serialize};
use zkryptium::{
    keys::pair::KeyPair,
    schemes::algorithms::{BbsBls12381Sha256, BbsBls12381Shake256},
};

use super::{
    alg_parameters::{Algorithm, JwkAlgorithmParameters, JwkEllipticCurveKeyParameters},
    types::KeyPairSubtype,
};

/// JWK parameters defined at https://datatracker.ietf.org/doc/html/rfc7517#section-4
#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
pub struct Jwk {
    /// ID of the key
    #[serde(skip_serializing_if = "Option::is_none", default)]
    pub kid: Option<String>,

    /// The intended use of the public key
    #[serde(rename = "use", skip_serializing_if = "Option::is_none", default)]
    pub pk_use: Option<PKUse>,

    /// The "key_ops" (key operations) parameter identifies the operation(s) for which the key is intended to be used

    /// The "use" and "key_ops" JWK members SHOULD NOT be used together;
    /// however, if both are used, the information they convey MUST be
    /// consistent.  Applications should specify which of these members they
    /// use, if either is to be used by the application.
    #[serde(skip_serializing_if = "Option::is_none", default)]
    pub key_ops: Option<Vec<KeyOps>>,

    /// The algorithm intended to be used
    #[serde(skip_serializing_if = "Option::is_none", default)]
    pub alg: Option<Algorithm>,

    /// X.509 Public key cerfificate URL.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub x5u: Option<String>,

    /// X.509 public key certificate chain.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub x5c: Option<Vec<String>>,

    /// X.509 Certificate thumbprint.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub x5t: Option<String>,

    //Key parameters based on the algorithm
    #[serde(flatten)]
    pub key_params: JwkAlgorithmParameters,
}

impl Jwk {
    pub fn generate(key_type: KeyPairSubtype) -> Result<Self, CustomError> {
        match key_type {
            KeyPairSubtype::BLS12381G2Sha256 => {
                let keypair = KeyPair::<BbsBls12381Sha256>::random()
                    .map_err(|_| CustomError::JwkGenerationError("Keygen failed".to_owned()))?;
                let pk = keypair.public_key();
                let (x, y) = pk.to_coordinates();
                let sk = keypair.private_key().to_bytes();
                let okp_params = JwkEllipticCurveKeyParameters::new(
                    super::curves::EllipticCurveTypes::BLS12381G2,
                    &x,
                    &y,
                    Some(&sk),
                );
                let jwk_params = JwkAlgorithmParameters::EllipticCurve(okp_params);
                Ok(Self {
                    kid: None,
                    pk_use: None,
                    key_ops: None,
                    alg: None,
                    x5u: None,
                    x5c: None,
                    x5t: None,
                    key_params: jwk_params,
                })
            }
            KeyPairSubtype::BLS12381G2Shake256 => {
                let keypair = KeyPair::<BbsBls12381Shake256>::random()
                    .map_err(|_| CustomError::JwkGenerationError("Keygen failed".to_owned()))?;
                let pk = keypair.public_key();
                let (x, y) = pk.to_coordinates();
                let sk = keypair.private_key().to_bytes();
                let okp_params = JwkEllipticCurveKeyParameters::new(
                    super::curves::EllipticCurveTypes::BLS12381G2,
                    &x,
                    &y,
                    Some(&sk),
                );
                let jwk_params = JwkAlgorithmParameters::EllipticCurve(okp_params);
                Ok(Self {
                    kid: None,
                    pk_use: None,
                    key_ops: None,
                    alg: None,
                    x5u: None,
                    x5c: None,
                    x5t: None,
                    key_params: jwk_params,
                })
            }
        }
    }

    pub fn set_kid(&mut self, kid: &str) {
        self.kid = Some(kid.to_string());
    }

    pub fn set_pk_use(&mut self, pk_use: PKUse) {
        self.pk_use = Some(pk_use);
    }

    pub fn set_key_ops(&mut self, key_ops: Vec<KeyOps>) {
        self.key_ops = Some(key_ops);
    }

    pub fn set_alg(&mut self, alg: Algorithm) {
        self.alg = Some(alg);
    }

    pub fn set_x5u(&mut self, x5u: &str) {
        self.x5u = Some(x5u.to_string());
    }

    pub fn set_x5c(&mut self, x5c: Vec<&str>) {
        self.x5c = Some(x5c.iter().map(|v| v.to_string()).collect());
    }

    pub fn set_x5t(&mut self, x5t: &str) {
        self.x5t = Some(x5t.to_string());
    }

    /// Returns `true` if _all_ private key components of the key are unset, `false` otherwise.
    pub fn is_public(&self) -> bool {
        self.key_params.is_public()
    }

    /// Returns `true` if _all_ private key components of the key are set, `false` otherwise.
    pub fn is_private(&self) -> bool {
        self.key_params.is_private()
    }

    pub fn from_key_params(key_params: JwkAlgorithmParameters) -> Self {
        let params: JwkAlgorithmParameters = key_params.into();
        Self {
            kid: None,
            pk_use: None,
            key_ops: None,
            alg: None,
            x5u: None,
            x5c: None,
            x5t: None,
            key_params: params,
        }
    }

    pub fn to_public(&self) -> Option<Jwk> {
        let mut public: Jwk = Jwk::from_key_params(self.key_params.to_public()?);

        if let Some(value) = &self.kid {
            public.set_kid(value);
        }

        if let Some(value) = self.pk_use {
            public.set_pk_use(value);
        }

        if let Some(value) = &self.key_ops {
            public.set_key_ops(value.iter().map(|op| op.inverse()).collect());
        }

        if let Some(value) = self.alg {
            public.set_alg(value);
        }

        if let Some(value) = &self.x5u {
            public.set_x5u(value);
        }

        if let Some(value) = &self.x5c {
            public.set_x5c(value.iter().map(|x| x.as_str()).collect());
        }

        if let Some(value) = &self.x5t {
            public.set_x5t(value);
        }

        Some(public)
    }
}

#[derive(
    Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, serde::Deserialize, serde::Serialize,
)]
#[serde(rename_all = "camelCase")]
pub enum KeyOps {
    /// Compute digital signature or MAC.
    Sign,
    /// Verify digital signature or MAC.
    Verify,
    /// Encrypt content.
    Encrypt,
    /// Decrypt content and validate decryption, if applicable.
    Decrypt,
    /// Encrypt key.
    WrapKey,
    /// Decrypt key and validate decryption, if applicable.
    UnwrapKey,
    /// Derive key.
    DeriveKey,
    /// Derive bits not to be used as a key.
    DeriveBits,

    /// Generate a proof
    ProofGeneration,
    ///Verify a proof
    ProofVerification,
}

impl KeyOps {
    pub const fn inverse(&self) -> Self {
        match self {
            Self::Sign => Self::Verify,
            Self::Verify => Self::Sign,
            Self::Encrypt => Self::Decrypt,
            Self::Decrypt => Self::Encrypt,
            Self::WrapKey => Self::UnwrapKey,
            Self::UnwrapKey => Self::WrapKey,
            Self::DeriveKey => Self::DeriveKey,
            Self::DeriveBits => Self::DeriveBits,
            Self::ProofGeneration => Self::ProofVerification,
            Self::ProofVerification => Self::ProofVerification,
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum PKUse {
    #[serde(rename = "sig")]
    Signature,
    #[serde(rename = "enc")]
    Encryption,
    #[serde(rename = "proof")]
    Proof,
}