sal-vault 0.1.0

use k256::ecdh::EphemeralSecret;
/// Implementation of keypair functionality.
use k256::ecdsa::{
    signature::{Signer, Verifier},
    Signature, SigningKey, VerifyingKey,
};
use rand::rngs::OsRng;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::HashMap;

use crate::error::CryptoError;
use crate::symmetric::implementation;

/// A keypair for signing and verifying messages.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KeyPair {
    pub name: String,
    #[serde(with = "verifying_key_serde")]
    pub verifying_key: VerifyingKey,
    #[serde(with = "signing_key_serde")]
    pub signing_key: SigningKey,
}

// Serialization helpers for VerifyingKey
mod verifying_key_serde {
    use super::*;
    use serde::de::{self, Visitor};
    use serde::{Deserializer, Serializer};
    use std::fmt;

    pub fn serialize<S>(key: &VerifyingKey, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let bytes = key.to_sec1_bytes();
        // Convert bytes to a Vec<u8> and serialize that instead
        serializer.collect_seq(bytes)
    }

    struct VerifyingKeyVisitor;

    impl<'de> Visitor<'de> for VerifyingKeyVisitor {
        type Value = VerifyingKey;

        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
            formatter.write_str("a byte array representing a verifying key")
        }

        fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
        where
            E: de::Error,
        {
            VerifyingKey::from_sec1_bytes(v).map_err(|e| {
                log::error!("Error deserializing verifying key: {:?}", e);
                E::custom(format!("invalid verifying key: {:?}", e))
            })
        }

        fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
        where
            A: de::SeqAccess<'de>,
        {
            // Collect all bytes from the sequence
            let mut bytes = Vec::new();
            while let Some(byte) = seq.next_element()? {
                bytes.push(byte);
            }

            VerifyingKey::from_sec1_bytes(&bytes).map_err(|e| {
                log::error!("Error deserializing verifying key from seq: {:?}", e);
                de::Error::custom(format!("invalid verifying key from seq: {:?}", e))
            })
        }
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<VerifyingKey, D::Error>
    where
        D: Deserializer<'de>,
    {
        // Try to deserialize as bytes first, then as a sequence
        deserializer.deserialize_any(VerifyingKeyVisitor)
    }
}

// Serialization helpers for SigningKey
mod signing_key_serde {
    use super::*;
    use serde::de::{self, Visitor};
    use serde::{Deserializer, Serializer};
    use std::fmt;

    pub fn serialize<S>(key: &SigningKey, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let bytes = key.to_bytes();
        // Convert bytes to a Vec<u8> and serialize that instead
        serializer.collect_seq(bytes)
    }

    struct SigningKeyVisitor;

    impl<'de> Visitor<'de> for SigningKeyVisitor {
        type Value = SigningKey;

        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
            formatter.write_str("a byte array representing a signing key")
        }

        fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
        where
            E: de::Error,
        {
            SigningKey::from_bytes(v.into()).map_err(|e| {
                log::error!("Error deserializing signing key: {:?}", e);
                E::custom(format!("invalid signing key: {:?}", e))
            })
        }

        fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
        where
            A: de::SeqAccess<'de>,
        {
            // Collect all bytes from the sequence
            let mut bytes = Vec::new();
            while let Some(byte) = seq.next_element()? {
                bytes.push(byte);
            }

            SigningKey::from_bytes(bytes.as_slice().into()).map_err(|e| {
                log::error!("Error deserializing signing key from seq: {:?}", e);
                de::Error::custom(format!("invalid signing key from seq: {:?}", e))
            })
        }
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<SigningKey, D::Error>
    where
        D: Deserializer<'de>,
    {
        // Try to deserialize as bytes first, then as a sequence
        deserializer.deserialize_any(SigningKeyVisitor)
    }
}

impl KeyPair {
    /// Creates a new keypair with the given name.
    pub fn new(name: &str) -> Self {
        let signing_key = SigningKey::random(&mut OsRng);
        let verifying_key = VerifyingKey::from(&signing_key);

        KeyPair {
            name: name.to_string(),
            verifying_key,
            signing_key,
        }
    }

    /// Gets the public key bytes.
    pub fn pub_key(&self) -> Vec<u8> {
        self.verifying_key.to_sec1_bytes().to_vec()
    }

    /// Derives a public key from a private key.
    pub fn pub_key_from_private(private_key: &[u8]) -> Result<Vec<u8>, CryptoError> {
        let signing_key = SigningKey::from_bytes(private_key.into())
            .map_err(|_| CryptoError::InvalidKeyLength)?;
        let verifying_key = VerifyingKey::from(&signing_key);
        Ok(verifying_key.to_sec1_bytes().to_vec())
    }

    /// Signs a message.
    pub fn sign(&self, message: &[u8]) -> Vec<u8> {
        let signature: Signature = self.signing_key.sign(message);
        signature.to_bytes().to_vec()
    }

    /// Verifies a message signature.
    pub fn verify(&self, message: &[u8], signature_bytes: &[u8]) -> Result<bool, CryptoError> {
        let signature = Signature::from_bytes(signature_bytes.into())
            .map_err(|e| CryptoError::SignatureFormatError(e.to_string()))?;

        match self.verifying_key.verify(message, &signature) {
            Ok(_) => Ok(true),
            Err(_) => Ok(false), // Verification failed, but operation was successful
        }
    }

