dotenvx-rs 0.4.31

use crate::commands::trim_private_key;
use aes_gcm::aead::OsRng;
use aes_gcm::aead::{Aead, KeyInit};
use aes_gcm::{Aes256Gcm, Key, Nonce};
use argon2::password_hash::SaltString;
use argon2::{self, Argon2, PasswordHasher};
use base64ct::{Base64, Base64UrlUnpadded, Encoding};
use colored::Colorize;
use dotenvx_rs::dotenvx::get_private_key;
use ecies::utils::generate_keypair;
use ecies::{PublicKey, SecretKey};
use k256::ecdsa::{SigningKey, VerifyingKey, Signature as EcdsaSignature, signature::hazmat::PrehashVerifier};
use native_tls::{HandshakeError, TlsConnector};
use serde_json::json;
use sha2::{Digest, Sha256};
use std::fs;
use std::fs::File;
use std::io::Write;
use std::net::TcpStream;
use std::path::Path;
use totp_rs::TOTP;

pub struct EcKeyPair {
    pub public_key: PublicKey,
    pub secret_key: SecretKey,
}

impl EcKeyPair {
    pub fn generate() -> Self {
        let (sk, pk) = generate_keypair();
        EcKeyPair {
            public_key: pk,
            secret_key: sk,
        }
    }

    pub fn from_secret_key(sk_hex: &str) -> Self {
        let sk_bytes = hex::decode(check_sk_hex(&trim_private_key(sk_hex.to_string()))).unwrap();
        let sk = SecretKey::parse_slice(&sk_bytes).unwrap();
        let pk = PublicKey::from_secret_key(&sk);
        EcKeyPair {
            public_key: pk,
            secret_key: sk,
        }
    }

    pub fn from_input(sk_hex: &str) -> anyhow::Result<Self> {
        let sk_bytes = hex::decode(trim_private_key(sk_hex.to_string()))?;
        let sk =
            SecretKey::parse_slice(&sk_bytes).map_err(|_| anyhow::anyhow!("Invalid secret key"))?;
        let pk = PublicKey::from_secret_key(&sk);
        Ok(EcKeyPair {
            public_key: pk,
            secret_key: sk,
        })
    }

    pub fn get_pk_hex(&self) -> String {
        let pk_compressed_bytes = self.public_key.serialize_compressed();
        hex::encode(pk_compressed_bytes)
    }

    pub fn get_sk_hex(&self) -> String {
        let sk_bytes = self.secret_key.serialize();
        hex::encode(sk_bytes)
    }
}

pub fn check_sk_hex(sk_hex: &str) -> &str {
    let key_len = sk_hex.len();
    if key_len > 64 {
        &sk_hex[(key_len - 64)..]
    } else {
        sk_hex
    }
}

pub fn encrypt_env_item(
    public_key: &str,
    value_plain: &str,
) -> Result<String, Box<dyn std::error::Error>> {
    let pk_bytes = hex::decode(public_key).unwrap();
    let encrypted_bytes = ecies::encrypt(&pk_bytes, value_plain.as_bytes()).unwrap();
    let base64_text = Base64::encode_string(&encrypted_bytes);
    Ok(format!("encrypted:{base64_text}"))
}

pub fn decrypt_env_item(
    private_key: &str,
    encrypted_text: &str,
) -> Result<String, Box<dyn std::error::Error>> {
    let encrypted_bytes = if encrypted_text.starts_with("encrypted:") {
        Base64::decode_vec(encrypted_text.strip_prefix("encrypted:").unwrap()).unwrap()
    } else {
        Base64::decode_vec(encrypted_text).unwrap()
    };
    let sk = hex::decode(check_sk_hex(private_key)).unwrap();
    let decrypted_bytes = ecies::decrypt(&sk, &encrypted_bytes).unwrap();
    Ok(String::from_utf8(decrypted_bytes)?)
}

pub fn decrypt_value(profile: &Option<String>, encrypted_value: &str) {
    if let Ok(private_key) = get_private_key(&None, profile) {
        if let Ok(plain_text) = decrypt_env_item(check_sk_hex(&private_key), encrypted_value) {
            println!("{plain_text}");
        } else {
            eprintln!(
                "{}",
                "Failed to decrypt the value, please check the private key and profile.".red()
            );
        }
    } else {
        eprintln!("{}",
                  "Private key not found, please check the DOTENV_PRIVATE_KEY environment variable or '.env.key' file.".red()
        );
    }
}

pub fn sha256(input: &[u8]) -> String {
    let mut hasher = Sha256::new();
    hasher.update(input);
    let result = hasher.finalize();
    hex::encode(result)
}

