arcis-compiler 0.9.4

use crate::{
    core::{
        circuits::{
            boolean::{boolean_value::BooleanValue, ed25519::Ed25519SecretKey, sha3::SHA3_256},
            key_recovery::{utils::reed_solomon::KeyRecoveryReedSolomonInit, MXE_KEYS_ENC_COUNT},
        },
        global_value::{curve_array::CurveArray, field_array::FieldArray, value::FieldValue},
    },
    traits::{FromLeBits, GetBit, Reveal, ToLeBytes},
    utils::{
        crypto::{
            key::{
                AES128Key,
                AES192Key,
                AES256Key,
                RescueKey,
                X25519PrivateKey,
                X25519PublicKey,
                RESCUE_KEY_COUNT,
            },
            rescue_cipher::RescueCipher,
        },
        field::{BaseField, ScalarField},
        zkp::elgamal::ElGamalSecretKey,
    },
};
use ff::PrimeField;

/// Threshold secret-shares the MXE Rescue key and encrypts under the provided public keys.
/// TODOs:
/// 1. the nodes must check n <= N
/// 2. the nodes must compute g = KeyRecoveryReedSolomonInit::compute_generator_polynomial::<D,
///    BaseField>(d), with D = N - (N - 1) / 3 and d = n - (n - 1)/ 3
/// 3. the nodes must check that the n first public keys are valid and distinct (otherwise some
///    nodes would hold more than one secret-share) and that the remaining N - n public keys
///    correspond to PublicKey::default()
#[allow(dead_code)]
pub fn key_recovery_init<const N: usize, const D: usize>(
    base_mxe_x25519_private_key: X25519PrivateKey<FieldValue<ScalarField>>,
    mxe_rescue_base_field_key: RescueKey<FieldValue<BaseField>>,
    peer_x25519_pubkeys: X25519PublicKey<CurveArray<N>>,
    nonce: FieldValue<BaseField>,
    n: FieldValue<BaseField>,
    g: [FieldValue<BaseField>; D],
) -> [FieldArray<N, BaseField>; RESCUE_KEY_COUNT] {
    assert_eq!(D, N - (N - 1) / 3);
    // threshold secret-share the Rescue key
    let rescue_base_field_key_shares = mxe_rescue_base_field_key.inner().map(|key| {
        FieldArray::from(KeyRecoveryReedSolomonInit::encode::<
            N,
            D,
            BooleanValue,
            FieldValue<BaseField>,
        >(key, g, (n - 1) / 3 + 1))
    });

    // for each provided public key, encrypt the 5 key shares
    let base_mxe_x25519_private_key_vec = X25519PrivateKey::new(
        FieldArray::<N, ScalarField>::from(base_mxe_x25519_private_key.inner()),
        base_mxe_x25519_private_key.is_expected_non_zero,
    );
    let peer_ciphers = RescueCipher::new_with_client_from_keys(
        base_mxe_x25519_private_key_vec,
        peer_x25519_pubkeys,
    );
    let rescue_base_field_key_shares_enc = peer_ciphers
        .encrypt(
            rescue_base_field_key_shares.to_vec(),
            FieldArray::from(nonce),
        )
        .try_into()
        .unwrap_or_else(|v: Vec<FieldArray<N, BaseField>>| {
            panic!(
                "Expected a Vec of length {} (found {})",
                RESCUE_KEY_COUNT,
                v.len()
            )
        });

    rescue_base_field_key_shares_enc
}

