blsful 4.0.0

use crate::helpers::{KEYGEN_SALT, get_crypto_rng};
use crate::impls::inner_types::*;
use crate::*;
use core::fmt::{self, Formatter};
use rand::CryptoRng;
use rand::RngExt;
use serde::de::{SeqAccess, Visitor};
use subtle::CtOption;
use vsss_rs::*;

/// Number of bytes needed to represent the secret key
pub const SECRET_KEY_BYTES: usize = 32;

/// A BLS secret key implementation
///
/// This does not expose the underlying curve
/// and signature scheme and can be used in situations where the specific
/// implementation is not known at compile time and where trait objects
/// are desirable but cannot be used because they lack the `Sized` trait.
/// The downside is that the type is indicated by a byte or string
/// for serialization and deserialization. If this is not desirable,
/// then use [`SecretKey<C>`](struct.SecretKey.html) instead.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum SecretKeyEnum {
    /// A secret key for signatures in G1 and public keys in G2
    G1(SecretKey<Bls12381G1Impl>),
    /// A secret key for signatures in G2 and public keys in G1
    G2(SecretKey<Bls12381G2Impl>),
}

impl Serialize for SecretKeyEnum {
    fn serialize<S: Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
        match self {
            SecretKeyEnum::G1(sk) => (Bls12381::G1, sk).serialize(s),
            SecretKeyEnum::G2(sk) => (Bls12381::G2, sk).serialize(s),
        }
    }
}

impl<'de> Deserialize<'de> for SecretKeyEnum {
    fn deserialize<D: Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
        struct SecretKeyEnumVisitor;

        impl<'de> Visitor<'de> for SecretKeyEnumVisitor {
            type Value = SecretKeyEnum;

            fn expecting(&self, f: &mut Formatter<'_>) -> fmt::Result {
                write!(f, "a tuple of the type and secret key")
            }

            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
            where
                A: SeqAccess<'de>,
            {
                let ee = seq
                    .next_element::<Bls12381>()?
                    .ok_or_else(|| serde::de::Error::invalid_length(0, &self))?;
                match ee {
                    Bls12381::G1 => {
                        let sk = seq
                            .next_element::<SecretKey<Bls12381G1Impl>>()?
                            .ok_or_else(|| serde::de::Error::invalid_length(1, &self))?;
                        Ok(SecretKeyEnum::G1(sk))
                    }
                    Bls12381::G2 => {
                        let sk = seq
                            .next_element::<SecretKey<Bls12381G2Impl>>()?
                            .ok_or_else(|| serde::de::Error::invalid_length(1, &self))?;
                        Ok(SecretKeyEnum::G2(sk))
                    }
                }
            }
        }
        d.deserialize_tuple(2, SecretKeyEnumVisitor)
    }
}

impl Default for SecretKeyEnum {
    fn default() -> Self {
        Self::G1(SecretKey(Scalar::default()))
    }
}

impl From<&SecretKeyEnum> for Vec<u8> {
    fn from(value: &SecretKeyEnum) -> Self {
        let (tt, output) = match value {
            SecretKeyEnum::G1(sk) => (Bls12381::G1, Vec::from(sk)),
            SecretKeyEnum::G2(sk) => (Bls12381::G2, Vec::from(sk)),
        };
        typed_bytes(tt, output)
    }
}

impl TryFrom<&[u8]> for SecretKeyEnum {
    type Error = BlsError;

    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        let (ee, value) = Bls12381::split_typed_bytes(value)?;
        match ee {
            Bls12381::G1 => {
                let sk = SecretKey::<Bls12381G1Impl>::try_from(value)?;
                Ok(SecretKeyEnum::G1(sk))
            }
            Bls12381::G2 => {
                let sk = SecretKey::<Bls12381G2Impl>::try_from(value)?;
                Ok(SecretKeyEnum::G2(sk))
            }
        }
    }
}

impl_from_derivatives!(SecretKeyEnum);

impl SecretKeyEnum {
    /// Return the concrete BLS12-381 signature group for this secret key.
    pub fn curve(&self) -> Bls12381 {
        match self {
            Self::G1(_) => Bls12381::G1,
            Self::G2(_) => Bls12381::G2,
        }
    }

    /// Create a new random secret key
    pub fn new(t: Bls12381) -> Self {
        match t {
            Bls12381::G1 => SecretKeyEnum::G1(SecretKey::new()),
            Bls12381::G2 => SecretKeyEnum::G2(SecretKey::new()),
        }
    }