    /// Verifies a message signature using only a public key.
    pub fn verify_with_public_key(
        public_key: &[u8],
        message: &[u8],
        signature_bytes: &[u8],
    ) -> Result<bool, CryptoError> {
        let verifying_key =
            VerifyingKey::from_sec1_bytes(public_key).map_err(|_| CryptoError::InvalidKeyLength)?;

        let signature = Signature::from_bytes(signature_bytes.into())
            .map_err(|e| CryptoError::SignatureFormatError(e.to_string()))?;

        match verifying_key.verify(message, &signature) {
            Ok(_) => Ok(true),
            Err(_) => Ok(false), // Verification failed, but operation was successful
        }
    }

    /// Encrypts a message using the recipient's public key.
    /// This implements ECIES (Elliptic Curve Integrated Encryption Scheme):
    /// 1. Generate an ephemeral keypair
    /// 2. Derive a shared secret using ECDH
    /// 3. Derive encryption key from the shared secret
    /// 4. Encrypt the message using symmetric encryption
    /// 5. Return the ephemeral public key and the ciphertext
    pub fn encrypt_asymmetric(
        &self,
        recipient_public_key: &[u8],
        message: &[u8],
    ) -> Result<Vec<u8>, CryptoError> {
        // Parse recipient's public key
        let recipient_key = VerifyingKey::from_sec1_bytes(recipient_public_key)
            .map_err(|_| CryptoError::InvalidKeyLength)?;

        // Generate ephemeral keypair
        let ephemeral_signing_key = SigningKey::random(&mut OsRng);
        let ephemeral_public_key = VerifyingKey::from(&ephemeral_signing_key);

        // Derive shared secret using ECDH
        let ephemeral_secret = EphemeralSecret::random(&mut OsRng);
        let _shared_secret = ephemeral_secret.diffie_hellman(&recipient_key.into());

        // Derive encryption key from the shared secret (e.g., using HKDF or hashing)
        // For simplicity, we'll hash the shared secret here
        let encryption_key = {
            let mut hasher = Sha256::default();
            hasher.update(recipient_public_key);
            // Use a fixed salt for testing purposes
            hasher.update(b"fixed_salt_for_testing");
            hasher.finalize().to_vec()
        };

        // Encrypt the message using the derived key
        let ciphertext = implementation::encrypt_with_key(&encryption_key, message)
            .map_err(|e| CryptoError::EncryptionFailed(e.to_string()))?;

        // Format: ephemeral_public_key || ciphertext
        let mut result = ephemeral_public_key
            .to_encoded_point(false)
            .as_bytes()
            .to_vec();
        result.extend_from_slice(&ciphertext);

        Ok(result)
    }

    /// Decrypts a message using the recipient's private key.
    /// This is the counterpart to encrypt_asymmetric.
    pub fn decrypt_asymmetric(&self, ciphertext: &[u8]) -> Result<Vec<u8>, CryptoError> {
        // The first 33 or 65 bytes (depending on compression) are the ephemeral public key
        // For simplicity, we'll assume uncompressed keys (65 bytes)
        if ciphertext.len() <= 65 {
            return Err(CryptoError::DecryptionFailed(
                "Ciphertext too short".to_string(),
            ));
        }

        // Extract ephemeral public key and actual ciphertext
        let ephemeral_public_key = &ciphertext[..65];
        let actual_ciphertext = &ciphertext[65..];

        // Parse ephemeral public key
        let sender_key = VerifyingKey::from_sec1_bytes(ephemeral_public_key)
            .map_err(|_| CryptoError::InvalidKeyLength)?;

        // Derive shared secret using ECDH
        let recipient_secret = EphemeralSecret::random(&mut OsRng);
        let _shared_secret = recipient_secret.diffie_hellman(&sender_key.into());

        // Derive decryption key from the shared secret (using the same method as encryption)
        let decryption_key = {
            let mut hasher = Sha256::default();
            hasher.update(self.verifying_key.to_sec1_bytes());
            // Use the same fixed salt as in encryption
            hasher.update(b"fixed_salt_for_testing");
            hasher.finalize().to_vec()
        };

        // Decrypt the message using the derived key
        implementation::decrypt_with_key(&decryption_key, actual_ciphertext)
            .map_err(|e| CryptoError::DecryptionFailed(e.to_string()))
    }
}

/// A collection of keypairs.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct KeySpace {
    pub name: String,
    pub keypairs: HashMap<String, KeyPair>,
}

impl KeySpace {
    /// Creates a new key space with the given name.
    pub fn new(name: &str) -> Self {
        KeySpace {
            name: name.to_string(),
            keypairs: HashMap::new(),
        }
    }

    /// Adds a new keypair to the space.
    pub fn add_keypair(&mut self, name: &str) -> Result<(), CryptoError> {
        if self.keypairs.contains_key(name) {
            return Err(CryptoError::KeypairAlreadyExists(name.to_string()));
        }

        let keypair = KeyPair::new(name);
        self.keypairs.insert(name.to_string(), keypair);
        Ok(())
    }

    /// Gets a keypair by name.
    pub fn get_keypair(&self, name: &str) -> Result<&KeyPair, CryptoError> {
        self.keypairs
            .get(name)
            .ok_or(CryptoError::KeypairNotFound(name.to_string()))
    }

    /// Lists all keypair names in the space.
    pub fn list_keypairs(&self) -> Vec<String> {
        self.keypairs.keys().cloned().collect()
    }
}