rialo-cdk 0.2.0-alpha.0

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! File-based keyring provider implementation.

#[cfg(feature = "file-storage")]
use std::fs::{self, File};
#[cfg(feature = "file-storage")]
use std::io::{Read, Write};
#[cfg(feature = "file-storage")]
use std::path::{Path, PathBuf};
use std::str::FromStr;

#[cfg(feature = "file-storage")]
use async_trait::async_trait;
#[cfg(feature = "file-storage")]
use ed25519_dalek::SigningKey as Keypair;
#[cfg(feature = "file-storage")]
use serde::{Deserialize, Serialize};

#[cfg(feature = "file-storage")]
use crate::constants::{BASE_DERIVATION_PATH, DEFAULT_NUM_ACCOUNTS};
#[cfg(feature = "file-storage")]
use crate::error::{Result, RialoError};
#[cfg(all(feature = "file-storage", feature = "encryption"))]
use crate::keyring::encryption;
#[cfg(all(feature = "file-storage", feature = "mnemonic"))]
use crate::keyring::mnemonic;
#[cfg(feature = "file-storage")]
use crate::keyring::provider_base::BaseKeyringProvider;
#[cfg(feature = "file-storage")]
use crate::keyring::traits::{Keyring, KeyringProvider};
use crate::rpc::types::Pubkey;

/// `FileKeyringProvider` implements keyring storage and management using the local filesystem.
/// It stores keyring data in JSON files with encrypted private keys.
#[cfg(feature = "file-storage")]
pub struct FileKeyringProvider {
    /// Directory where keyring files are stored
    keyring_dir: PathBuf,
}

/// Represents a single keypair within a keyring file
#[cfg(feature = "file-storage")]
#[derive(Serialize, Deserialize)]
struct StoredKeypair {
    /// Base58-encoded public key
    pubkey: Pubkey,
    /// Encrypted private key (AES-GCM)
    encrypted_keypair: Vec<u8>,
    /// BIP32/44 derivation path if this is an HD keypair
    derivation_path: Option<String>,
    /// Keypair index within the keyring
    index: u32,
}

/// Represents the structure of a keyring file on disk
#[cfg(feature = "file-storage")]
#[derive(Serialize, Deserialize)]
struct KeyringFile {
    /// Name of the keyring
    name: String,
    /// Optional BIP39 mnemonic phrase (for HD keyrings)
    mnemonic: Option<String>,
    /// List of keypairs in this keyring
    keypairs: Vec<StoredKeypair>,
}

#[cfg(feature = "file-storage")]
impl FileKeyringProvider {
    /// Creates a new FileKeyringProvider that stores keyrings in the specified directory.
    ///
    /// # Arguments
    ///
    /// * `keyring_dir` - Path where keyring files will be stored
    ///
    /// # Returns
    ///
    /// A new FileKeyringProvider instance
    pub fn new(keyring_dir: impl AsRef<Path>) -> Self {
        Self {
            keyring_dir: keyring_dir.as_ref().to_path_buf(),
        }
    }

    /// Returns the default path for storing keyring files.
    ///
    /// The default location is `<os-specific-config-dir>/rialo/keyrings/`
    ///
    /// # Returns
    ///
    /// The default keyring directory path or an error
    pub fn default_path() -> Result<PathBuf> {
        let mut path = dirs::config_dir()
            .ok_or_else(|| RialoError::Keyring("Could not find config directory".to_string()))?;
        path.push("rialo");
        path.push("keyrings");
        std::fs::create_dir_all(&path)?;
        Ok(path)
    }

    /// Constructs the full path to a specific keyring file.
    fn keyring_path(&self, name: &str) -> PathBuf {
        let mut path = self.keyring_dir.clone();
        path.push(format!("{name}.keyring"));
        path
    }