    /// Compute a secret key from a hash
    pub fn from_hash<B: AsRef<[u8]>>(t: Bls12381, data: B) -> Self {
        match t {
            Bls12381::G1 => SecretKeyEnum::G1(SecretKey::from_hash(data)),
            Bls12381::G2 => SecretKeyEnum::G2(SecretKey::from_hash(data)),
        }
    }

    /// Compute a secret key from a CS-PRNG
    pub fn random(t: Bls12381, rng: impl CryptoRng) -> Self {
        match t {
            Bls12381::G1 => SecretKeyEnum::G1(SecretKey::random(rng)),
            Bls12381::G2 => SecretKeyEnum::G2(SecretKey::random(rng)),
        }
    }

    /// Get the big-endian byte representation of this key
    pub fn to_be_bytes(&self) -> Vec<u8> {
        let (t, output) = match self {
            SecretKeyEnum::G1(sk) => (Bls12381::G1, sk.to_be_bytes()),
            SecretKeyEnum::G2(sk) => (Bls12381::G2, sk.to_be_bytes()),
        };
        typed_bytes(t, output)
    }

    /// Get the little-endian byte representation of this key
    pub fn to_le_bytes(&self) -> Vec<u8> {
        let (t, output) = match self {
            SecretKeyEnum::G1(sk) => (Bls12381::G1, sk.to_le_bytes()),
            SecretKeyEnum::G2(sk) => (Bls12381::G2, sk.to_le_bytes()),
        };
        typed_bytes(t, output)
    }

    /// Convert a big-endian representation of the secret key.
    pub fn from_be_bytes(bytes: &[u8]) -> CtOption<Self> {
        let (t, bytes) = match Bls12381::split_typed_bytes(bytes) {
            Ok(parts) => parts,
            Err(_) => return CtOption::new(Self::default(), Choice::from(0u8)),
        };
        match bytes.try_into() {
            Ok(sk) => match t {
                Bls12381::G1 => {
                    let ct_sk = SecretKey::from_be_bytes(&sk);
                    let choice = ct_sk.is_some();
                    let val = SecretKeyEnum::G1(Option::from(ct_sk).unwrap_or_default());
                    CtOption::new(val, choice)
                }
                Bls12381::G2 => {
                    let ct_sk = SecretKey::from_be_bytes(&sk);
                    let choice = ct_sk.is_some();
                    let val = SecretKeyEnum::G2(Option::from(ct_sk).unwrap_or_default());
                    CtOption::new(val, choice)
                }
            },
            Err(_) => CtOption::new(Self::default(), Choice::from(0u8)),
        }
    }

    /// Convert a little-endian representation of the secret key.
    pub fn from_le_bytes(bytes: &[u8]) -> CtOption<Self> {
        let (t, bytes) = match Bls12381::split_typed_bytes(bytes) {
            Ok(parts) => parts,
            Err(_) => return CtOption::new(Self::default(), Choice::from(0u8)),
        };
        match bytes.try_into() {
            Ok(sk) => match t {
                Bls12381::G1 => {
                    let ct_sk = SecretKey::from_le_bytes(&sk);
                    let choice = ct_sk.is_some();
                    let val = SecretKeyEnum::G1(Option::from(ct_sk).unwrap_or_default());
                    CtOption::new(val, choice)
                }
                Bls12381::G2 => {
                    let ct_sk = SecretKey::from_le_bytes(&sk);
                    let choice = ct_sk.is_some();
                    let val = SecretKeyEnum::G2(Option::from(ct_sk).unwrap_or_default());
                    CtOption::new(val, choice)
                }
            },
            Err(_) => CtOption::new(Self::default(), Choice::from(0u8)),
        }
    }

    /// Compute the public key for this secret key.
    pub fn public_key(&self) -> PublicKeyEnum {
        match self {
            SecretKeyEnum::G1(sk) => PublicKeyEnum::G1(sk.public_key()),
            SecretKeyEnum::G2(sk) => PublicKeyEnum::G2(sk.public_key()),
        }
    }

    /// Create a proof of possession for this secret key.
    pub fn proof_of_possession(&self) -> BlsResult<ProofOfPossessionEnum> {
        match self {
            SecretKeyEnum::G1(sk) => sk.proof_of_possession().map(ProofOfPossessionEnum::G1),
            SecretKeyEnum::G2(sk) => sk.proof_of_possession().map(ProofOfPossessionEnum::G2),
        }
    }

