synor/apps/cli/src/wallet.rs

//! Wallet management.
//!
//! All wallets use Hybrid keys (Ed25519 + Dilithium) for quantum-resistant security.

use std::fs;
use std::path::Path;

use aes_gcm::{
    aead::{Aead, KeyInit},
    Aes256Gcm, Nonce,
};
use argon2::{Argon2, Params};
use rand::RngCore;
use serde::{Deserialize, Serialize};
use synor_crypto::{HybridKeypair, Mnemonic, Network};

/// Wallet data.
///
/// All Synor wallets use Hybrid keys combining Ed25519 (classical) and
/// Dilithium (post-quantum) for maximum security against both classical
/// and quantum attacks.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Wallet {
    /// Wallet name.
    pub name: String,

    /// Network (mainnet, testnet).
    pub network: String,

    /// Encrypted seed (if HD wallet).
    /// Format: salt (16 bytes) || nonce (12 bytes) || ciphertext
    pub encrypted_seed: Option<String>,

    /// Addresses.
    pub addresses: Vec<WalletAddress>,

    /// Creation timestamp.
    pub created_at: u64,
}

/// Wallet address.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct WalletAddress {
    /// Address string.
    pub address: String,

    /// Derivation path (for HD wallets).
    pub path: Option<String>,

    /// Encrypted private key (Ed25519 component).
    /// Format: salt (16 bytes) || nonce (12 bytes) || ciphertext
    pub encrypted_ed25519_key: String,

    /// Encrypted private key (Dilithium component).
    /// Format: salt (16 bytes) || nonce (12 bytes) || ciphertext
    pub encrypted_dilithium_key: String,

    /// Public key (Ed25519 component, hex).
    pub ed25519_public_key: String,

    /// Public key (Dilithium component, hex).
    pub dilithium_public_key: String,

    /// Label.
    pub label: Option<String>,

    /// Is default address.
    pub is_default: bool,
}

/// Current encryption key for session (not persisted).
static mut CURRENT_PASSWORD: Option<Vec<u8>> = None;

impl Wallet {
    /// Creates a new wallet with Hybrid keys.
    ///
    /// Returns (wallet, mnemonic_phrase) so user can back up the phrase.
    pub fn create(name: &str, network: &str, password: &str) -> anyhow::Result<(Self, String)> {
        // Generate mnemonic (24 words for maximum security)
        let mnemonic = Mnemonic::generate(24)
            .map_err(|e| anyhow::anyhow!("Failed to generate mnemonic: {}", e))?;
        let phrase = mnemonic.phrase().to_string();

        // Derive seed from mnemonic
        let seed = mnemonic.to_seed("");

        // Store password for session
        set_password(password);

        // Generate first address (always Hybrid)
        let net = parse_network(network)?;
        let addr = generate_hybrid_address(&seed, 0, net, password)?;

        let wallet = Wallet {
            name: name.to_string(),
            network: network.to_string(),
            encrypted_seed: Some(encrypt_data(&seed, password)?),
            addresses: vec![WalletAddress {
                address: addr.address,
                path: Some("m/44'/21337'/0'/0/0".to_string()),
                encrypted_ed25519_key: addr.encrypted_ed25519_key,
                encrypted_dilithium_key: addr.encrypted_dilithium_key,
                ed25519_public_key: addr.ed25519_public_key,
                dilithium_public_key: addr.dilithium_public_key,
                label: Some("Default".to_string()),
                is_default: true,
            }],
            created_at: current_timestamp(),
        };

        Ok((wallet, phrase))
    }

