newton-core 0.4.16

//! BLS key generation and loading utilities

use std::path::PathBuf;

use crate::keys::{
    error::{KeyError, KeyResult},
    get_bls_key_path,
    storage::{read_key_file, write_key_file},
};
use ark_bn254::{Fr, G1Projective};
use ark_ec::PrimeGroup;
use ark_ff::{BigInteger, PrimeField, UniformRand};
use newton_crypto_bls::BlsKeyPair;
use rand_core::OsRng;
use rust_bls_bn254::keystores::base_keystore::Keystore;
use tracing::{debug, info};

use alloy::primitives::{B256, U256};

/// Load a BLS private key from storage and create a keypair
pub fn load_bls_key_from_path(path: &str) -> KeyResult<BlsKeyPair> {
    debug!("Loading BLS key: {}", path);

    let hex_key = read_key_file(&PathBuf::from(path))?;

    // Remove 0x prefix if present
    let hex_key = hex_key.strip_prefix("0x").unwrap_or(&hex_key);

    // Decode hex to bytes
    let key_bytes = hex::decode(hex_key).map_err(KeyError::HexDecode)?;

    // Convert to string format expected by BlsKeyPair::new (same as existing keystore format)
    let fr_key: String = key_bytes.iter().map(|&b| b as char).collect();

    // Create BlsKeyPair using the same method as the existing code
    let key_pair = BlsKeyPair::new(fr_key)
        .map_err(|e| KeyError::BlsKeypair(eyre::eyre!("Failed to create BLS keypair: {}", e)))?;

    debug!("Loaded BLS key: {}", path);
    Ok(key_pair)
}

/// Get the public key for an existing BLS key without loading the full keypair
pub fn get_bls_public_from_path(path: &str) -> KeyResult<G1Projective> {
    let hex_key = read_key_file(&PathBuf::from(path))?;

    // Remove 0x prefix if present
    let hex_key = hex_key.strip_prefix("0x").unwrap_or(&hex_key);

    // Decode hex to bytes
    let key_bytes = hex::decode(hex_key).map_err(KeyError::HexDecode)?;

    // Convert bytes back to Fr
    let sk = Fr::from_le_bytes_mod_order(&key_bytes);

    // Generate public key
    let pk = G1Projective::generator().mul_bigint((sk).into_bigint());

    Ok(pk)
}

/// Load BLS key from keystore file
pub fn load_bls_key_from_keystore(keystore_path: &str, keystore_password: &str) -> KeyResult<BlsKeyPair> {
    let keystore = Keystore::from_file(keystore_path)
        .map_err(|e| KeyError::InvalidKey(format!("Failed to load BLS keystore: {}", e)))?
        .decrypt(keystore_password)
        .map_err(|e| KeyError::InvalidKey(format!("Failed to decrypt BLS keystore: {}", e)))?;

    let key_bytes: Vec<u8> = keystore;
    let big_int = num_bigint::BigUint::from_bytes_be(&key_bytes);
    let decimal_key = big_int.to_string();

    BlsKeyPair::new(decimal_key).map_err(|e| KeyError::BlsKeypair(eyre::eyre!("Failed to create BLS keypair: {}", e)))
}

/// Load BLS key from decimal string
pub fn load_operator_bls_key(bls_private_key: &str) -> KeyResult<BlsKeyPair> {
    BlsKeyPair::new(bls_private_key.to_string())
        .map_err(|e| KeyError::BlsKeypair(eyre::eyre!("Failed to create BLS keypair: {}", e)))
}

#[cfg(test)]
mod tests {
    use crate::keys::ecdsa::load_operator_ecdsa_key;

    use super::*;
    use crate::{
        newton_prover_task_manager::{INewtonPolicy, INewtonProverTaskManager::TaskResponse, NewtonMessage},
        r#newton_prover_task_manager::NewtonProverTaskManager,
    };
    use alloy::{
        primitives::{keccak256, Address, Bytes, B256, U256},
        sol_types::SolValue,
    };
    use newton_crypto_bn254::utils::verify_message;