    /// Reads and parses a keyring file from disk.
    fn read_keyring_file(&self, name: &str) -> Result<KeyringFile> {
        let mut file = File::open(self.keyring_path(name))
            .map_err(|e| RialoError::Keyring(format!("Failed to open keyring file: {e}")))?;

        let mut contents = String::new();
        file.read_to_string(&mut contents)
            .map_err(|e| RialoError::Keyring(format!("Failed to read keyring file: {e}")))?;

        serde_json::from_str(&contents)
            .map_err(|e| RialoError::Keyring(format!("Failed to parse keyring file: {e}")))
    }

    /// Saves a keyring file structure to disk.
    fn save_keyring_file(&self, name: &str, keyring_file: &KeyringFile) -> Result<()> {
        let json = serde_json::to_string_pretty(keyring_file)
            .map_err(|e| RialoError::Keyring(format!("Failed to serialize keyring: {e}")))?;

        let mut file = File::create(self.keyring_path(name))
            .map_err(|e| RialoError::Keyring(format!("Failed to create keyring file: {e}")))?;

        file.write_all(json.as_bytes())
            .map_err(|e| RialoError::Keyring(format!("Failed to write keyring file: {e}")))?;

        Ok(())
    }
}

#[cfg(feature = "file-storage")]
#[async_trait]
impl BaseKeyringProvider for FileKeyringProvider {}

#[cfg(feature = "file-storage")]
#[async_trait]
impl KeyringProvider for FileKeyringProvider {
    /// Creates a new keyring with a single randomly generated keypair.
    async fn create(&self, name: &str, password: &str) -> Result<Keyring> {
        if self.exists(name).await? {
            return Err(RialoError::Keyring(format!(
                "Keyring already exists: {name}"
            )));
        }

        fs::create_dir_all(&self.keyring_dir)?;

        let keypair = Keypair::generate(&mut rand::rngs::OsRng);

        #[cfg(feature = "encryption")]
        let encrypted_keypair = encryption::encrypt_keypair(&keypair, password)?;
        #[cfg(not(feature = "encryption"))]
        let encrypted_keypair = {
            let _ = password;
            keypair.as_bytes().to_vec()
        };

        let keyring = Keyring::new(name.to_string(), keypair, None, None);

        let stored = StoredKeypair {
            pubkey: Pubkey::from_str(&keyring.pubkey_string()).unwrap(),
            encrypted_keypair,
            derivation_path: None,
            index: 0,
        };

        let keyring_file = KeyringFile {
            name: name.to_string(),
            mnemonic: None,
            keypairs: vec![stored],
        };

        self.save_keyring_file(name, &keyring_file)?;

        Ok(keyring)
    }

    /// Creates a new HD keyring with a generated mnemonic phrase.
    #[cfg(feature = "mnemonic")]
    async fn create_with_mnemonic(&self, name: &str, password: &str) -> Result<Keyring> {
        if self.exists(name).await? {
            return Err(RialoError::Keyring(format!(
                "Keyring already exists: {name}"
            )));
        }

        fs::create_dir_all(&self.keyring_dir)?;

        let mnemonic_phrase = mnemonic::generate_mnemonic()?;

        let mut keypairs = Vec::new();
        let mut primary_keypair = None;

        for i in 0..DEFAULT_NUM_ACCOUNTS {
            let path = format!("{}{}'/{}'", BASE_DERIVATION_PATH, i, 0);
            let keypair = mnemonic::keypair_from_mnemonic(&mnemonic_phrase, Some(&path))?;

            #[cfg(feature = "encryption")]
            let encrypted_keypair = encryption::encrypt_keypair(&keypair, password)?;
            #[cfg(not(feature = "encryption"))]
            let encrypted_keypair = {
                let _ = password;
                keypair.as_bytes().to_vec()
            };

            let pubkey =
                Pubkey::from_str(&bs58::encode(keypair.verifying_key().as_bytes()).into_string())
                    .unwrap();

            if i == 0 {
                primary_keypair = Some(keypair);
            }

            let stored = StoredKeypair {
                pubkey,
                encrypted_keypair,
                derivation_path: Some(path),
                index: i,
            };

            keypairs.push(stored);
        }

        let primary_path = format!("{}{}'/{}'", BASE_DERIVATION_PATH, 0, 0);
        let keyring = Keyring::new(
            name.to_string(),
            primary_keypair.unwrap(),
            Some(mnemonic_phrase.clone()),
            Some(primary_path),
        );

        let keyring_file = KeyringFile {
            name: name.to_string(),
            mnemonic: Some(mnemonic_phrase),
            keypairs,
        };

        self.save_keyring_file(name, &keyring_file)?;

        Ok(keyring)
    }

