spark-cryptography 0.1.11

use bitcoin::secp256k1::{PublicKey, SecretKey};

/// Encrypt a message using ECIES with secp256k1
///
/// # Arguments
///
/// * `public_key` - The public key (of the receiver) to encrypt the message to
/// * `msg` - The message to encrypt
///
/// # Returns
///
/// A Result containing the encrypted data or an error
pub fn encrypt_ecies(public_key: &PublicKey, message: &[u8]) -> Result<Vec<u8>, &'static str> {
    // Initialize secp context
    // let secp = Secp256k1::new();

    // // Generate ephemeral key pair
    // let mut rng = OsRng;
    // let ephemeral_secret_key = SecretKey::new(&mut rng);
    // let ephemeral_public_key = PublicKey::from_secret_key(&secp, &ephemeral_secret_key);

    // // Perform ECDH to derive shared secret
    // let shared_point = match public_key.mul_tweak(&secp, &ephemeral_secret_key.into()) {
    //     Ok(point) => point,
    //     Err(_) => return Err("ECDH operation failed"),
    // };

    // // Use the shared secret for encryption
    // let shared_secret = shared_point.serialize();
    // let key_material = sha256::Hash::hash(&shared_secret).to_byte_array();

    // // Encrypt using XOR with key derivation
    // let mut ciphertext = message.to_vec();
    // for i in 0..ciphertext.len() {
    //     let key_byte = key_material[i % key_material.len()];
    //     ciphertext[i] ^= key_byte;
    // }

    // // Format output: ephemeral_pubkey (33 bytes) + ciphertext
    // let mut output = Vec::with_capacity(33 + ciphertext.len());
    // output.extend_from_slice(&ephemeral_public_key.serialize());
    // output.extend_from_slice(&ciphertext);

    let pk_bytes = public_key.serialize();
    ecies::encrypt(&pk_bytes, message).map_err(|_| "ECIES encryption failed")
}

/// Decrypt a message using ECIES with secp256k1
///
/// # Arguments
///
/// * `secret_key` - The secret key (of the receiver) to decrypt the message
/// * `ciphertext` - The encrypted message
///
/// # Returns
///
/// A Result containing the decrypted data or an error
pub fn decrypt_ecies(secret_key: SecretKey, ciphertext: &[u8]) -> Result<Vec<u8>, &'static str> {
    // Check if ciphertext is long enough to contain required components
    // if ciphertext.len() <= 33 {
    //     return Err("Ciphertext too short");
    // }

    // // Initialize secp context
    // let secp = Secp256k1::new();

    // // Extract the ephemeral public key and encrypted data
    // let ephemeral_public_key = match PublicKey::from_slice(&ciphertext[0..33]) {
    //     Ok(key) => key,
    //     Err(_) => return Err("Invalid ephemeral public key in ciphertext"),
    // };
    // let encrypted_data = &ciphertext[33..];

    // // Perform ECDH to derive shared secret
    // let shared_point = match ephemeral_public_key.mul_tweak(&secp, &secret_key.into()) {
    //     Ok(point) => point,
    //     Err(_) => return Err("ECDH operation failed"),
    // };

    // // Use the shared secret for decryption
    // let shared_secret = shared_point.serialize();
    // let key_material = sha256::Hash::hash(&shared_secret).to_byte_array();

    // // Decrypt using XOR (same operation as encryption)
    // let mut plaintext = encrypted_data.to_vec();
    // for i in 0..plaintext.len() {
    //     let key_byte = key_material[i % key_material.len()];
    //     plaintext[i] ^= key_byte;
    // }

    let sk_bytes = secret_key.secret_bytes();
    ::ecies::decrypt(&sk_bytes, ciphertext).map_err(|_| "ECIES decryption failed")
}

#[cfg(test)]
mod tests {
    use super::*;
    use bitcoin::secp256k1::rand::rngs::OsRng;
    use bitcoin::secp256k1::{PublicKey, Secp256k1, SecretKey};
    use rand::RngCore;

    #[test]
    fn test_encrypt_decrypt() {
        let secp = Secp256k1::new();

        // Generate key pair for receiver
        let mut rng = OsRng;
        let receiver_secret_key = SecretKey::new(&mut rng);
        let receiver_public_key = PublicKey::from_secret_key(&secp, &receiver_secret_key);

        // Test message
        let plaintext = b"This is a test message for ECIES encryption";

        // Encrypt
        let ciphertext = encrypt_ecies(&receiver_public_key, plaintext).unwrap();

        // Decrypt
        let decrypted = decrypt_ecies(receiver_secret_key, &ciphertext).unwrap();

        assert_eq!(plaintext.to_vec(), decrypted);
    }