    /// Imports wallet from seed phrase.
    pub fn import(
        name: &str,
        network: &str,
        seed_phrase: &str,
        password: &str,
    ) -> anyhow::Result<Self> {
        // Validate and parse mnemonic
        let mnemonic = Mnemonic::from_phrase(seed_phrase)
            .map_err(|e| anyhow::anyhow!("Invalid mnemonic phrase: {}", e))?;

        // Derive seed from mnemonic
        let seed = mnemonic.to_seed("");

        // Store password for session
        set_password(password);

        // Generate first address (always Hybrid)
        let net = parse_network(network)?;
        let addr = generate_hybrid_address(&seed, 0, net, password)?;

        let wallet = Wallet {
            name: name.to_string(),
            network: network.to_string(),
            encrypted_seed: Some(encrypt_data(&seed, password)?),
            addresses: vec![WalletAddress {
                address: addr.address,
                path: Some("m/44'/21337'/0'/0/0".to_string()),
                encrypted_ed25519_key: addr.encrypted_ed25519_key,
                encrypted_dilithium_key: addr.encrypted_dilithium_key,
                ed25519_public_key: addr.ed25519_public_key,
                dilithium_public_key: addr.dilithium_public_key,
                label: Some("Default".to_string()),
                is_default: true,
            }],
            created_at: current_timestamp(),
        };

        Ok(wallet)
    }

    /// Loads wallet from file.
    pub fn load(path: &Path) -> anyhow::Result<Self> {
        let content = fs::read_to_string(path)?;
        let wallet: Wallet = serde_json::from_str(&content)?;
        Ok(wallet)
    }

    /// Saves wallet to file.
    pub fn save(&self, path: &Path) -> anyhow::Result<()> {
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)?;
        }

        let content = serde_json::to_string_pretty(self)?;
        fs::write(path, content)?;

        // Set restrictive permissions on Unix
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let perms = fs::Permissions::from_mode(0o600);
            fs::set_permissions(path, perms)?;
        }

        Ok(())
    }

    /// Gets default address.
    pub fn default_address(&self) -> Option<&WalletAddress> {
        self.addresses.iter().find(|a| a.is_default)
    }

    /// Gets all addresses.
    pub fn all_addresses(&self) -> &[WalletAddress] {
        &self.addresses
    }

    /// Generates a new address.
    pub fn new_address(
        &mut self,
        label: Option<String>,
        password: &str,
    ) -> anyhow::Result<&WalletAddress> {
        let seed = self
            .encrypted_seed
            .as_ref()
            .map(|s| decrypt_data(s, password))
            .transpose()?
            .ok_or_else(|| anyhow::anyhow!("No seed in wallet"))?;

        let seed_array: [u8; 64] = seed
            .try_into()
            .map_err(|_| anyhow::anyhow!("Invalid seed length"))?;

        let index = self.addresses.len() as u32;
        let net = parse_network(&self.network)?;
        let addr = generate_hybrid_address(&seed_array, index, net, password)?;

        self.addresses.push(WalletAddress {
            address: addr.address,
            path: Some(format!("m/44'/21337'/0'/0/{}", index)),
            encrypted_ed25519_key: addr.encrypted_ed25519_key,
            encrypted_dilithium_key: addr.encrypted_dilithium_key,
            ed25519_public_key: addr.ed25519_public_key,
            dilithium_public_key: addr.dilithium_public_key,
            label,
            is_default: false,
        });

        Ok(self.addresses.last().unwrap())
    }

    /// Exports seed phrase.
    pub fn export_seed_phrase(&self, password: &str) -> anyhow::Result<String> {
        let seed = self
            .encrypted_seed
            .as_ref()
            .map(|s| decrypt_data(s, password))
            .transpose()?
            .ok_or_else(|| anyhow::anyhow!("No seed in wallet"))?;

        // We can't recover the mnemonic from the seed directly without entropy.
        // For security, we should store the encrypted entropy instead.
        // For now, return an error explaining this limitation.
        anyhow::bail!(
            "Cannot export mnemonic from derived seed. \
             Please use the original mnemonic phrase you wrote down during wallet creation. \
             Seed hex: {}",
            hex::encode(&seed)
        )
    }

    /// Unlocks the wallet for signing operations.
    pub fn unlock(&self, password: &str) -> anyhow::Result<()> {
        // Verify password by trying to decrypt the seed
        if let Some(ref encrypted) = self.encrypted_seed {
            decrypt_data(encrypted, password)?;
        }
        set_password(password);
        Ok(())
    }

    /// Signs a transaction with hybrid signature (Ed25519 + Dilithium).
    pub fn sign_transaction(
        &self,
        address: &str,
        tx_bytes: &[u8],
        password: &str,
    ) -> anyhow::Result<HybridSignatureBytes> {
        let addr = self
            .addresses
            .iter()
            .find(|a| a.address == address)
            .ok_or_else(|| anyhow::anyhow!("Address not found in wallet"))?;

        // Decrypt Ed25519 private key
        let ed25519_seed = decrypt_data(&addr.encrypted_ed25519_key, password)?;
        let _ed25519_seed: [u8; 32] = ed25519_seed
            .try_into()
            .map_err(|_| anyhow::anyhow!("Invalid Ed25519 key length"))?;

        // We need to reconstruct the keypair to sign
        // For this, we need the full 64-byte seed. Let's derive from the wallet seed.
        let wallet_seed = self
            .encrypted_seed
            .as_ref()
            .map(|s| decrypt_data(s, password))
            .transpose()?
            .ok_or_else(|| anyhow::anyhow!("No seed in wallet"))?;

        let wallet_seed: [u8; 64] = wallet_seed
            .try_into()
            .map_err(|_| anyhow::anyhow!("Invalid seed length"))?;

        // Find the index of this address
        let index = self
            .addresses
            .iter()
            .position(|a| a.address == address)
            .ok_or_else(|| anyhow::anyhow!("Address not found"))? as u32;

        // Derive the keypair for this index
        let derived_seed = derive_key_at_index(&wallet_seed, index);
        let keypair = HybridKeypair::from_seed(&derived_seed)
            .map_err(|e| anyhow::anyhow!("Failed to derive keypair: {:?}", e))?;

        // Sign the transaction
        let signature = keypair.sign(tx_bytes);

        Ok(HybridSignatureBytes {
            ed25519: signature.ed25519_signature.to_vec(),
            dilithium: signature.dilithium_signature.clone(),
        })
    }
}