    /// Recovers a keyring from an existing mnemonic phrase.
    #[cfg(feature = "mnemonic")]
    async fn recover_from_mnemonic(
        &self,
        name: &str,
        mnemonic: &str,
        password: &str,
    ) -> Result<Keyring> {
        if self.exists(name).await? {
            return Err(RialoError::Keyring(format!(
                "Keyring already exists: {name}"
            )));
        }

        fs::create_dir_all(&self.keyring_dir)?;

        let mut keypairs = Vec::new();
        let mut primary_keypair = None;

        for i in 0..DEFAULT_NUM_ACCOUNTS {
            let path = format!("{}{}'/{}'", BASE_DERIVATION_PATH, i, 0);
            let keypair = mnemonic::keypair_from_mnemonic(mnemonic, Some(&path))?;

            let pubkey =
                Pubkey::from_str(&bs58::encode(keypair.verifying_key().as_bytes()).into_string())
                    .unwrap();

            #[cfg(feature = "encryption")]
            let encrypted_keypair = encryption::encrypt_keypair(&keypair, password)?;
            #[cfg(not(feature = "encryption"))]
            let encrypted_keypair = {
                let _ = password;
                keypair.as_bytes().to_vec()
            };

            if i == 0 {
                primary_keypair = Some(keypair);
            }

            let stored = StoredKeypair {
                pubkey,
                encrypted_keypair,
                derivation_path: Some(path),
                index: i,
            };

            keypairs.push(stored);
        }

        let primary_path = format!("{}{}'/{}'", BASE_DERIVATION_PATH, 0, 0);
        let keyring = Keyring::new(
            name.to_string(),
            primary_keypair.unwrap(),
            Some(mnemonic.to_string()),
            Some(primary_path),
        );

        let keyring_file = KeyringFile {
            name: name.to_string(),
            mnemonic: Some(mnemonic.to_string()),
            keypairs,
        };

        self.save_keyring_file(name, &keyring_file)?;

        Ok(keyring)
    }

    /// Loads an existing keyring from storage.
    async fn load(&self, name: &str, password: &str) -> Result<Keyring> {
        if !self.exists(name).await? {
            return Err(RialoError::Keyring(format!("Keyring not found: {name}")));
        }

        let keyring_file = self.read_keyring_file(name)?;

        if keyring_file.keypairs.is_empty() {
            return Err(RialoError::Keyring("Keyring has no keypairs".to_string()));
        }

        let primary = &keyring_file.keypairs[0];

        #[cfg(feature = "encryption")]
        let keypair = encryption::decrypt_keypair(&primary.encrypted_keypair, password)?;
        #[cfg(not(feature = "encryption"))]
        let keypair = {
            let _ = password;
            let keypair_bytes: [u8; 32] = primary.encrypted_keypair[..32]
                .try_into()
                .map_err(|_| RialoError::Keyring("Invalid keypair size".to_string()))?;
            Keypair::from_bytes(&keypair_bytes)
        };

        Ok(Keyring::new(
            name.to_string(),
            keypair,
            keyring_file.mnemonic,
            primary.derivation_path.clone(),
        ))
    }

    /// Lists all available keyrings in the storage directory.
    async fn list(&self) -> Result<Vec<String>> {
        if !self.keyring_dir.exists() {
            fs::create_dir_all(&self.keyring_dir)?;
            return Ok(Vec::new());
        }

        let entries = fs::read_dir(&self.keyring_dir)
            .map_err(|e| RialoError::Keyring(format!("Failed to read keyring directory: {e}")))?;

        let mut names = Vec::new();
        for entry in entries {
            let entry = entry
                .map_err(|e| RialoError::Keyring(format!("Failed to read directory entry: {e}")))?;

            let path = entry.path();
            if let Some(ext) = path.extension() {
                if ext == "keyring" {
                    if let Some(stem) = path.file_stem() {
                        if let Some(name) = stem.to_str() {
                            names.push(name.to_string());
                        }
                    }
                }
            }
        }

        Ok(names)
    }

