NAC_Blockchain/sdk/nac-cli/src/utils/keystore.rs

use crate::error::{CliError, Result};
use crate::utils::crypto::{encrypt_private_key, decrypt_private_key, private_key_to_address};
use serde::{Deserialize, Serialize};
use std::fs;
use std::path::{Path, PathBuf};
use chrono::Utc;

/// Keystore文件格式
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KeystoreFile {
    /// 版本号
    pub version: u32,
    /// 地址
    pub address: String,
    /// 加密的私钥
    pub encrypted_key: String,
    /// 创建时间
    pub created_at: String,
    /// 备注
    pub note: Option<String>,
}

impl KeystoreFile {
    /// 创建新的keystore文件
    pub fn new(private_key: &str, password: &str, note: Option<String>) -> Result<Self> {
        let address = private_key_to_address(private_key)?;
        let encrypted_key = encrypt_private_key(private_key, password)?;
        
        Ok(Self {
            version: 1,
            address,
            encrypted_key,
            created_at: Utc::now().to_rfc3339(),
            note,
        })
    }

    /// 解密私钥
    pub fn decrypt(&self, password: &str) -> Result<String> {
        decrypt_private_key(&self.encrypted_key, password)
    }

    /// 保存到文件
    pub fn save(&self, path: &Path) -> Result<()> {
        let json = serde_json::to_string_pretty(self)
            .map_err(|e| CliError::Io(format!("序列化keystore失败: {}", e)))?;
        
        fs::write(path, json)
            .map_err(|e| CliError::Io(format!("写入keystore文件失败: {}", e)))?;
        
        Ok(())
    }

    /// 从文件加载
    pub fn load(path: &Path) -> Result<Self> {
        let json = fs::read_to_string(path)
            .map_err(|e| CliError::Io(format!("读取keystore文件失败: {}", e)))?;
        
        serde_json::from_str(&json)
            .map_err(|e| CliError::Io(format!("解析keystore文件失败: {}", e)))
    }
}

/// Keystore管理器
pub struct KeystoreManager {
    keystore_dir: PathBuf,
}

impl KeystoreManager {
    /// 创建新的keystore管理器
    pub fn new(keystore_dir: PathBuf) -> Result<Self> {
        // 确保keystore目录存在
        if !keystore_dir.exists() {
            fs::create_dir_all(&keystore_dir)
                .map_err(|e| CliError::Io(format!("创建keystore目录失败: {}", e)))?;
        }
        
        Ok(Self { keystore_dir })
    }

    /// 获取默认keystore目录
    pub fn default_dir() -> Result<PathBuf> {
        let home = dirs::home_dir()
            .ok_or_else(|| CliError::Io("无法获取用户主目录".to_string()))?;
        Ok(home.join(".nac").join("keystore"))
    }

    /// 创建默认keystore管理器
    pub fn default() -> Result<Self> {
        Self::new(Self::default_dir()?)
    }

    /// 生成keystore文件名
    fn generate_filename(&self, address: &str) -> String {
        let timestamp = Utc::now().format("%Y%m%d_%H%M%S");
        let addr_short = &address[2..10]; // 取地址前8个字符
        format!("UTC--{}--{}.json", timestamp, addr_short)
    }

    /// 导入私钥
    pub fn import(&self, private_key: &str, password: &str, note: Option<String>) -> Result<String> {
        let keystore = KeystoreFile::new(private_key, password, note)?;
        let filename = self.generate_filename(&keystore.address);
        let path = self.keystore_dir.join(&filename);
        
        keystore.save(&path)?;
        
        Ok(keystore.address)
    }

    /// 导出私钥
    pub fn export(&self, address: &str, password: &str) -> Result<String> {
        let keystore = self.find_by_address(address)?;
        keystore.decrypt(password)
    }

    /// 列出所有账户
    pub fn list(&self) -> Result<Vec<KeystoreFile>> {
        let mut keystores = Vec::new();
        
        let entries = fs::read_dir(&self.keystore_dir)
            .map_err(|e| CliError::Io(format!("读取keystore目录失败: {}", e)))?;
        
        for entry in entries {
            let entry = entry.map_err(|e| CliError::Io(format!("读取目录项失败: {}", e)))?;
            let path = entry.path();
            
            if path.extension().and_then(|s| s.to_str()) == Some("json") {
                if let Ok(keystore) = KeystoreFile::load(&path) {
                    keystores.push(keystore);
                }
            }
        }
        
        // 按创建时间排序
        keystores.sort_by(|a, b| b.created_at.cmp(&a.created_at));
        
        Ok(keystores)
    }

    /// 根据地址查找keystore
    pub fn find_by_address(&self, address: &str) -> Result<KeystoreFile> {
        let keystores = self.list()?;
        
        keystores
            .into_iter()
            .find(|k| k.address.eq_ignore_ascii_case(address))
            .ok_or_else(|| CliError::Io(format!("未找到地址 {} 的keystore", address)))
    }