    /// Sign a message with this secret key using the specified scheme.
    pub fn sign(&self, scheme: SignatureSchemes, msg: &[u8]) -> BlsResult<SignatureEnum> {
        match self {
            SecretKeyEnum::G1(sk) => sk.sign(scheme, msg).map(SignatureEnum::G1),
            SecretKeyEnum::G2(sk) => sk.sign(scheme, msg).map(SignatureEnum::G2),
        }
    }

    /// Sign a message using the basic signature scheme.
    pub fn sign_basic(&self, msg: impl AsRef<[u8]>) -> BlsResult<SignatureEnum> {
        self.sign(SignatureSchemes::Basic, msg.as_ref())
    }

    /// Sign a message using the message-augmentation signature scheme.
    pub fn sign_augmented(&self, msg: impl AsRef<[u8]>) -> BlsResult<SignatureEnum> {
        self.sign(SignatureSchemes::MessageAugmentation, msg.as_ref())
    }

    /// Sign a message using the proof-of-possession signature scheme.
    pub fn sign_pop(&self, msg: impl AsRef<[u8]>) -> BlsResult<SignatureEnum> {
        self.sign(SignatureSchemes::ProofOfPossession, msg.as_ref())
    }
}

/// The secret key is a field element `x` where 0 < `x` < `r`
/// and `r` is the curve order. See Section 4.3 of
/// <https://eprint.iacr.org/2016/663.pdf>
#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
pub struct SecretKey<C: BlsSignatureImpl>(
    /// The secret key raw value
    #[serde(serialize_with = "traits::scalar::serialize::<C, _>")]
    #[serde(deserialize_with = "traits::scalar::deserialize::<C, _>")]
    pub <<C as Pairing>::PublicKey as Group>::Scalar,
);

impl<C: BlsSignatureImpl> From<SecretKey<C>> for [u8; SECRET_KEY_BYTES] {
    fn from(sk: SecretKey<C>) -> [u8; SECRET_KEY_BYTES] {
        sk.to_be_bytes()
    }
}

impl<'a, C: BlsSignatureImpl> From<&'a SecretKey<C>> for [u8; SECRET_KEY_BYTES] {
    fn from(sk: &'a SecretKey<C>) -> [u8; SECRET_KEY_BYTES] {
        sk.to_be_bytes()
    }
}

impl_from_derivatives_generic!(SecretKey);

impl<C: BlsSignatureImpl> From<&SecretKey<C>> for Vec<u8> {
    fn from(value: &SecretKey<C>) -> Self {
        value.to_be_bytes().to_vec()
    }
}

impl<C: BlsSignatureImpl> TryFrom<&[u8]> for SecretKey<C> {
    type Error = BlsError;

    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        let bytes = <[u8; 32]>::try_from(value)
            .map_err(|_| BlsError::InvalidInputs("Invalid secret key bytes".to_string()))?;
        Option::from(Self::from_be_bytes(&bytes))
            .ok_or_else(|| BlsError::InvalidInputs("Invalid secret key bytes".to_string()))
    }
}

impl<C: BlsSignatureImpl> SecretKey<C> {
    /// Create a new random secret key
    pub fn new() -> Self {
        Self::random(get_crypto_rng())
    }

    /// Compute a secret key from a hash
    pub fn from_hash<B: AsRef<[u8]>>(data: B) -> Self {
        Self(<C as HashToScalar>::hash_to_scalar(
            data.as_ref(),
            KEYGEN_SALT,
        ))
    }

    /// Compute a secret key from a CS-PRNG
    pub fn random(mut rng: impl CryptoRng) -> Self {
        Self(<C as HashToScalar>::hash_to_scalar(
            rng.random::<[u8; SECRET_KEY_BYTES]>(),
            KEYGEN_SALT,
        ))
    }

    /// Get the big-endian byte representation of this key
    pub fn to_be_bytes(&self) -> [u8; SECRET_KEY_BYTES] {
        scalar_to_be_bytes::<C, SECRET_KEY_BYTES>(self.0)
    }

    /// Get the little-endian byte representation of this key
    pub fn to_le_bytes(&self) -> [u8; SECRET_KEY_BYTES] {
        scalar_to_le_bytes::<C, SECRET_KEY_BYTES>(self.0)
    }

