arcis-compiler 0.9.1

use crate::{
    core::{
        circuits::{
            boolean::{
                boolean_array::BooleanArray,
                boolean_value::BooleanValue,
                byte::Byte,
                ed25519::{Ed25519SecretKey, Ed25519SigningKey},
            },
            key_recovery::{utils::reed_solomon::KeyRecoveryReedSolomonFinal, MXE_KEYS_ENC_COUNT},
        },
        global_value::{
            curve_array::CurveArray,
            curve_value::CurveValue,
            field_array::FieldArray,
            value::FieldValue,
        },
    },
    traits::{FromLeBits, GetBit, GreaterEqual, Reveal, Select},
    utils::{
        crypto::{
            key::{
                AES128Key,
                AES192Key,
                AES256Key,
                RescueKey,
                X25519PrivateKey,
                X25519PublicKey,
                AES_128_KEY_COUNT,
                AES_192_KEY_COUNT,
                RESCUE_KEY_COUNT,
            },
            rescue_cipher::RescueCipher,
        },
        field::{BaseField, ScalarField},
        zkp::elgamal::{ElGamalKeypair, ElGamalSecretKey},
    },
};
use ff::PrimeField;

/// Reconstruction of the MXE Rescue key. Each input consists of:
/// - the 5 threshold secret-shares of the node in the super-cluster
/// - the public key of the node in the super-cluster
///
/// TODOs:
/// 1. the nodes must check that the n first public keys are distinct (otherwise some nodes would
///    hold more than one secret-share) and that the remaining N - n public keys correspond to
///    PublicKey::default()
/// 2. the nodes must check that the K - k last elements of alpha_pows and scaled_polynomials are 0
#[allow(dead_code)]
pub fn key_recovery_finalize<const N: usize, const K: usize, const DMINUSONE: usize>(
    backup_mxe_x25519_private_key: X25519PrivateKey<FieldValue<ScalarField>>,
    peer_x25519_pubkeys: X25519PublicKey<CurveArray<N>>,
    nonce: FieldValue<BaseField>,
    rescue_base_field_key_shares_enc: [FieldArray<N, BaseField>; RESCUE_KEY_COUNT],
    n: FieldValue<BaseField>,
    alpha_pows: FieldArray<K, BaseField>,
    scaled_polynomials: FieldArray<K, BaseField>,
) -> (
    RescueKey<FieldValue<BaseField>>,
    [FieldArray<N, BaseField>; RESCUE_KEY_COUNT],
) {
    assert_eq!(K, (N - 1) / 3 + 1);
    assert_eq!(DMINUSONE, N - (N - 1) / 3 - 1);
    let backup_mxe_x25519_private_key_vec = X25519PrivateKey::new(
        FieldArray::<N, ScalarField>::from(backup_mxe_x25519_private_key.inner()),
        backup_mxe_x25519_private_key.is_expected_non_zero,
    );
    let cipher = RescueCipher::new_with_client_from_keys(
        backup_mxe_x25519_private_key_vec,
        peer_x25519_pubkeys,
    );
    // first n entires are true, remaining ones are false
    let n_indicator = BooleanArray::<N>::from(
        TryInto::<[BooleanValue; N]>::try_into(
            (0..N)
                .map(|i| FieldValue::<BaseField>::from(i).lt(n))
                .collect::<Vec<BooleanValue>>(),
        )
        .unwrap_or_else(|v: Vec<BooleanValue>| {
            panic!("Expected a Vec of length {} (found {})", N, v.len())
        }),
    );
    let rescue_base_field_key_shares =
        TryInto::<[FieldArray<N, BaseField>; RESCUE_KEY_COUNT]>::try_into(cipher.decrypt(
            rescue_base_field_key_shares_enc.to_vec(),
            FieldArray::from(nonce),
        ))
        .unwrap_or_else(|v: Vec<FieldArray<N, BaseField>>| {
            panic!(
                "Expected a Vec of length {} (found {})",
                RESCUE_KEY_COUNT,
                v.len()
            )
        })
        .map(|shares| {
            // we need to set the trailing N - n key shares to 0
            n_indicator.select(shares, FieldArray::<N, BaseField>::from(0))
        });

    let k = (n - 1) / 3 + 1;
    let d_minus_one = n - k;

    let syndromes = rescue_base_field_key_shares.map(|shares| {
        FieldArray::from(KeyRecoveryReedSolomonFinal::compute_syndromes::<
            N,
            DMINUSONE,
            BaseField,
            BooleanValue,
            FieldValue<BaseField>,
        >(
            <[FieldValue<BaseField>; N]>::from(shares), d_minus_one
        ))
    });

    let errors = syndromes.map(|synd| {
        FieldArray::from(KeyRecoveryReedSolomonFinal::compute_errors::<N, DMINUSONE>(
            d_minus_one,
            <[FieldValue<BaseField>; DMINUSONE]>::from(synd),
        ))
    });

    let mxe_rescue_base_field_key = RescueKey::new_from_inner(
        rescue_base_field_key_shares
            .into_iter()
            .zip(errors)
            .map(|(r, e)| {
                let s = KeyRecoveryReedSolomonFinal::subtract_errors::<N, FieldValue<BaseField>>(
                    <[FieldValue<BaseField>; N]>::from(r),
                    <[FieldValue<BaseField>; N]>::from(e),
                );
                KeyRecoveryReedSolomonFinal::decode::<N, K, BooleanValue, FieldValue<BaseField>>(
                    s,
                    <[FieldValue<BaseField>; K]>::from(alpha_pows),
                    <[FieldValue<BaseField>; K]>::from(scaled_polynomials),
                )
            })
            .collect::<Vec<FieldValue<BaseField>>>()
            .try_into()
            .unwrap_or_else(|v: Vec<FieldValue<BaseField>>| {
                panic!(
                    "Expected a Vec of length {} (found {})",
                    RESCUE_KEY_COUNT,
                    v.len()
                )
            }),
    );

    (mxe_rescue_base_field_key, errors)
}

