anttp 0.26.0

use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use utoipa::ToSchema;
use saorsa_pqc::api::sig::MlDsaVariant;
use saorsa_pqc::ml_dsa_65;
use saorsa_pqc::api::{MlDsaPublicKey, MlDsaSignature};
use crate::config::anttp_config::AntTpConfig;

#[derive(Debug, Serialize, Deserialize, ToSchema, Clone)]
pub struct Crypto {
    pub signature: Option<String>,
    pub verified: Option<bool>,
}

#[derive(Debug, Serialize, Deserialize, ToSchema, Clone)]
pub struct CryptoContent {
    pub content: Option<String>,
}

#[derive(Debug, Clone)]
pub struct CryptoService {
    ant_tp_config: AntTpConfig,
}

impl CryptoService {
    pub fn new(ant_tp_config: AntTpConfig) -> Self {
        Self { ant_tp_config }
    }

    pub fn sign_map(&self, mut data_map: HashMap<String, Crypto>) -> HashMap<String, Crypto> {
        for (data_hex, crypto_struct) in data_map.iter_mut() {
            match self.sign(data_hex) {
                Some(signature_hex) => {
                    crypto_struct.signature = Some(signature_hex);
                    crypto_struct.verified = Some(true);
                }
                None => {
                    crypto_struct.verified = Some(false);
                }
            }
        }
        data_map
    }

    pub fn sign(&self, data_hex: &str) -> Option<String> {
        match self.ant_tp_config.get_app_private_key() {
            Ok(app_private_key) => {
                let dsa = ml_dsa_65();
                match hex::decode(data_hex) {
                    Ok(data_bytes) => {
                        let signature = dsa.sign(&app_private_key, data_bytes.as_slice());
                        Some(hex::encode(signature.unwrap().to_bytes()))
                    }
                    Err(_) => None,
                }
            }
            Err(_) => None,
        }
    }

    pub fn verify_map(&self, public_key: String, mut data_map: HashMap<String, Crypto>) -> HashMap<String, Crypto> {
        for (data_hex, crypto_struct) in data_map.iter_mut() {
            let signature = crypto_struct.signature.clone().unwrap_or_default();
            let is_verified = self.verify(&public_key, &signature, data_hex);
            crypto_struct.verified = Some(is_verified);
        }
        data_map
    }

    pub fn verify(&self, public_key_hex: &str, signature_hex: &str, data_hex: &str) -> bool {
        let public_key_bytes = match hex::decode(public_key_hex) {
            Ok(bytes) => bytes,
            Err(_) => return false,
        };
        let signature_bytes = match hex::decode(signature_hex) {
            Ok(bytes) => bytes,
            Err(_) => return false,
        };
        let data_bytes = match hex::decode(data_hex) {
            Ok(bytes) => bytes,
            Err(_) => return false,
        };

        if public_key_bytes.len() != 48 {
            return false;
        }
        let public_key = match MlDsaPublicKey::from_bytes(MlDsaVariant::MlDsa65, public_key_bytes.as_slice()) {
            Ok(pk) => pk,
            Err(_) => return false,
        };

        if signature_bytes.len() != 96 {
            return false;
        }
        let signature = match MlDsaSignature::from_bytes(MlDsaVariant::MlDsa65, signature_bytes.as_slice()) {
            Ok(sig) => sig,
            Err(_) => return false,
        };

        let dsa = ml_dsa_65();
        dsa.verify(&public_key, &data_bytes, &signature).unwrap()
    }

    /*pub fn encrypt_map(&self, public_key: String, mut data_map: HashMap<String, CryptoContent>) -> HashMap<String, CryptoContent> {
        for (data_base64, crypto_content_struct) in data_map.iter_mut() {
            match self.encrypt(&public_key, data_base64) {
                Some(encrypted_bytes) => {
                    crypto_content_struct.content = Some(general_purpose::STANDARD.encode(encrypted_bytes));
                }
                None => {
                    crypto_content_struct.content = None;
                }
            }
        }
        data_map
    }*/

