synor/apps/desktop-wallet/src-tauri/src/wallet.rs

//! Wallet state management for the desktop wallet

use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::RwLock;
use serde::{Deserialize, Serialize};
use zeroize::Zeroize;

use crate::crypto::{self, EncryptedWallet, SeedData};
use crate::{Error, Result};

/// Wallet file name
const WALLET_FILE: &str = "wallet.json";

/// Represents a derived address in the wallet
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WalletAddress {
    /// Bech32 encoded address
    pub address: String,
    /// Derivation index
    pub index: u32,
    /// Whether this is a change address
    pub is_change: bool,
    /// Human-readable label
    pub label: Option<String>,
}

/// Network connection state
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NetworkConnection {
    /// RPC endpoint URL
    pub rpc_url: String,
    /// WebSocket endpoint URL
    pub ws_url: Option<String>,
    /// Whether currently connected
    pub connected: bool,
    /// Network type (mainnet/testnet)
    pub network: String,
}

/// Persisted wallet metadata (non-sensitive)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WalletMetadata {
    /// Encrypted wallet data
    pub encrypted: EncryptedWallet,
    /// Network (mainnet/testnet)
    pub network: String,
    /// Creation timestamp
    pub created_at: i64,
    /// Derived addresses (public only)
    pub addresses: Vec<WalletAddress>,
}

/// Unlocked wallet data (sensitive, only in memory)
pub struct UnlockedWallet {
    /// Decrypted seed
    seed: SeedData,
    /// Cached Ed25519 private keys (index -> key)
    ed25519_keys: Vec<[u8; 32]>,
    /// Original mnemonic (encrypted in file, decrypted here for export)
    mnemonic: Option<String>,
}

impl UnlockedWallet {
    /// Create from seed
    pub fn new(seed: SeedData, mnemonic: Option<String>) -> Self {
        Self {
            seed,
            ed25519_keys: Vec::new(),
            mnemonic,
        }
    }

    /// Get the seed
    pub fn seed(&self) -> &[u8; 64] {
        &self.seed.seed
    }

    /// Get or derive Ed25519 key for index
    pub fn get_ed25519_key(&mut self, index: u32) -> [u8; 32] {
        // Extend keys vector if needed
        while self.ed25519_keys.len() <= index as usize {
            let (private_key, _) = crypto::derive_ed25519_keypair(&self.seed.seed, self.ed25519_keys.len() as u32);
            self.ed25519_keys.push(private_key);
        }
        self.ed25519_keys[index as usize]
    }

    /// Get mnemonic if available
    pub fn mnemonic(&self) -> Option<&str> {
        self.mnemonic.as_deref()
    }
}

impl Drop for UnlockedWallet {
    fn drop(&mut self) {
        // Zero out all sensitive data
        for key in &mut self.ed25519_keys {
            key.zeroize();
        }
        if let Some(ref mut m) = self.mnemonic {
            m.zeroize();
        }
        // SeedData has ZeroizeOnDrop, so seed is auto-zeroed
    }
}

/// Main wallet state managed by Tauri
pub struct WalletState {
    /// Wallet data directory
    pub data_dir: Arc<RwLock<Option<PathBuf>>>,
    /// Network connection
    pub connection: Arc<RwLock<Option<NetworkConnection>>>,
    /// Unlocked wallet (None if locked)
    pub unlocked: Arc<RwLock<Option<UnlockedWallet>>>,
    /// Derived addresses
    pub addresses: Arc<RwLock<Vec<WalletAddress>>>,
    /// Wallet metadata (loaded from file)
    pub metadata: Arc<RwLock<Option<WalletMetadata>>>,
}

impl WalletState {
    /// Create a new wallet state
    pub fn new() -> Self {
        Self {
            data_dir: Arc::new(RwLock::new(None)),
            connection: Arc::new(RwLock::new(None)),
            unlocked: Arc::new(RwLock::new(None)),
            addresses: Arc::new(RwLock::new(Vec::new())),
            metadata: Arc::new(RwLock::new(None)),
        }
    }

    /// Set the data directory
    pub async fn set_data_dir(&self, path: PathBuf) -> Result<()> {
        // Create directory if it doesn't exist
        tokio::fs::create_dir_all(&path).await
            .map_err(|e| Error::Io(e))?;

        let mut data_dir = self.data_dir.write().await;
        *data_dir = Some(path);
        Ok(())
    }

    /// Get wallet file path
    pub async fn wallet_path(&self) -> Result<PathBuf> {
        let data_dir = self.data_dir.read().await;
        data_dir
            .as_ref()
            .map(|p| p.join(WALLET_FILE))
            .ok_or_else(|| Error::Internal("Data directory not set".to_string()))
    }

    /// Check if a wallet file exists
    pub async fn wallet_exists(&self) -> bool {
        if let Ok(path) = self.wallet_path().await {
            path.exists()
        } else {
            false
        }
    }

    /// Save wallet metadata to file
    pub async fn save_metadata(&self) -> Result<()> {
        let path = self.wallet_path().await?;
        let metadata = self.metadata.read().await;

        if let Some(meta) = metadata.as_ref() {
            let json = serde_json::to_string_pretty(meta)
                .map_err(|e| Error::Serialization(e.to_string()))?;
            tokio::fs::write(&path, json).await
                .map_err(|e| Error::Io(e))?;
        }

        Ok(())
    }