    /// Lists all keypairs in a specific keyring.
    async fn list_keypairs(&self, name: &str) -> Result<Vec<(u32, Pubkey)>> {
        if !self.exists(name).await? {
            return Err(RialoError::Keyring(format!("Keyring not found: {name}")));
        }

        let keyring_file = self.read_keyring_file(name)?;

        let keypairs: Vec<(u32, Pubkey)> = keyring_file
            .keypairs
            .iter()
            .map(|kp| (kp.index, kp.pubkey))
            .collect();

        Ok(keypairs)
    }

    /// Checks if a keyring with the given name exists.
    async fn exists(&self, name: &str) -> Result<bool> {
        Ok(self.keyring_path(name).exists())
    }

    /// Derives a new keyring from an existing HD keyring's mnemonic.
    #[cfg(feature = "hd-wallet")]
    async fn derive_keyring(
        &self,
        source_keyring_name: &str,
        new_keyring_name: &str,
        keypair_index: u32,
        password: &str,
    ) -> Result<Keyring> {
        self.validate_mnemonic_operation(source_keyring_name, new_keyring_name)
            .await?;

        let source_keyring = self.load(source_keyring_name, password).await?;

        let mnemonic = source_keyring.mnemonic().ok_or_else(|| {
            RialoError::Keyring("Source keyring does not have a mnemonic phrase".to_string())
        })?;

        let path = format!("{}{}'/{}'", BASE_DERIVATION_PATH, keypair_index, 0);
        let keypair = mnemonic::keypair_from_mnemonic(mnemonic, Some(&path))?;

        #[cfg(feature = "encryption")]
        let encrypted_keypair = encryption::encrypt_keypair(&keypair, password)?;
        #[cfg(not(feature = "encryption"))]
        let encrypted_keypair = keypair.as_bytes().to_vec();

        let keyring = Keyring::new(
            new_keyring_name.to_string(),
            keypair,
            Some(mnemonic.to_string()),
            Some(path.clone()),
        );

        let stored = StoredKeypair {
            pubkey: Pubkey::from_str(&keyring.pubkey_string()).unwrap(),
            encrypted_keypair,
            derivation_path: Some(path),
            index: keypair_index,
        };

        let keyring_file = KeyringFile {
            name: new_keyring_name.to_string(),
            mnemonic: Some(mnemonic.to_string()),
            keypairs: vec![stored],
        };

        self.save_keyring_file(new_keyring_name, &keyring_file)?;

        Ok(keyring)
    }

    /// Derives a new keypair within an existing HD keyring.
    #[cfg(feature = "hd-wallet")]
    async fn derive_keypair(
        &self,
        keyring_name: &str,
        keypair_index: u32,
        password: &str,
    ) -> Result<(u32, Pubkey)> {
        if !self.exists(keyring_name).await? {
            return Err(RialoError::Keyring(format!(
                "Keyring not found: {keyring_name}"
            )));
        }

        let mut keyring_file = self.read_keyring_file(keyring_name)?;

        if keyring_file
            .keypairs
            .iter()
            .any(|kp| kp.index == keypair_index)
        {
            return Err(RialoError::Keyring(format!(
                "Keypair with index {keypair_index} already exists"
            )));
        }

        let mnemonic = keyring_file.mnemonic.as_ref().ok_or_else(|| {
            RialoError::Keyring("Keyring does not have a mnemonic phrase".to_string())
        })?;

        let path = format!("{}{}'/{}'", BASE_DERIVATION_PATH, keypair_index, 0);
        let keypair = mnemonic::keypair_from_mnemonic(mnemonic, Some(&path))?;

        #[cfg(feature = "encryption")]
        let encrypted_keypair = encryption::encrypt_keypair(&keypair, password)?;
        #[cfg(not(feature = "encryption"))]
        let encrypted_keypair = {
            let _ = password;
            keypair.as_bytes().to_vec()
        };

        let pubkey =
            Pubkey::from_str(&bs58::encode(keypair.verifying_key().as_bytes()).into_string())
                .unwrap();

        let stored = StoredKeypair {
            pubkey,
            encrypted_keypair,
            derivation_path: Some(path),
            index: keypair_index,
        };

        keyring_file.keypairs.push(stored);
        self.save_keyring_file(keyring_name, &keyring_file)?;

        Ok((keypair_index, pubkey))
    }