    #[test]
    fn test_load_operator_keys_and_verify_signature() {
        // Test BLS private key (decimal string)
        let bls_private_key = "5665911737697037570291183107345399045458787373029916535446553888385923852145";
        // Test ECDSA private key (hex string)
        let ecdsa_private_key = "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80";

        // Load BLS key
        let bls_keypair = load_operator_bls_key(bls_private_key).expect("Failed to load BLS key");
        info!("Successfully loaded BLS key");

        // Load ECDSA key
        let ecdsa_signer = load_operator_ecdsa_key(ecdsa_private_key).expect("Failed to load ECDSA key");
        info!("Successfully loaded ECDSA key");
        info!("ECDSA Address: {}", ecdsa_signer.address());

        // Create a TaskResponse using the actual binding
        let task_response = TaskResponse {
            taskId: B256::from([1u8; 32]),
            policyClient: Address::from([2u8; 20]),
            policyId: B256::from([3u8; 32]),
            policyAddress: Address::from([4u8; 20]),
            intent: NewtonMessage::Intent {
                from: Address::from([5u8; 20]),
                to: Address::from([6u8; 20]),
                value: U256::from(1000),
                data: [0x01, 0x02, 0x03].into(),
                chainId: U256::from(1),
                functionSignature: [0x04, 0x05].into(),
            },
            intentSignature: Bytes::default(),
            evaluationResult: [0x01].into(), // true
            policyTaskData: NewtonMessage::PolicyTaskData {
                policyId: B256::from([3u8; 32]),
                policyAddress: Address::from([4u8; 20]),
                policy: Bytes::default(),
                policyData: vec![],
            },
            policyConfig: INewtonPolicy::PolicyConfig {
                policyParams: Bytes::default(),
                expireAfter: 0,
            },
            initializationTimestamp: U256::ZERO,
        };

        // Encode the task response
        let encoded_response = task_response.abi_encode();
        info!("Encoded response length: {}", encoded_response.len());

        // Hash the encoded response
        let message_hash = keccak256(&encoded_response);
        info!("Message hash: {:?}", message_hash);

        // Sign the message with BLS key
        let signature = bls_keypair.sign_message(&message_hash.into());
        info!("Successfully signed message with BLS key");

        // Verify the signature using verify_message from eigensdk
        let is_valid = verify_message(
            bls_keypair.public_key_g2().g2(),
            &message_hash,
            signature.g1_point().g1(),
        );

        assert!(is_valid, "BLS signature verification failed");
        info!("  BLS signature verification successful");

        // Additional verification - check public key derivation
        let public_key_g1 = bls_keypair.public_key().g1();
        info!("BLS Public Key G1: {:?}", public_key_g1);

        let public_key_g2 = bls_keypair.public_key_g2().g2();
        info!("BLS Public Key G2: {:?}", public_key_g2);
    }

    #[test]
    fn test_load_bls_key_with_different_formats() {
        // Test with the same BLS private key
        let bls_private_key = "5665911737697037570291183107345399045458787373029916535446553888385923852145";

        let keypair1 = load_operator_bls_key(bls_private_key).expect("Failed to load BLS key");

        // Verify the public key is consistent
        let pk1 = keypair1.public_key().g1();
        info!("Public key loaded successfully: {:?}", pk1);

        // Load again to ensure consistency
        let keypair2 = load_operator_bls_key(bls_private_key).expect("Failed to load BLS key second time");
        let pk2 = keypair2.public_key().g1();

        assert_eq!(pk1, pk2, "Public keys should be identical for the same private key");
    }

    #[test]
    fn test_multiple_message_signatures() {
        let bls_private_key = "5665911737697037570291183107345399045458787373029916535446553888385923852145";
        let bls_keypair = load_operator_bls_key(bls_private_key).expect("Failed to load BLS key");

        // Test signing multiple different messages
        let messages = [B256::from([1u8; 32]), B256::from([2u8; 32]), B256::from([255u8; 32])];

        for (i, message) in messages.iter().enumerate() {
            let signature = bls_keypair.sign_message(&(*message).into());

            let is_valid = verify_message(bls_keypair.public_key_g2().g2(), message, signature.g1_point().g1());

            assert!(is_valid, "Message {} verification failed", i);
            info!("  Message {} signature verified", i);
        }
    }