#[allow(dead_code)]
fn get_https_cert_sha256(host: &str, port: u16) -> anyhow::Result<String> {
    let connector = TlsConnector::new()?;
    let stream = TcpStream::connect(format!("{host}:{port}"))?;

    match connector.connect(host, stream) {
        Ok(tls_stream) => {
            if let Some(cert) = tls_stream.peer_certificate()? {
                let der = cert.to_der()?;
                return Ok(sha256(&der));
            }
        }
        Err(HandshakeError::Failure(e)) => {
            return Err(anyhow::anyhow!(e.to_string()));
        }
        Err(HandshakeError::WouldBlock(_)) => {
            return Err(anyhow::anyhow!("Failed to get certificate"));
        }
    }
    Err(anyhow::anyhow!("Failed to get certificate"))
}

/// trim the message and sign it using the private key and return the signature in base64 format
pub fn sign_message(private_key: &str, message: &str) -> anyhow::Result<String> {
    // Step 1: Hash the message using SHA-256
    let mut hasher = Sha256::new();
    hasher.update(message.trim());
    let message_hash = hasher.finalize();
    // Step 2: Sign the message hash with the private key
    let sk_bytes = hex::decode(check_sk_hex(private_key))?;
    let signing_key = SigningKey::from_slice(&sk_bytes)
        .map_err(|_| anyhow::anyhow!("Invalid private key format"))?;
    let (signature, _) = signing_key.sign_prehash_recoverable(&message_hash)
        .map_err(|_| anyhow::anyhow!("Signing failed"))?;
    Ok(Base64::encode_string(&signature.to_bytes()))
}

/// trim the message and sign it using the private key and return the signature in bytes format
pub fn sign_message_bytes(private_key: &str, message: &str) -> anyhow::Result<Vec<u8>> {
    // Step 1: Hash the message using SHA-256
    let mut hasher = Sha256::new();
    hasher.update(message.trim());
    let message_hash = hasher.finalize();
    // Step 2: Sign the message hash with the private key
    let sk_bytes = hex::decode(check_sk_hex(private_key))?;
    let signing_key = SigningKey::from_slice(&sk_bytes)
        .map_err(|_| anyhow::anyhow!("Invalid private key format"))?;
    let (signature, _) = signing_key.sign_prehash_recoverable(&message_hash)
        .map_err(|_| anyhow::anyhow!("Signing failed"))?;
    Ok(signature.to_bytes().to_vec())
}

/// trim the message and verify the signature using the public key
pub fn verify_signature(public_key: &str, message: &str, signature: &str) -> anyhow::Result<bool> {
    // Step 1: Hash the message using SHA-256
    let mut hasher = Sha256::new();
    hasher.update(message.trim());
    let message_hash = hasher.finalize();
    // Step 2: Verify the signature with the public key
    let pk_bytes = hex::decode(public_key)?;
    let verifying_key = VerifyingKey::from_sec1_bytes(&pk_bytes)
        .map_err(|_| anyhow::anyhow!("Invalid public key format"))?;
    let signature_bytes = Base64::decode_vec(signature)?;
    let sig = EcdsaSignature::from_slice(&signature_bytes)
        .map_err(|_| anyhow::anyhow!("Invalid signature format"))?;
    verifying_key.verify_prehash(&message_hash, &sig)
        .map_err(|_| anyhow::anyhow!("Signature verification failed"))?;
    Ok(true)
}

/// generate a JWT token using the private key and claims, and algorithm ES256K(secp256k1)
pub fn generate_jwt_token(
    private_key_hex: &str,
    claims: serde_json::Value,
) -> anyhow::Result<String> {
    let header_obj = json!({"typ": "JWT","alg": "ES256K"});
    let header = Base64UrlUnpadded::encode_string(serde_json::to_string(&header_obj)?.as_bytes());
    let payload = Base64UrlUnpadded::encode_string(serde_json::to_string(&claims)?.as_bytes());
    let message = format!("{header}.{payload}");
    let signature_bytes = sign_message_bytes(private_key_hex, &message)?;
    let signature = Base64UrlUnpadded::encode_string(signature_bytes.as_slice());
    Ok(format!("{header}.{payload}.{signature}"))
}