/// Hybrid signature bytes (Ed25519 + Dilithium).
pub struct HybridSignatureBytes {
    /// Ed25519 signature (64 bytes).
    pub ed25519: Vec<u8>,
    /// Dilithium signature (~2420 bytes).
    pub dilithium: Vec<u8>,
}

impl HybridSignatureBytes {
    /// Returns the combined signature bytes.
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(self.ed25519.len() + self.dilithium.len());
        bytes.extend_from_slice(&self.ed25519);
        bytes.extend_from_slice(&self.dilithium);
        bytes
    }
}

// ==================== Helper Functions ====================

/// Generated hybrid address data.
struct HybridAddressData {
    address: String,
    encrypted_ed25519_key: String,
    encrypted_dilithium_key: String,
    ed25519_public_key: String,
    dilithium_public_key: String,
}

fn parse_network(network: &str) -> anyhow::Result<Network> {
    match network.to_lowercase().as_str() {
        "mainnet" => Ok(Network::Mainnet),
        "testnet" => Ok(Network::Testnet),
        "devnet" => Ok(Network::Devnet),
        _ => anyhow::bail!("Unknown network: {}", network),
    }
}

fn derive_key_at_index(master_seed: &[u8; 64], index: u32) -> [u8; 64] {
    // Use HKDF-like derivation to get a unique seed for each index
    let mut derived = [0u8; 64];

    // Derive first 32 bytes (for Ed25519)
    let mut input1 = Vec::with_capacity(68);
    input1.extend_from_slice(&master_seed[..32]);
    input1.extend_from_slice(&index.to_le_bytes());
    let hash1: [u8; 32] = blake3::hash(&input1).into();
    derived[..32].copy_from_slice(&hash1);

    // Derive second 32 bytes (for Dilithium)
    let mut input2 = Vec::with_capacity(68);
    input2.extend_from_slice(&master_seed[32..64]);
    input2.extend_from_slice(&index.to_le_bytes());
    let hash2: [u8; 32] = blake3::hash(&input2).into();
    derived[32..64].copy_from_slice(&hash2);

    derived
}