/// Decrypt the encrypted MXE keys with the recovered MXE Rescue key.
#[allow(dead_code, clippy::type_complexity)]
pub fn key_decryption_finalize(
    mxe_rescue_base_field_key: RescueKey<FieldValue<BaseField>>,
    nonce: FieldValue<BaseField>,
    mxe_keys_enc: [FieldValue<BaseField>; MXE_KEYS_ENC_COUNT],
) -> (
    X25519PrivateKey<FieldValue<ScalarField>>,
    X25519PublicKey<CurveValue>,
    RescueKey<FieldValue<ScalarField>>,
    AES128Key<BooleanValue>,
    AES192Key<BooleanValue>,
    AES256Key<BooleanValue>,
    Ed25519SecretKey,
    Ed25519SigningKey,
    ElGamalKeypair,
) {
    let mxe_cipher = RescueCipher::new(mxe_rescue_base_field_key);
    let mxe_keys = mxe_cipher.decrypt(mxe_keys_enc.to_vec(), nonce);

    // convert the x25519 private key key from BaseField to ScalarField
    let mxe_x25519_private_key_base_field = mxe_keys[0];
    let mxe_x25519_private_key = X25519PrivateKey::new(
        FieldValue::<ScalarField>::from_le_bits(
            (0..ScalarField::NUM_BITS as usize)
                .map(|i| mxe_x25519_private_key_base_field.get_bit(i, false))
                .collect::<Vec<BooleanValue>>(),
            false,
        ),
        true,
    );
    let mxe_x25519_pubkey = X25519PublicKey::new_from_private_key(mxe_x25519_private_key).reveal();

    // convert the Rescue scalar field key from BaseField to ScalarField
    let mxe_rescue_scalar_field_key = RescueKey::<FieldValue<ScalarField>>::new_from_inner(
        [
            mxe_keys[1],
            mxe_keys[2],
            mxe_keys[3],
            mxe_keys[4],
            mxe_keys[5],
        ]
        .map(|key| {
            FieldValue::<ScalarField>::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 = mxe_keys[6];
    let mxe_aes_128_key = AES128Key::new_from_inner(
        (0..AES_128_KEY_COUNT)
            .map(|i| mxe_aes_128_key_compressed.get_bit(i, false))
            .collect::<Vec<BooleanValue>>()
            .chunks(8)
            .map(|chunk| {
                Byte::new(
                    chunk
                        .to_vec()
                        .try_into()
                        .unwrap_or_else(|v: Vec<BooleanValue>| {
                            panic!("Expected a Vec of length 8 (found {})", v.len())
                        }),
                )
            })
            .collect::<Vec<Byte<BooleanValue>>>()
            .try_into()
            .unwrap_or_else(|v: Vec<Byte<BooleanValue>>| {
                panic!(
                    "Expected a Vec of length {} (found {})",
                    AES_128_KEY_COUNT / 8,
                    v.len()
                )
            }),
    );

    let mxe_aes_192_key_compressed = mxe_keys[7];
    let mxe_aes_192_key = AES192Key::new_from_inner(
        (0..AES_192_KEY_COUNT)
            .map(|i| mxe_aes_192_key_compressed.get_bit(i, false))
            .collect::<Vec<BooleanValue>>()
            .chunks(8)
            .map(|chunk| {
                Byte::new(
                    chunk
                        .to_vec()
                        .try_into()
                        .unwrap_or_else(|v: Vec<BooleanValue>| {
                            panic!("Expected a Vec of length 8 (found {})", v.len())
                        }),
                )
            })
            .collect::<Vec<Byte<BooleanValue>>>()
            .try_into()
            .unwrap_or_else(|v: Vec<Byte<BooleanValue>>| {
                panic!(
                    "Expected a Vec of length {} (found {})",
                    AES_192_KEY_COUNT / 8,
                    v.len()
                )
            }),
    );

    let mxe_aes_256_key_lo_compressed = mxe_keys[8];
    let mxe_aes_256_key_hi_compressed = mxe_keys[9];
    let mut mxe_aes_256_key_bytes = [Byte::<BooleanValue>::from(0u8); 32];
    mxe_aes_256_key_bytes[..16].copy_from_slice(
        &(0..128)
            .map(|i| mxe_aes_256_key_lo_compressed.get_bit(i, false))
            .collect::<Vec<BooleanValue>>()
            .chunks(8)
            .map(|chunk| {
                Byte::new(
                    chunk
                        .to_vec()
                        .try_into()
                        .unwrap_or_else(|v: Vec<BooleanValue>| {
                            panic!("Expected a Vec of length 8 (found {})", v.len())
                        }),
                )
            })
            .collect::<Vec<Byte<BooleanValue>>>(),
    );
    mxe_aes_256_key_bytes[16..].copy_from_slice(
        &(0..128)
            .map(|i| mxe_aes_256_key_hi_compressed.get_bit(i, false))
            .collect::<Vec<BooleanValue>>()
            .chunks(8)
            .map(|chunk| {
                Byte::new(
                    chunk
                        .to_vec()
                        .try_into()
                        .unwrap_or_else(|v: Vec<BooleanValue>| {
                            panic!("Expected a Vec of length 8 (found {})", v.len())
                        }),
                )
            })
            .collect::<Vec<Byte<BooleanValue>>>(),
    );
    let mxe_aes_256_key = AES256Key::<BooleanValue>::new_from_inner(mxe_aes_256_key_bytes);

    let mxe_ed25519_secret_key_lo_compressed = mxe_keys[10];
    let mxe_ed25519_secret_key_hi_compressed = mxe_keys[11];
    let mut mxe_ed25519_secret_key_bytes = [Byte::<BooleanValue>::from(0u8); 32];
    mxe_ed25519_secret_key_bytes[..16].copy_from_slice(
        &(0..128)
            .map(|i| mxe_ed25519_secret_key_lo_compressed.get_bit(i, false))
            .collect::<Vec<BooleanValue>>()
            .chunks(8)
            .map(|chunk| {
                Byte::new(
                    chunk
                        .to_vec()
                        .try_into()
                        .unwrap_or_else(|v: Vec<BooleanValue>| {
                            panic!("Expected a Vec of length 8 (found {})", v.len())
                        }),
                )
            })
            .collect::<Vec<Byte<BooleanValue>>>(),
    );
    mxe_ed25519_secret_key_bytes[16..].copy_from_slice(
        &(0..128)
            .map(|i| mxe_ed25519_secret_key_hi_compressed.get_bit(i, false))
            .collect::<Vec<BooleanValue>>()
            .chunks(8)
            .map(|chunk| {
                Byte::new(
                    chunk
                        .to_vec()
                        .try_into()
                        .unwrap_or_else(|v: Vec<BooleanValue>| {
                            panic!("Expected a Vec of length 8 (found {})", v.len())
                        }),
                )
            })
            .collect::<Vec<Byte<BooleanValue>>>(),
    );
    let mxe_ed25519_secret_key = Ed25519SecretKey::new_from_inner(mxe_ed25519_secret_key_bytes);
    let mxe_ed25519_signing_key = Ed25519SigningKey::from(mxe_ed25519_secret_key);
    let mxe_ed25519_signing_key = Ed25519SigningKey::new(
        mxe_ed25519_signing_key.s,
        mxe_ed25519_signing_key.hash_prefix,
        mxe_ed25519_signing_key.verifying_key.reveal(),
    );

    // convert the ElGamal secret key from BaseField to ScalarField
    let mxe_elgamal_secret_key_base_field = mxe_keys[12];
    let mxe_elgamal_secret_key = ElGamalSecretKey::new(FieldValue::<ScalarField>::from_le_bits(
        (0..ScalarField::NUM_BITS as usize)
            .map(|i| mxe_elgamal_secret_key_base_field.get_bit(i, false))
            .collect::<Vec<BooleanValue>>(),
        false,
    ));
    let mxe_elgamal_keypair = ElGamalKeypair::new(mxe_elgamal_secret_key);

    (
        mxe_x25519_private_key,
        mxe_x25519_pubkey,
        mxe_rescue_scalar_field_key,
        mxe_aes_128_key,
        mxe_aes_192_key,
        mxe_aes_256_key,
        mxe_ed25519_secret_key,
        mxe_ed25519_signing_key,
        mxe_elgamal_keypair,
    )
}