//============= aes_gcm =======
pub fn encrypt_file<P: AsRef<Path>>(
    input_file: P,
    output_file: P,
    password: &str,
) -> anyhow::Result<()> {
    let plain_bytes = std::fs::read(input_file)?;
    // password hashing with Argon2
    let argon2 = Argon2::default();
    let salt = SaltString::generate(&mut OsRng);
    let password_hash = argon2.hash_password(password.as_bytes(), &salt).unwrap();
    let hash = password_hash.hash.unwrap();

    // Initialize AES-GCM with the derived key
    let aes_key = Key::<Aes256Gcm>::from_slice(hash.as_bytes()); // Use the first 32 bytes of the hash
    let cipher = Aes256Gcm::new(aes_key);

    // Generate a random nonce
    let random_nonce = rand::random::<[u8; 12]>();
    // Encrypt the plaintext
    let ciphertext = cipher
        .encrypt(Nonce::from_slice(&random_nonce), plain_bytes.as_ref())
        .expect("encryption failure!");

    // // Write the salt, nonce, and ciphertext to the output file
    let mut output = File::create(output_file)?;
    let mut salt_bytes: [u8; 16] = [0; 16];
    salt.decode_b64(&mut salt_bytes).unwrap();
    output.write_all(&salt_bytes)?; // First 16 bytes: salt
    output.write_all(&random_nonce)?; // Next 12 bytes: nonce
    output.write_all(&ciphertext)?; // Remaining bytes: ciphertext
    Ok(())
}

pub fn decrypt_file<P: AsRef<Path>>(
    encrypted_file: P,
    output_file: P,
    password: &str,
) -> anyhow::Result<()> {
    // Read the encrypted file
    let encrypted_file_content = fs::read(encrypted_file)?;

    // Extract the salt, nonce, and ciphertext
    let salt_bytes = &encrypted_file_content[0..16]; // First 16 bytes: salt
    let salt = SaltString::encode_b64(salt_bytes).unwrap();
    let nonce_bytes = &encrypted_file_content[16..28]; // Next 12 bytes: nonce
    let ciphertext = &encrypted_file_content[28..]; // Remaining bytes: ciphertext

    // password hashing with Argon2
    let argon2 = Argon2::default();
    let password_hash = argon2.hash_password(password.as_bytes(), &salt).unwrap();
    let hash = password_hash.hash.unwrap();

    // Initialize AES-GCM with the derived key
    let aes_key = Key::<Aes256Gcm>::from_slice(hash.as_bytes());
    let cipher = Aes256Gcm::new(aes_key);

    // Decrypt the ciphertext
    let plain_bytes = cipher
        .decrypt(Nonce::from_slice(nonce_bytes), ciphertext)
        .expect("decryption failure!");

    // Write the decrypted bytes to the output file
    fs::write(output_file, plain_bytes)?;
    Ok(())
}

pub fn generate_totp_password(totp_url: &str) -> anyhow::Result<String> {
    let totp = TOTP::from_url(totp_url)?;
    totp.generate_current()
        .map_err(|e| anyhow::anyhow!(e.to_string()))
}

#[cfg(test)]
mod tests {
    use super::*;
    use testresult::TestResult;

    #[test]
    fn test_signature_and_verify() {
        let public_key = "02f6e5c1a348cd70ee9ebcdf271892878d83d7bb9c1cd0644ae0da0a04904b83e4";
        let private_key = "00d3a1bf3a9a989e3ae11a58e89d95e26d32e0445d825d7fe7ef162ffad2706580";
        let message = "Hello, secp256k1!";
        // Sign the message
        let signature = sign_message(private_key, message).unwrap();
        println!("Signature: {signature}");
        let verify_result = verify_signature(public_key, message, &signature).unwrap();
        assert!(verify_result, "Signature verification failed");
    }

    #[test]
    fn test_jwt_generate() {
        use chrono::Utc;
        let now = Utc::now().timestamp();
        let private_key = "c81efd721a711661296a53b768c780e0d9ec9d597e49d8ed53eed0b638b958cf";
        let claims = json!({
            "sub": "linux-china",
            "kid": "b895c58f944855010fa88f7a76642e2005d51705cb27597c7a85347628ac5dcf",
            "exp": now + 60*60*24*365, // Expiration time (e.g., 2026-01-01T00:00:00Z)
            "iat": now, // now (e.g., 2021-06-01T00:00:00Z)
            "iss": "dotenvx"
        });
        let jwt_token = generate_jwt_token(private_key, claims).unwrap();
        println!("JWT: {jwt_token}");
    }

    #[test]
    fn test_encrypt_file() -> TestResult {
        // Input file and password
        let input_file = "tests/example.txt";
        let output_file = "tests/example.txt.aes";
        let password = "your_secure_password";
        // Encrypt the file
        encrypt_file(input_file, output_file, password).unwrap();
        Ok(())
    }

    #[test]
    fn test_decrypt_file() -> TestResult {
        // Input file and password
        let encrypted_file = "tests/example.txt.aes";
        let output_file = "tests/example.txt";
        let password = "your_secure_password";
        // Encrypt the file
        decrypt_file(encrypted_file, output_file, password).unwrap();
        Ok(())
    }

    #[test]
    fn test_https_cert() {
        let finger_print = get_https_cert_sha256("dotenvx.microservices.club", 443).unwrap();
        println!("finger_print : {}", finger_print);
    }
}