    /// Load wallet metadata from file
    pub async fn load_metadata(&self) -> Result<()> {
        let path = self.wallet_path().await?;

        if !path.exists() {
            return Err(Error::WalletNotFound);
        }

        let json = tokio::fs::read_to_string(&path).await
            .map_err(|e| Error::Io(e))?;

        let meta: WalletMetadata = serde_json::from_str(&json)
            .map_err(|e| Error::Serialization(e.to_string()))?;

        // Load addresses
        {
            let mut addresses = self.addresses.write().await;
            *addresses = meta.addresses.clone();
        }

        // Store metadata
        {
            let mut metadata = self.metadata.write().await;
            *metadata = Some(meta);
        }

        Ok(())
    }

    /// Check if wallet is unlocked
    pub async fn is_unlocked(&self) -> bool {
        self.unlocked.read().await.is_some()
    }

    /// Lock the wallet (clear sensitive data)
    pub async fn lock(&self) {
        let mut unlocked = self.unlocked.write().await;
        // Drop triggers zeroization
        *unlocked = None;
    }

    /// Unlock wallet with password
    pub async fn unlock(&self, password: &str) -> Result<()> {
        let metadata = self.metadata.read().await;
        let meta = metadata.as_ref().ok_or(Error::WalletNotFound)?;

        // Decrypt seed
        let seed = crypto::decrypt_seed(&meta.encrypted, password)?;

        // Create unlocked wallet
        let wallet = UnlockedWallet::new(seed, None);

        let mut unlocked = self.unlocked.write().await;
        *unlocked = Some(wallet);

        Ok(())
    }

    /// Create a new wallet
    pub async fn create(&self, password: &str, testnet: bool) -> Result<(String, String)> {
        // Generate mnemonic
        let mnemonic = crypto::generate_mnemonic()?;

        // Derive seed
        let seed = crypto::mnemonic_to_seed(&mnemonic, "")?;

        // Encrypt seed
        let encrypted = crypto::encrypt_seed(&seed.seed, password)?;

        // Derive first address
        let (_, pubkey) = crypto::derive_ed25519_keypair(&seed.seed, 0);
        let address = crypto::pubkey_to_address(&pubkey, testnet)?;

        let first_address = WalletAddress {
            address: address.clone(),
            index: 0,
            is_change: false,
            label: Some("Default".to_string()),
        };

        // Create metadata
        let metadata = WalletMetadata {
            encrypted,
            network: if testnet { "testnet".to_string() } else { "mainnet".to_string() },
            created_at: chrono_timestamp(),
            addresses: vec![first_address.clone()],
        };

        // Store in state
        {
            let mut meta = self.metadata.write().await;
            *meta = Some(metadata);
        }
        {
            let mut addresses = self.addresses.write().await;
            *addresses = vec![first_address];
        }

        // Save to file
        self.save_metadata().await?;

        // Unlock wallet
        let wallet = UnlockedWallet::new(seed, Some(mnemonic.clone()));
        {
            let mut unlocked = self.unlocked.write().await;
            *unlocked = Some(wallet);
        }

        Ok((mnemonic, address))
    }

    /// Import wallet from mnemonic
    pub async fn import(&self, mnemonic: &str, password: &str, testnet: bool) -> Result<String> {
        // Validate mnemonic
        crypto::validate_mnemonic(mnemonic)?;

        // Derive seed
        let seed = crypto::mnemonic_to_seed(mnemonic, "")?;

        // Encrypt seed
        let encrypted = crypto::encrypt_seed(&seed.seed, password)?;

        // Derive first address
        let (_, pubkey) = crypto::derive_ed25519_keypair(&seed.seed, 0);
        let address = crypto::pubkey_to_address(&pubkey, testnet)?;

        let first_address = WalletAddress {
            address: address.clone(),
            index: 0,
            is_change: false,
            label: Some("Default".to_string()),
        };

        // Create metadata
        let metadata = WalletMetadata {
            encrypted,
            network: if testnet { "testnet".to_string() } else { "mainnet".to_string() },
            created_at: chrono_timestamp(),
            addresses: vec![first_address.clone()],
        };

        // Store in state
        {
            let mut meta = self.metadata.write().await;
            *meta = Some(metadata);
        }
        {
            let mut addresses = self.addresses.write().await;
            *addresses = vec![first_address];
        }

        // Save to file
        self.save_metadata().await?;

        // Unlock wallet
        let wallet = UnlockedWallet::new(seed, Some(mnemonic.to_string()));
        {
            let mut unlocked = self.unlocked.write().await;
            *unlocked = Some(wallet);
        }

        Ok(address)
    }

    /// Generate a new address
    pub async fn generate_address(&self, label: Option<String>, is_change: bool) -> Result<WalletAddress> {
        // Get next index
        let addresses = self.addresses.read().await;
        let index = addresses.len() as u32;
        drop(addresses);

        // Get unlocked wallet
        let mut unlocked = self.unlocked.write().await;
        let wallet = unlocked.as_mut().ok_or(Error::WalletLocked)?;

        // Derive keypair
        let (_, pubkey) = crypto::derive_ed25519_keypair(wallet.seed(), index);

        // Determine network
        let metadata = self.metadata.read().await;
        let testnet = metadata.as_ref()
            .map(|m| m.network == "testnet")
            .unwrap_or(true);

        let address_str = crypto::pubkey_to_address(&pubkey, testnet)?;

        let address = WalletAddress {
            address: address_str,
            index,
            is_change,
            label,
        };

        // Add to state
        drop(unlocked);
        {
            let mut addresses = self.addresses.write().await;
            addresses.push(address.clone());
        }

        // Update metadata and save
        {
            let mut metadata = self.metadata.write().await;
            if let Some(meta) = metadata.as_mut() {
                meta.addresses.push(address.clone());
            }
        }
        self.save_metadata().await?;

        Ok(address)
    }
}

impl Default for WalletState {
    fn default() -> Self {
        Self::new()
    }
}

/// Get current timestamp
fn chrono_timestamp() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs() as i64)
        .unwrap_or(0)
}