fn generate_hybrid_address(
    seed: &[u8; 64],
    index: u32,
    network: Network,
    password: &str,
) -> anyhow::Result<HybridAddressData> {
    // Derive seed for this index
    let derived_seed = derive_key_at_index(seed, index);

    // Generate hybrid keypair from derived seed
    let keypair = HybridKeypair::from_seed(&derived_seed)
        .map_err(|e| anyhow::anyhow!("Failed to generate keypair: {:?}", e))?;

    // Get public keys
    let pubkey = keypair.public_key();
    let ed25519_public_key = hex::encode(pubkey.ed25519_bytes());
    let dilithium_public_key = hex::encode(pubkey.dilithium_bytes());

    // Get address
    let address = keypair.address(network).to_string();

    // Encrypt private keys
    let secret = keypair.secret_key();
    let encrypted_ed25519_key = encrypt_data(secret.ed25519_seed(), password)?;
    let encrypted_dilithium_key = encrypt_data(&derived_seed[32..64], password)?;

    Ok(HybridAddressData {
        address,
        encrypted_ed25519_key,
        encrypted_dilithium_key,
        ed25519_public_key,
        dilithium_public_key,
    })
}

/// Encrypts data using AES-256-GCM with Argon2 key derivation.
///
/// Output format: salt (16 bytes) || nonce (12 bytes) || ciphertext
fn encrypt_data(data: &[u8], password: &str) -> anyhow::Result<String> {
    // Generate random salt and nonce
    let mut salt = [0u8; 16];
    let mut nonce_bytes = [0u8; 12];
    rand::thread_rng().fill_bytes(&mut salt);
    rand::thread_rng().fill_bytes(&mut nonce_bytes);

    // Derive encryption key using Argon2id
    let key = derive_encryption_key(password.as_bytes(), &salt)?;

    // Encrypt with AES-256-GCM
    let cipher = Aes256Gcm::new_from_slice(&key)
        .map_err(|e| anyhow::anyhow!("Failed to create cipher: {}", e))?;
    let nonce = Nonce::from_slice(&nonce_bytes);
    let ciphertext = cipher
        .encrypt(nonce, data)
        .map_err(|e| anyhow::anyhow!("Encryption failed: {}", e))?;

    // Combine: salt || nonce || ciphertext
    let mut result = Vec::with_capacity(16 + 12 + ciphertext.len());
    result.extend_from_slice(&salt);
    result.extend_from_slice(&nonce_bytes);
    result.extend_from_slice(&ciphertext);

    Ok(hex::encode(result))
}

/// Decrypts data encrypted with encrypt_data.
fn decrypt_data(encrypted_hex: &str, password: &str) -> anyhow::Result<Vec<u8>> {
    let encrypted = hex::decode(encrypted_hex)?;

    if encrypted.len() < 28 {
        // 16 (salt) + 12 (nonce) = 28 minimum
        anyhow::bail!("Invalid encrypted data: too short");
    }

    // Extract salt, nonce, and ciphertext
    let salt = &encrypted[..16];
    let nonce_bytes = &encrypted[16..28];
    let ciphertext = &encrypted[28..];

    // Derive encryption key using Argon2id
    let key = derive_encryption_key(password.as_bytes(), salt)?;

    // Decrypt with AES-256-GCM
    let cipher = Aes256Gcm::new_from_slice(&key)
        .map_err(|e| anyhow::anyhow!("Failed to create cipher: {}", e))?;
    let nonce = Nonce::from_slice(nonce_bytes);
    let plaintext = cipher
        .decrypt(nonce, ciphertext)
        .map_err(|_| anyhow::anyhow!("Decryption failed: invalid password or corrupted data"))?;

    Ok(plaintext)
}

/// Derives a 32-byte encryption key from password using Argon2id.
fn derive_encryption_key(password: &[u8], salt: &[u8]) -> anyhow::Result<[u8; 32]> {
    // Argon2id parameters (OWASP recommendations)
    let params = Params::new(
        65536, // 64 MiB memory
        3,     // 3 iterations
        4,     // 4 parallel threads
        Some(32),
    )
    .map_err(|e| anyhow::anyhow!("Invalid Argon2 params: {}", e))?;

    let argon2 = Argon2::new(argon2::Algorithm::Argon2id, argon2::Version::V0x13, params);

    let mut key = [0u8; 32];
    argon2
        .hash_password_into(password, salt, &mut key)
        .map_err(|e| anyhow::anyhow!("Key derivation failed: {}", e))?;

    Ok(key)
}