/// Encrypt the MXE keys with the MXE Rescue cipher.
#[allow(dead_code, clippy::too_many_arguments)]
pub fn key_encryption_init(
    mxe_rescue_base_field_key: RescueKey<FieldValue<BaseField>>,
    nonce: FieldValue<BaseField>,
    mxe_x25519_private_key: X25519PrivateKey<FieldValue<ScalarField>>,
    mxe_rescue_scalar_field_key: RescueKey<FieldValue<ScalarField>>,
    mxe_aes_128_key: AES128Key<BooleanValue>,
    mxe_aes_192_key: AES192Key<BooleanValue>,
    mxe_aes_256_key: AES256Key<BooleanValue>,
    mxe_ed25519_secret_key: Ed25519SecretKey,
    mxe_elgamal_secret_key: ElGamalSecretKey,
) -> [FieldValue<BaseField>; MXE_KEYS_ENC_COUNT] {
    let mxe_cipher = RescueCipher::new(mxe_rescue_base_field_key);

    // MXE-encrypt the x25519 private key, Rescue scalar field key, AES keys, ed25519 secret key and
    // the ElGamal secret key
    // convert the x25519 private key from ScalarField to BaseField
    let mxe_x25519_private_key = FieldValue::<BaseField>::from_le_bits(
        (0..ScalarField::NUM_BITS as usize)
            .map(|i| mxe_x25519_private_key.inner().get_bit(i, false))
            .collect::<Vec<BooleanValue>>(),
        false,
    );

    // convert the Rescue scalar field key from ScalarField to BaseField
    let mxe_rescue_scalar_field_key = mxe_rescue_scalar_field_key.inner().map(|key| {
        FieldValue::<BaseField>::from_le_bits(
            (0..ScalarField::NUM_BITS as usize)
                .map(|i| key.get_bit(i, false))
                .collect::<Vec<BooleanValue>>(),
            false,
        )
    });

    let mxe_aes_128_key_compressed = FieldValue::<BaseField>::from_le_bits(
        mxe_aes_128_key
            .inner()
            .into_iter()
            .flat_map(|byte| byte.to_vec())
            .collect::<Vec<BooleanValue>>(),
        false,
    );
    let mxe_aes_192_key_compressed = FieldValue::<BaseField>::from_le_bits(
        mxe_aes_192_key
            .inner()
            .into_iter()
            .flat_map(|byte| byte.to_vec())
            .collect::<Vec<BooleanValue>>(),
        false,
    );
    let mxe_aes_256_key_bits = mxe_aes_256_key
        .inner()
        .into_iter()
        .flat_map(|byte| byte.to_vec())
        .collect::<Vec<BooleanValue>>();
    let mxe_aes_256_key_lo_compressed = FieldValue::<BaseField>::from_le_bits(
        mxe_aes_256_key_bits
            .iter()
            .copied()
            .take(128)
            .collect::<Vec<BooleanValue>>(),
        false,
    );
    let mxe_aes_256_key_hi_compressed = FieldValue::<BaseField>::from_le_bits(
        mxe_aes_256_key_bits
            .iter()
            .copied()
            .skip(128)
            .collect::<Vec<BooleanValue>>(),
        false,
    );

    let mxe_ed25519_secret_key_bits = mxe_ed25519_secret_key
        .inner()
        .into_iter()
        .flat_map(|byte| byte.to_vec())
        .collect::<Vec<BooleanValue>>();
    let mxe_ed25519_secret_key_lo_compressed = FieldValue::<BaseField>::from_le_bits(
        mxe_ed25519_secret_key_bits
            .iter()
            .copied()
            .take(128)
            .collect::<Vec<BooleanValue>>(),
        false,
    );
    let mxe_ed25519_secret_key_hi_compressed = FieldValue::<BaseField>::from_le_bits(
        mxe_ed25519_secret_key_bits
            .iter()
            .copied()
            .skip(128)
            .collect::<Vec<BooleanValue>>(),
        false,
    );

    // convert the ElGamal key from ScalarField to BaseField
    let mxe_elgamal_secret_key = FieldValue::<BaseField>::from_le_bits(
        (0..ScalarField::NUM_BITS as usize)
            .map(|i| mxe_elgamal_secret_key.get_scalar().get_bit(i, false))
            .collect::<Vec<BooleanValue>>(),
        false,
    );

    let mut mxe_keys = vec![mxe_x25519_private_key];
    mxe_keys.extend(mxe_rescue_scalar_field_key);
    mxe_keys.extend(vec![
        mxe_aes_128_key_compressed,
        mxe_aes_192_key_compressed,
        mxe_aes_256_key_lo_compressed,
        mxe_aes_256_key_hi_compressed,
        mxe_ed25519_secret_key_lo_compressed,
        mxe_ed25519_secret_key_hi_compressed,
        mxe_elgamal_secret_key,
    ]);
    let mxe_keys_enc = mxe_cipher
        .encrypt(mxe_keys, nonce)
        .try_into()
        .unwrap_or_else(|v: Vec<FieldValue<BaseField>>| {
            panic!(
                "Expected a Vec of length {} (found {})",
                MXE_KEYS_ENC_COUNT,
                v.len()
            )
        });

    mxe_keys_enc
}

/// Compute a SHA3-256 digest of the concatenation of all MXE keys.
#[allow(dead_code, clippy::too_many_arguments)]
pub fn hash_keys(
    mxe_x25519_private_key: X25519PrivateKey<FieldValue<ScalarField>>,
    mxe_rescue_base_field_key: RescueKey<FieldValue<BaseField>>,
    mxe_rescue_scalar_field_key: RescueKey<FieldValue<ScalarField>>,
    mxe_aes_128_key: AES128Key<BooleanValue>,
    mxe_aes_192_key: AES192Key<BooleanValue>,
    mxe_aes_256_key: AES256Key<BooleanValue>,
    mxe_ed25519_secret_key: Ed25519SecretKey,
    mxe_elgamal_secret_key: ElGamalSecretKey,
) -> [FieldValue<BaseField>; 32] {
    let mut mxe_keys = mxe_x25519_private_key.inner().to_le_bytes().to_vec();
    mxe_keys.extend(
        mxe_rescue_base_field_key
            .inner()
            .into_iter()
            .flat_map(|key| key.to_le_bytes().to_vec()),
    );
    mxe_keys.extend(
        mxe_rescue_scalar_field_key
            .inner()
            .into_iter()
            .flat_map(|key| key.to_le_bytes().to_vec()),
    );
    mxe_keys.extend(mxe_aes_128_key.inner());
    mxe_keys.extend(mxe_aes_192_key.inner());
    mxe_keys.extend(mxe_aes_256_key.inner());
    mxe_keys.extend(mxe_ed25519_secret_key.inner());
    mxe_keys.extend(mxe_elgamal_secret_key.get_scalar().to_le_bytes());

    let hasher = SHA3_256::new();
    let digest = hasher.digest(mxe_keys);

    digest.map(|byte| FieldValue::<BaseField>::from(byte.reveal()))
}