    /// Verify that BLS private keys derive to the expected G1/G2 pubkeys
    /// This test verifies the Sepolia stagef operator keys against on-chain registered values
    #[test]
    fn test_verify_operator_bls_keys_sepolia_stagef() {
        use ark_ec::AffineRepr;
        use newton_crypto_bls::convert_to_g1_point;

        // Helper to compute operator ID from G1 pubkey (keccak256 of ABI-encoded G1 point)
        fn compute_operator_id(g1: ark_bn254::G1Affine) -> B256 {
            let g1_point = convert_to_g1_point(g1).expect("Failed to convert G1 point");
            // ABI encode: two uint256 values (X, Y)
            let mut encoded = Vec::with_capacity(64);
            encoded.extend_from_slice(&g1_point.X.to_be_bytes::<32>());
            encoded.extend_from_slice(&g1_point.Y.to_be_bytes::<32>());
            keccak256(&encoded)
        }

        // Operator configurations from environment
        // Each tuple: (name, bls_private_key, expected_operator_id, expected_g1_x, expected_g1_y)
        let operators = [
            (
                "Operator 1 (0x0B26205c99A7669cCC6Dfe4D95CaCC5BfE62e2fe)",
                "3437392234444159768111770309164609957102690468096034096776626309734712123977",
                "0xfd8c33b1612b4bccb01926d9ff26aa56ac768bb493346adeec36b2538b51e8c1",
                "3921767720995179277909625244037822258046941948782863550120879716817940452717",
                "9466587799414527615865516490417180554460539059429958111726479271529872103751",
            ),
            (
                "Operator 2 (0x84C2662C0359dA7ac2101c7392ab767A4A80Cf96)",
                "7019086643264615414456756527632272354703097262265640370649595594745347218247",
                "0x03f232d0b0d939f8987208c7196f46bcd79310137277480595e140bab0b8a2f0",
                "5785384538136481768800205166274937495894624497780844126780399836164351973025",
                "21407700507482687489947338356096094467844731053681107079477243210556238999880",
            ),
            (
                "Operator 2 (0x6FAf85Dda4e07300418beDc31AF6621383de65E8)",
                "17338602372173648149501880755323073193946756536467840304826772396306052737333",
                "0xfe0742d124d105dcfb28b845a89035e7c92de3ff36d32606caf2099017a31974",
                "13185909042843528655162451917016584196037912218358625181215090182994157198053",
                "13575221952311471737620708106833236715951854541195392557376226506693517828827",
            ),
        ];

        println!("\n=== BLS Key Derivation Verification ===\n");

        for (name, bls_private_key, expected_op_id, expected_g1_x, expected_g1_y) in operators {
            println!("Verifying: {}", name);
            println!("  BLS Private Key: {}", bls_private_key);

            // Load the BLS keypair
            let keypair = load_operator_bls_key(bls_private_key).expect("Failed to load BLS key");

            // Get derived G1 pubkey
            let g1 = keypair.public_key().g1();
            let (g1_x, g1_y) = if let (Some(x), Some(y)) = (g1.x(), g1.y()) {
                (x.to_string(), y.to_string())
            } else {
                ("infinity".to_string(), "infinity".to_string())
            };

            // Get derived G2 pubkey
            let g2 = keypair.public_key_g2().g2();
            let g2_coords = if let (Some(x), Some(y)) = (g2.x(), g2.y()) {
                format!("X_c0={}, X_c1={}, Y_c0={}, Y_c1={}", x.c0, x.c1, y.c0, y.c1)
            } else {
                "infinity".to_string()
            };

            // Compute operator ID from G1 pubkey
            let derived_op_id = compute_operator_id(keypair.public_key().g1());
            let derived_op_id_hex = format!("0x{}", hex::encode(derived_op_id.as_slice()));

            println!("  Derived G1_X: {}", g1_x);
            println!("  Derived G1_Y: {}", g1_y);
            println!("  Derived G2: {}", g2_coords);
            println!("  Derived Operator ID: {}", derived_op_id_hex);
            println!("  Expected Operator ID: {}", expected_op_id);

            // Verify G1 matches expected
            let g1_x_matches = g1_x == expected_g1_x;
            let g1_y_matches = g1_y == expected_g1_y;
            let op_id_matches = derived_op_id_hex.to_lowercase() == expected_op_id.to_lowercase();

            println!("  G1_X matches: {}", g1_x_matches);
            println!("  G1_Y matches: {}", g1_y_matches);
            println!("  Operator ID matches: {}", op_id_matches);

            if g1_x_matches && g1_y_matches && op_id_matches {
                println!("  ✓ VERIFIED: Private key derives to on-chain G1 pubkey\n");
            } else {
                println!("  ✗ MISMATCH: Private key does NOT derive to on-chain G1 pubkey!");
                println!("    Expected G1_X: {}", expected_g1_x);
                println!("    Expected G1_Y: {}", expected_g1_y);
                println!();
            }
        }
    }