fn set_password(password: &str) {
    unsafe {
        CURRENT_PASSWORD = Some(password.as_bytes().to_vec());
    }
}

fn current_timestamp() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs()
}

/// Lists wallets in directory.
pub fn list_wallets(dir: &Path) -> anyhow::Result<Vec<String>> {
    if !dir.exists() {
        return Ok(vec![]);
    }

    let mut wallets = Vec::new();
    for entry in fs::read_dir(dir)? {
        let entry = entry?;
        let path = entry.path();
        if path.extension().is_some_and(|e| e == "wallet") {
            if let Some(name) = path.file_stem() {
                wallets.push(name.to_string_lossy().to_string());
            }
        }
    }

    Ok(wallets)
}

/// Validates a mnemonic phrase.
pub fn validate_mnemonic(phrase: &str) -> bool {
    Mnemonic::validate(phrase)
}

/// Suggests word completions for mnemonic entry.
pub fn suggest_word(partial: &str) -> Vec<&'static str> {
    synor_crypto::mnemonic::suggest_word(partial)
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    #[test]
    fn test_wallet_create() {
        let (wallet, phrase) = Wallet::create("test", "mainnet", "testpassword123").unwrap();
        assert_eq!(wallet.name, "test");
        assert_eq!(wallet.addresses.len(), 1);
        assert!(wallet.addresses[0].is_default);
        // Address should be hybrid format
        assert!(wallet.addresses[0].address.starts_with("synor1"));
        // Phrase should be 24 words
        assert_eq!(phrase.split_whitespace().count(), 24);
    }

    #[test]
    fn test_wallet_import() {
        let phrase = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon art";
        let wallet = Wallet::import("test", "mainnet", phrase, "testpassword123").unwrap();
        assert_eq!(wallet.name, "test");
        assert_eq!(wallet.addresses.len(), 1);
    }

    #[test]
    fn test_wallet_save_load() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("test.wallet");

        let (wallet, _) = Wallet::create("test", "mainnet", "testpassword123").unwrap();
        wallet.save(&path).unwrap();

        let loaded = Wallet::load(&path).unwrap();
        assert_eq!(loaded.name, wallet.name);
        assert_eq!(loaded.addresses.len(), wallet.addresses.len());
    }

    #[test]
    fn test_new_address() {
        let (mut wallet, _) = Wallet::create("test", "mainnet", "testpassword123").unwrap();
        let addr = wallet
            .new_address(Some("Second".to_string()), "testpassword123")
            .unwrap()
            .clone();

        assert_eq!(wallet.addresses.len(), 2);
        assert!(!addr.is_default);
        // All addresses should be hybrid format
        assert!(addr.address.starts_with("synor1"));
    }

    #[test]
    fn test_sign_transaction() {
        let (wallet, _) = Wallet::create("test", "mainnet", "testpassword123").unwrap();
        let default_addr = wallet.default_address().unwrap();

        let tx_data = b"test transaction";
        let sig = wallet
            .sign_transaction(&default_addr.address, tx_data, "testpassword123")
            .unwrap();

        // Hybrid signature has both components
        assert_eq!(sig.ed25519.len(), 64);
        assert!(!sig.dilithium.is_empty());
    }

    #[test]
    fn test_encryption_decryption() {
        let data = b"secret data to encrypt";
        let password = "strong_password_123";

        let encrypted = encrypt_data(data, password).unwrap();
        let decrypted = decrypt_data(&encrypted, password).unwrap();

        assert_eq!(decrypted, data);
    }

    #[test]
    fn test_wrong_password() {
        let data = b"secret data";
        let encrypted = encrypt_data(data, "correct_password").unwrap();

        let result = decrypt_data(&encrypted, "wrong_password");
        assert!(result.is_err());
    }

    #[test]
    fn test_validate_mnemonic() {
        assert!(validate_mnemonic(
            "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about"
        ));
        assert!(!validate_mnemonic("invalid phrase here"));
    }
}