    /*pub fn encrypt(&self, public_key_hex: &str, data_base64: &str) -> Option<Vec<u8>> {
        let public_key_bytes = match hex::decode(public_key_hex) {
            Ok(bytes) => bytes,
            Err(_) => return None,
        };
        let data_bytes = match general_purpose::STANDARD.decode(data_base64) {
            Ok(bytes) => bytes,
            Err(_) => return None,
        };

        if public_key_bytes.len() != 48 {
            return None;
        }
        let public_key = match MlDsaPublicKey::from_bytes(MlDsaVariant::MlDsa65, public_key_bytes.as_slice()) {
            Ok(pk) => pk,
            Err(_) => return None,
        };
        
        Some(public_key.encrypt(&data_bytes).to_bytes())
    }*/

    /*pub fn decrypt_map(&self, mut data_map: HashMap<String, CryptoContent>) -> HashMap<String, CryptoContent> {
        for (data_base64, crypto_content_struct) in data_map.iter_mut() {
            match self.decrypt(data_base64) {
                Some(decrypted_bytes) => {
                    crypto_content_struct.content = Some(general_purpose::STANDARD.encode(decrypted_bytes));
                }
                None => {
                    crypto_content_struct.content = None;
                }
            }
        }
        data_map
    }*/

    /*pub fn decrypt(&self, data_base64: &str) -> Option<Vec<u8>> {
        let app_private_key = match self.ant_tp_config.get_app_private_key() {
            Ok(key) => key,
            Err(_) => return None,
        };
        let data_bytes = match general_purpose::STANDARD.decode(data_base64) {
            Ok(bytes) => bytes,
            Err(_) => return None,
        };

        let ciphertext = match blsttc::Ciphertext::from_bytes(&data_bytes) {
            Ok(ct) => ct,
            Err(_) => return None,
        };

        app_private_key.decrypt(&ciphertext)
    }*/
}