    /// Verify G2 point addition for the two signers
    /// This test computes Operator1_G2 + Operator2_G2 and compares with the logged aggregated G2
    #[test]
    fn test_verify_g2_aggregation() {
        use ark_ec::{AffineRepr, CurveGroup};

        // Operator 1 BLS key
        let op1_bls_key = "3437392234444159768111770309164609957102690468096034096776626309734712123977";
        let op1_keypair = load_operator_bls_key(op1_bls_key).expect("Failed to load Operator 1 BLS key");

        // Operator 2 BLS key
        let op2_bls_key = "7019086643264615414456756527632272354703097262265640370649595594745347218247";
        let op2_keypair = load_operator_bls_key(op2_bls_key).expect("Failed to load Operator 2 BLS key");

        // Get G2 pubkeys
        let op1_g2 = op1_keypair.public_key_g2().g2();
        let op2_g2 = op2_keypair.public_key_g2().g2();

        // Compute aggregated G2 (Operator 1 + Operator 2)
        let aggregated_g2 = (op1_g2 + op2_g2).into_affine();

        // Print results
        println!("\n=== G2 Aggregation Verification ===\n");

        if let (Some(x), Some(y)) = (op1_g2.x(), op1_g2.y()) {
            println!("Operator 1 G2:");
            println!("  X_c0 = {}", x.c0);
            println!("  X_c1 = {}", x.c1);
            println!("  Y_c0 = {}", y.c0);
            println!("  Y_c1 = {}", y.c1);
        }

        if let (Some(x), Some(y)) = (op2_g2.x(), op2_g2.y()) {
            println!("\nOperator 2 G2:");
            println!("  X_c0 = {}", x.c0);
            println!("  X_c1 = {}", x.c1);
            println!("  Y_c0 = {}", y.c0);
            println!("  Y_c1 = {}", y.c1);
        }

        if let (Some(x), Some(y)) = (aggregated_g2.x(), aggregated_g2.y()) {
            println!("\nExpected Aggregated G2 (Op1 + Op2):");
            println!("  X_c0 = {}", x.c0);
            println!("  X_c1 = {}", x.c1);
            println!("  Y_c0 = {}", y.c0);
            println!("  Y_c1 = {}", y.c1);

            // Compare with logged value from Datadog
            let logged_x_c0 = "14353610553595172534056806588837466963076630321468849559816767082177709710634";
            let logged_x_c1 = "8614678016146218371314685919997481237934655769445601302232162555523220039328";
            let logged_y_c0 = "11923179702360463644198799684718882158376252611210558687905593530626203466916";
            let logged_y_c1 = "21202544982869411316901470582335588974306682123429582574592263080885762817812";

            println!("\nLogged Aggregated G2 (from Datadog):");
            println!("  X_c0 = {}", logged_x_c0);
            println!("  X_c1 = {}", logged_x_c1);
            println!("  Y_c0 = {}", logged_y_c0);
            println!("  Y_c1 = {}", logged_y_c1);

            let x_c0_match = x.c0.to_string() == logged_x_c0;
            let x_c1_match = x.c1.to_string() == logged_x_c1;
            let y_c0_match = y.c0.to_string() == logged_y_c0;
            let y_c1_match = y.c1.to_string() == logged_y_c1;

            println!("\nComparison:");
            println!("  X_c0 matches: {}", x_c0_match);
            println!("  X_c1 matches: {}", x_c1_match);
            println!("  Y_c0 matches: {}", y_c0_match);
            println!("  Y_c1 matches: {}", y_c1_match);

            if x_c0_match && x_c1_match && y_c0_match && y_c1_match {
                println!("\n✓ AGGREGATED G2 IS CORRECT");
            } else {
                println!("\n✗ AGGREGATED G2 MISMATCH - BN254 pairing will fail!");
            }
        }
    }
}