    #[test]
    fn test_decrypt_secret_key() {
        let secp = Secp256k1::new();

        // Generate key pair for identity
        let mut rng = OsRng;
        let identity_secret_key = SecretKey::new(&mut rng);
        let identity_public_key = PublicKey::from_secret_key(&secp, &identity_secret_key);

        // Generate a new secret key to encrypt
        let new_secret_key = SecretKey::new(&mut rng);

        // Encrypt the new secret key
        let ciphertext =
            encrypt_ecies(&identity_public_key, &new_secret_key.secret_bytes()).unwrap();

        // Decrypt the secret key
        let decrypted_bytes = decrypt_ecies(identity_secret_key, &ciphertext).unwrap();
        let decrypted_secret_key = SecretKey::from_slice(&decrypted_bytes).unwrap();

        assert_eq!(new_secret_key, decrypted_secret_key);
    }

    #[test]
    fn test_empty_message() {
        let secp = Secp256k1::new();
        let mut rng = OsRng;
        let secret_key = SecretKey::new(&mut rng);
        let public_key = PublicKey::from_secret_key(&secp, &secret_key);

        // Encrypt and decrypt an empty message
        let empty_message = b"";
        let ciphertext = encrypt_ecies(&public_key, empty_message).unwrap();
        let decrypted = decrypt_ecies(secret_key, &ciphertext).unwrap();

        assert_eq!(empty_message.to_vec(), decrypted);
    }

    #[test]
    fn test_large_message() {
        let secp = Secp256k1::new();
        let mut rng = OsRng;
        let secret_key = SecretKey::new(&mut rng);
        let public_key = PublicKey::from_secret_key(&secp, &secret_key);

        // Create a large message (100KB)
        let large_message = vec![0xAA; 100 * 1024];

        let ciphertext = encrypt_ecies(&public_key, &large_message).unwrap();
        let decrypted = decrypt_ecies(secret_key, &ciphertext).unwrap();

        assert_eq!(large_message, decrypted);
    }

    #[test]
    fn test_binary_data() {
        let secp = Secp256k1::new();
        let mut rng = OsRng;
        let secret_key = SecretKey::new(&mut rng);
        let public_key = PublicKey::from_secret_key(&secp, &secret_key);

        // Test with binary data containing all possible byte values
        let binary_data: Vec<u8> = (0..=255).collect();

        let ciphertext = encrypt_ecies(&public_key, &binary_data).unwrap();
        let decrypted = decrypt_ecies(secret_key, &ciphertext).unwrap();

        assert_eq!(binary_data, decrypted);
    }

    #[test]
    fn test_different_key_fails() {
        let secp = Secp256k1::new();
        let mut rng = OsRng;

        // Generate two different key pairs
        let secret_key1 = SecretKey::new(&mut rng);
        let public_key1 = PublicKey::from_secret_key(&secp, &secret_key1);

        let secret_key2 = SecretKey::new(&mut rng);

        // Encrypt with public_key1
        let message = b"This message should not be decryptable with the wrong key";
        let ciphertext = encrypt_ecies(&public_key1, message).unwrap();

        // Try to decrypt with secret_key2 (should fail or produce incorrect output)
        let decrypted = decrypt_ecies(secret_key2, &ciphertext);

        // The decryption will "succeed" with XOR but produce garbage data
        assert!(decrypted.is_err());
    }

    #[test]
    fn test_invalid_ciphertext() {
        let mut rng = OsRng;
        let secret_key = SecretKey::new(&mut rng);

        // Test with ciphertext that's too short
        let result = decrypt_ecies(secret_key, &[1, 2, 3]);
        assert!(result.is_err());

        // Test with ciphertext that has an invalid public key
        let mut invalid_ciphertext = vec![0; 33 + 10]; // 33 bytes for pubkey + 10 for message
        rng.fill_bytes(&mut invalid_ciphertext); // Fill with random data
        let result = decrypt_ecies(secret_key, &invalid_ciphertext);
        assert!(result.is_err());
    }

    #[test]
    fn test_tampered_ciphertext() {
        let secp = Secp256k1::new();
        let mut rng = OsRng;
        let secret_key = SecretKey::new(&mut rng);
        let public_key = PublicKey::from_secret_key(&secp, &secret_key);

        let message = b"This is a secret message";
        let mut ciphertext = encrypt_ecies(&public_key, message).unwrap();

        // Tamper with the encrypted data portion (not the ephemeral public key)
        if ciphertext.len() > 40 {
            ciphertext[40] ^= 0xFF;
        }

        // The decryption will "succeed" with XOR but produce altered data
        let decrypted = decrypt_ecies(secret_key, &ciphertext);
        assert!(decrypted.is_err());
    }
}