/*#[cfg(test)]
mod tests {
    use super::*;
    use clap::Parser;

    #[test]
    fn test_verify_success() {
        let secret_key = SecretKey::random();
        let public_key = hex::encode(secret_key.public_key().to_bytes());
        let data = b"test data";
        let data_hex = hex::encode(data);
        let signature = hex::encode(secret_key.sign(data).to_bytes());

        let mut data_map = HashMap::new();
        data_map.insert(data_hex.clone(), Crypto {
            signature: Some(signature),
            verified: None,
        });

        let ant_tp_config = AntTpConfig::parse_from(&["anttp"]);
        let service = CryptoService::new(ant_tp_config);
        let result = service.verify_map(public_key, data_map);

        assert!(result.get(&data_hex).unwrap().verified.unwrap());
    }

    #[test]
    fn test_verify_failure() {
        let secret_key = SecretKey::random();
        let public_key = hex::encode(secret_key.public_key().to_bytes());
        let data_hex = hex::encode(b"test data");
        let signature = hex::encode(secret_key.sign(b"other data").to_bytes());

        let mut data_map = HashMap::new();
        data_map.insert(data_hex.clone(), Crypto {
            signature: Some(signature),
            verified: None,
        });

        let ant_tp_config = AntTpConfig::parse_from(&["anttp"]);
        let service = CryptoService::new(ant_tp_config);
        let result = service.verify_map(public_key, data_map);

        assert!(!result.get(&data_hex).unwrap().verified.unwrap());
    }

    #[test]
    fn test_sign_individual_success() {
        let secret_key = SecretKey::random();
        let app_private_key_hex = secret_key.to_hex();
        let data = b"test data";
        let data_hex = hex::encode(data);

        let ant_tp_config = AntTpConfig::parse_from(&["anttp", "--app-private-key", &app_private_key_hex]);
        let service = CryptoService::new(ant_tp_config);
        let signature = service.sign(&data_hex);

        assert!(signature.is_some());
        let is_verified = service.verify(
            &hex::encode(secret_key.public_key().to_bytes()),
            &signature.unwrap(),
            &data_hex
        );
        assert!(is_verified);
    }

    #[test]
    fn test_sign_success() {
        let secret_key = SecretKey::random();
        let app_private_key_hex = secret_key.to_hex();
        let data = b"test data";
        let data_hex = hex::encode(data);

        let mut data_map = HashMap::new();
        data_map.insert(data_hex.clone(), Crypto {
            signature: None,
            verified: None,
        });

        let ant_tp_config = AntTpConfig::parse_from(&["anttp", "--app-private-key", &app_private_key_hex]);
        let service = CryptoService::new(ant_tp_config);
        let result = service.sign_map(data_map);

        let crypto_struct = result.get(&data_hex).unwrap();
        assert!(crypto_struct.verified.unwrap());
        assert!(crypto_struct.signature.is_some());

        // Crypto the generated signature
        let is_verified = service.verify(
            &hex::encode(secret_key.public_key().to_bytes()),
            crypto_struct.signature.as_ref().unwrap(),
            &data_hex
        );
        assert!(is_verified);
    }

    #[test]
    fn test_encrypt_success() {
        let secret_key = SecretKey::random();
        let public_key_hex = hex::encode(secret_key.public_key().to_bytes());
        let data = b"test data";
        let data_base64 = general_purpose::STANDARD.encode(data);

        let ant_tp_config = AntTpConfig::parse_from(&["anttp"]);
        let service = CryptoService::new(ant_tp_config);
        let encrypted_bytes = service.encrypt(&public_key_hex, &data_base64);

        assert!(encrypted_bytes.is_some());
        let encrypted_bytes = encrypted_bytes.unwrap();
        
        let ciphertext = blsttc::Ciphertext::from_bytes(&encrypted_bytes).unwrap();
        
        let decrypted_data = secret_key.decrypt(&ciphertext).unwrap();
        assert_eq!(decrypted_data, data);
    }

    #[test]
    fn test_encrypt_map_success() {
        let secret_key = SecretKey::random();
        let public_key_hex = hex::encode(secret_key.public_key().to_bytes());
        let data = b"test data";
        let data_base64 = general_purpose::STANDARD.encode(data);

        let mut data_map = HashMap::new();
        data_map.insert(data_base64.clone(), CryptoContent {
            content: None,
        });

        let ant_tp_config = AntTpConfig::parse_from(&["anttp"]);
        let service = CryptoService::new(ant_tp_config);
        let result = service.encrypt_map(public_key_hex, data_map);

        let crypto_content_struct = result.get(&data_base64).unwrap();
        assert!(crypto_content_struct.content.is_some());
        
        let encrypted_base64 = crypto_content_struct.content.as_ref().unwrap();
        let encrypted_bytes = general_purpose::STANDARD.decode(encrypted_base64).unwrap();
        
        let ciphertext = blsttc::Ciphertext::from_bytes(&encrypted_bytes).unwrap();
        
        let decrypted_data = secret_key.decrypt(&ciphertext).unwrap();
        assert_eq!(decrypted_data, data);
    }

    #[test]
    fn test_decrypt_success() {
        let secret_key = SecretKey::random();
        let app_private_key_hex = secret_key.to_hex();
        let data = b"test data";
        let encrypted_data = secret_key.public_key().encrypt(data).to_bytes();
        let encrypted_data_base64 = general_purpose::STANDARD.encode(encrypted_data);

        let ant_tp_config = AntTpConfig::parse_from(&["anttp", "--app-private-key", &app_private_key_hex]);
        let service = CryptoService::new(ant_tp_config);
        let decrypted_bytes = service.decrypt(&encrypted_data_base64);

        assert!(decrypted_bytes.is_some());
        assert_eq!(decrypted_bytes.unwrap(), data);
    }

    #[test]
    fn test_decrypt_map_success() {
        let secret_key = SecretKey::random();
        let app_private_key_hex = secret_key.to_hex();
        let data = b"test data";
        let encrypted_data = secret_key.public_key().encrypt(data).to_bytes();
        let encrypted_data_base64 = general_purpose::STANDARD.encode(encrypted_data);

        let mut data_map = HashMap::new();
        data_map.insert(encrypted_data_base64.clone(), CryptoContent {
            content: None,
        });

        let ant_tp_config = AntTpConfig::parse_from(&["anttp", "--app-private-key", &app_private_key_hex]);
        let service = CryptoService::new(ant_tp_config);
        let result = service.decrypt_map(data_map);

        let decrypted_content_base64 = result.get(&encrypted_data_base64).unwrap().content.as_ref().unwrap();
        let decrypted_bytes = general_purpose::STANDARD.decode(decrypted_content_base64).unwrap();
        assert_eq!(decrypted_bytes, data);
    }
}
*/