    /// Gets the balance of a specific keypair.
    async fn get_keypair_balance(&self, name: &str, _keypair_index: u32) -> Result<u64> {
        if !self.exists(name).await? {
            return Err(RialoError::Keyring(format!("Keyring not found: {name}")));
        }
        Ok(0) // Placeholder - would query blockchain
    }

    /// Lists all keyring names and their primary public keys.
    async fn list_public_keys(&self) -> Result<Vec<(String, Pubkey)>> {
        let mut results = Vec::new();
        let keyrings = self.list().await?;

        for name in keyrings {
            let path = self.keyring_path(&name);
            if let Ok(mut file) = File::open(&path) {
                let mut contents = String::new();
                if file.read_to_string(&mut contents).is_ok() {
                    if let Ok(keyring_file) = serde_json::from_str::<KeyringFile>(&contents) {
                        if let Some(kp) = keyring_file.keypairs.first() {
                            results.push((name, kp.pubkey));
                        }
                    }
                }
            }
        }
        Ok(results)
    }

    /// Gets the primary public key of a keyring.
    async fn get_public_key(&self, name: &str) -> Result<Pubkey> {
        if !self.exists(name).await? {
            return Err(RialoError::Keyring(format!("Keyring not found: {name}")));
        }

        let keyring_file = self.read_keyring_file(name)?;

        if let Some(kp) = keyring_file.keypairs.first() {
            Ok(kp.pubkey)
        } else {
            Err(RialoError::Keyring("Keyring has no keypairs".to_string()))
        }
    }

    /// Gets information about all keypairs in a keyring.
    async fn get_keypairs_info(&self, name: &str) -> Result<Vec<(u32, Pubkey)>> {
        self.list_keypairs(name).await
    }

    /// Gets the public key of a specific keypair.
    async fn get_keypair_public_key(&self, name: &str, keypair_index: u32) -> Result<Pubkey> {
        if !self.exists(name).await? {
            return Err(RialoError::Keyring(format!("Keyring not found: {name}")));
        }

        let keyring_file = self.read_keyring_file(name)?;

        for kp in keyring_file.keypairs {
            if kp.index == keypair_index {
                return Ok(kp.pubkey);
            }
        }

        Err(RialoError::Keyring(format!(
            "Keypair with index {keypair_index} not found in keyring '{name}'"
        )))
    }

    /// Gets the next available keypair index for a keyring.
    async fn next_keypair_index(&self, name: &str) -> Result<u32> {
        if !self.exists(name).await? {
            return Err(RialoError::Keyring(format!("Keyring not found: {name}")));
        }

        let keyring_file = self.read_keyring_file(name)?;

        if keyring_file.keypairs.is_empty() {
            return Ok(0);
        }

        let max_index = keyring_file
            .keypairs
            .iter()
            .map(|kp| kp.index)
            .max()
            .unwrap_or(0);
        Ok(max_index + 1)
    }
}

// === Backward compatibility type alias ===

/// File-based wallet provider.
///
/// # Deprecated
///
/// Use `FileKeyringProvider` instead.
#[cfg(feature = "file-storage")]
#[deprecated(since = "0.2.0", note = "Use FileKeyringProvider instead")]
pub type FileWalletProvider = FileKeyringProvider;