    /// 删除账户
    pub fn delete(&self, address: &str) -> Result<()> {
        let entries = fs::read_dir(&self.keystore_dir)
            .map_err(|e| CliError::Io(format!("读取keystore目录失败: {}", e)))?;
        
        for entry in entries {
            let entry = entry.map_err(|e| CliError::Io(format!("读取目录项失败: {}", e)))?;
            let path = entry.path();
            
            if path.extension().and_then(|s| s.to_str()) == Some("json") {
                if let Ok(keystore) = KeystoreFile::load(&path) {
                    if keystore.address.eq_ignore_ascii_case(address) {
                        fs::remove_file(&path)
                            .map_err(|e| CliError::Io(format!("删除keystore文件失败: {}", e)))?;
                        return Ok(());
                    }
                }
            }
        }
        
        Err(CliError::Io(format!("未找到地址 {} 的keystore", address)))
    }

    /// 更新备注
    pub fn update_note(&self, address: &str, note: String) -> Result<()> {
        let entries = fs::read_dir(&self.keystore_dir)
            .map_err(|e| CliError::Io(format!("读取keystore目录失败: {}", e)))?;
        
        for entry in entries {
            let entry = entry.map_err(|e| CliError::Io(format!("读取目录项失败: {}", e)))?;
            let path = entry.path();
            
            if path.extension().and_then(|s| s.to_str()) == Some("json") {
                if let Ok(mut keystore) = KeystoreFile::load(&path) {
                    if keystore.address.eq_ignore_ascii_case(address) {
                        keystore.note = Some(note);
                        keystore.save(&path)?;
                        return Ok(());
                    }
                }
            }
        }
        
        Err(CliError::Io(format!("未找到地址 {} 的keystore", address)))
    }

    /// 修改密码
    pub fn change_password(&self, address: &str, old_password: &str, new_password: &str) -> Result<()> {
        let entries = fs::read_dir(&self.keystore_dir)
            .map_err(|e| CliError::Io(format!("读取keystore目录失败: {}", e)))?;
        
        for entry in entries {
            let entry = entry.map_err(|e| CliError::Io(format!("读取目录项失败: {}", e)))?;
            let path = entry.path();
            
            if path.extension().and_then(|s| s.to_str()) == Some("json") {
                if let Ok(keystore) = KeystoreFile::load(&path) {
                    if keystore.address.eq_ignore_ascii_case(address) {
                        // 用旧密码解密
                        let private_key = keystore.decrypt(old_password)?;
                        
                        // 用新密码加密
                        let new_keystore = KeystoreFile::new(&private_key, new_password, keystore.note)?;
                        new_keystore.save(&path)?;
                        
                        return Ok(());
                    }
                }
            }
        }
        
        Err(CliError::Io(format!("未找到地址 {} 的keystore", address)))
    }

    /// 验证密码
    pub fn verify_password(&self, address: &str, password: &str) -> Result<bool> {
        let keystore = self.find_by_address(address)?;
        match keystore.decrypt(password) {
            Ok(_) => Ok(true),
            Err(_) => Ok(false),
        }
    }
}

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

    #[test]
    fn test_keystore_file() {
        let private_key = "1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef";
        let password = "test_password";
        
        let keystore = KeystoreFile::new(private_key, password, Some("test account".to_string())).expect("mainnet: handle error");
        
        assert_eq!(keystore.version, 1);
        assert!(keystore.address.starts_with("0x"));
        assert!(keystore.note.is_some());
        
        let decrypted = keystore.decrypt(password).expect("mainnet: handle error");
        assert_eq!(decrypted, private_key);
    }

    #[test]
    fn test_keystore_manager() {
        let temp_dir = TempDir::new().expect("mainnet: handle error");
        let manager = KeystoreManager::new(temp_dir.path().to_path_buf()).expect("mainnet: handle error");
        
        let private_key = "1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef";
        let password = "test_password";
        
        // 导入
        let address = manager.import(private_key, password, Some("test".to_string())).expect("mainnet: handle error");
        
        // 列出
        let list = manager.list().expect("mainnet: handle error");
        assert_eq!(list.len(), 1);
        assert_eq!(list[0].address, address);
        
        // 导出
        let exported = manager.export(&address, password).expect("mainnet: handle error");
        assert_eq!(exported, private_key);
        
        // 删除
        manager.delete(&address).expect("mainnet: handle error");
        let list = manager.list().expect("mainnet: handle error");
        assert_eq!(list.len(), 0);
    }

    #[test]
    fn test_change_password() {
        let temp_dir = TempDir::new().expect("mainnet: handle error");
        let manager = KeystoreManager::new(temp_dir.path().to_path_buf()).expect("mainnet: handle error");
        
        let private_key = "1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef";
        let old_password = "old_password";
        let new_password = "new_password";
        
        let address = manager.import(private_key, old_password, None).expect("mainnet: handle error");
        
        // 修改密码
        manager.change_password(&address, old_password, new_password).expect("mainnet: handle error");
        
        // 用新密码导出
        let exported = manager.export(&address, new_password).expect("mainnet: handle error");
        assert_eq!(exported, private_key);
        
        // 用旧密码应该失败
        assert!(manager.export(&address, old_password).is_err());
    }
}