    /// Convert a big-endian representation of the secret key.
    pub fn from_be_bytes(bytes: &[u8; SECRET_KEY_BYTES]) -> CtOption<Self> {
        scalar_from_be_bytes::<C, SECRET_KEY_BYTES>(bytes).map(Self)
    }

    /// Convert a little-endian representation of the secret key.
    pub fn from_le_bytes(bytes: &[u8; SECRET_KEY_BYTES]) -> CtOption<Self> {
        scalar_from_le_bytes::<C, SECRET_KEY_BYTES>(bytes).map(Self)
    }

    /// Secret-share this key by creating `limit` shares, of which `threshold`
    /// are required to reconstruct the secret.
    pub fn split(&self, threshold: usize, limit: usize) -> BlsResult<Vec<SecretKeyShare<C>>> {
        self.split_with_rng(threshold, limit, get_crypto_rng())
    }

    /// Secret-share this key using the specified RNG by creating `limit` shares,
    /// of which `threshold` are required to reconstruct the secret.
    pub fn split_with_rng(
        &self,
        threshold: usize,
        limit: usize,
        mut rng: impl CryptoRng,
    ) -> BlsResult<Vec<SecretKeyShare<C>>> {
        let secret = IdentifierPrimeField(self.0);
        let shares = shamir::split_secret::<<C as Pairing>::SecretKeyShare>(
            threshold, limit, &secret, &mut rng,
        )?
        .into_iter()
        .map(SecretKeyShare)
        .collect::<Vec<_>>();
        Ok(shares)
    }

    /// Reconstruct a secret from shares created by [`Self::split`].
    pub fn combine(shares: &[SecretKeyShare<C>]) -> BlsResult<Self> {
        if shares.len() < 2 {
            return Err(BlsError::InvalidInputs(
                "at least two secret key shares are required".to_string(),
            ));
        }
        let ss = shares.iter().map(|s| s.0.clone()).collect::<Vec<_>>();
        let secret = ss.combine()?;
        Ok(Self(secret.0))
    }

    /// Compute the public key
    pub fn public_key(&self) -> PublicKey<C> {
        PublicKey(<C as BlsSignatureCore>::public_key(&self.0))
    }

    /// Create a proof of possession
    pub fn proof_of_possession(&self) -> BlsResult<ProofOfPossession<C>> {
        Ok(ProofOfPossession(<C as BlsSignaturePop>::pop_prove(
            &self.0,
        )?))
    }

    /// Sign a message with this secret key using the specified scheme
    pub fn sign(&self, scheme: SignatureSchemes, msg: &[u8]) -> BlsResult<Signature<C>> {
        match scheme {
            SignatureSchemes::Basic => {
                let inner = <C as BlsSignatureBasic>::sign(&self.0, msg)?;
                Ok(Signature::Basic(inner))
            }
            SignatureSchemes::MessageAugmentation => {
                let inner = <C as BlsSignatureMessageAugmentation>::sign(&self.0, msg)?;
                Ok(Signature::MessageAugmentation(inner))
            }
            SignatureSchemes::ProofOfPossession => {
                let inner = <C as BlsSignaturePop>::sign(&self.0, msg)?;
                Ok(Signature::ProofOfPossession(inner))
            }
        }
    }

    /// Sign a message using the basic signature scheme.
    pub fn sign_basic(&self, msg: impl AsRef<[u8]>) -> BlsResult<Signature<C>> {
        self.sign(SignatureSchemes::Basic, msg.as_ref())
    }

    /// Sign a message using the message-augmentation signature scheme.
    pub fn sign_augmented(&self, msg: impl AsRef<[u8]>) -> BlsResult<Signature<C>> {
        self.sign(SignatureSchemes::MessageAugmentation, msg.as_ref())
    }

    /// Sign a message using the proof-of-possession signature scheme.
    pub fn sign_pop(&self, msg: impl AsRef<[u8]>) -> BlsResult<Signature<C>> {
        self.sign(SignatureSchemes::ProofOfPossession, msg.as_ref())
    }

    /// Create a signcryption decryption key that hides the secret key and can
    /// decrypt the ciphertext.
    pub fn sign_decryption_key<B: AsRef<[u8]>>(
        &self,
        ciphertext: &SignCryptCiphertext<C>,
    ) -> SignCryptDecryptionKey<C> {
        SignCryptDecryptionKey(ciphertext.u * self.0)
    }
}