etopay_sdk/wallet/
wallet_manager.rs

1//! This module contains the implementation of the StardustWallet trait and the WalletImpl struct.
2//! The StardustWallet trait defines the methods for creating, migrating, backing up, restoring, and deleting a wallet.
3//! The WalletImpl struct represents an instantiated wallet and holds the necessary state and configuration.
4//!
5
6use super::share::Share;
7use super::wallet::WalletUser;
8use super::wallet_evm::{WalletImplEvm, WalletImplEvmErc20};
9use super::wallet_stardust::WalletImplStardust;
10use crate::core::{Config, UserRepoT};
11use crate::types::newtypes::{AccessToken, EncryptionPin, EncryptionSalt, PlainPassword};
12use crate::wallet::error::{ErrorKind, Result, WalletError};
13use api_types::api::networks::{ApiNetwork, ApiProtocol};
14use async_trait::async_trait;
15use iota_sdk::crypto::keys::bip39::Mnemonic;
16use log::{info, warn};
17use secrecy::{ExposeSecret, SecretBox};
18use std::marker::PhantomData;
19use std::ops::{Deref, DerefMut};
20
21/// Represents borrowing a [`WalletUser`] instance with a lifetime connected to the wallet manager.
22/// This prevents wallets to be stored and used later by another user.
23pub struct WalletBorrow<'a> {
24    inner: Box<dyn WalletUser + Send + Sync>,
25    /// with this we "attach" a lifetime to this object even though it is not "needed"
26    _lifetime: std::marker::PhantomData<&'a ()>,
27}
28
29#[cfg(test)]
30impl WalletBorrow<'_> {
31    /// test function to create [`WalletBorrow`] instances in mock objects
32    pub fn from(inner: impl WalletUser + Send + Sync + 'static) -> Self {
33        Self {
34            inner: Box::new(inner),
35            _lifetime: PhantomData,
36        }
37    }
38}
39
40impl Deref for WalletBorrow<'_> {
41    type Target = Box<dyn WalletUser + Send + Sync>;
42
43    fn deref(&self) -> &Self::Target {
44        &self.inner
45    }
46}
47impl DerefMut for WalletBorrow<'_> {
48    fn deref_mut(&mut self) -> &mut Self::Target {
49        &mut self.inner
50    }
51}
52
53/// Creates a wallet and returns an instance to work upon
54#[cfg_attr(test, mockall::automock)]
55#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
56#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
57pub trait WalletManager: std::fmt::Debug {
58    /// Get the recovery share
59    fn get_recovery_share(&self) -> Option<Share>;
60
61    /// Set the recovery share
62    fn set_recovery_share(&mut self, share: Option<Share>);
63
64    /// Generate a new mnemonic and create shares. Returns the new mnemonic.
65    async fn create_wallet_from_new_mnemonic(
66        &mut self,
67        config: &Config,
68        access_token: &Option<AccessToken>,
69        repo: &mut UserRepoT,
70        pin: &EncryptionPin,
71    ) -> Result<String>;
72
73    /// Create shares from a mnemonic
74    async fn create_wallet_from_existing_mnemonic(
75        &mut self,
76        config: &Config,
77        access_token: &Option<AccessToken>,
78        repo: &mut UserRepoT,
79        pin: &EncryptionPin,
80        mnemonic: &str,
81    ) -> Result<()>;
82
83    /// Create shares from a kdbx backup byte stream
84    async fn create_wallet_from_backup(
85        &mut self,
86        config: &Config,
87        access_token: &Option<AccessToken>,
88        repo: &mut UserRepoT,
89        pin: &EncryptionPin,
90        backup: &[u8],
91        backup_password: &PlainPassword,
92    ) -> Result<()>;
93
94    /// Create kdbx backup bytes from shares
95    async fn create_wallet_backup(
96        &mut self,
97        config: &Config,
98        access_token: &Option<AccessToken>,
99        repo: &mut UserRepoT,
100        pin: &EncryptionPin,
101        backup_password: &PlainPassword,
102    ) -> Result<Vec<u8>>;
103
104    /// deletes the user's wallet
105    async fn delete_wallet(
106        &mut self,
107        config: &Config,
108        access_token: &Option<AccessToken>,
109        repo: &mut UserRepoT,
110    ) -> Result<()>;
111
112    /// Checks if the mnemonic resembled by the shares is the same as the provided mnemonic.
113    async fn check_mnemonic(
114        &mut self,
115        config: &Config,
116        access_token: &Option<AccessToken>,
117        repo: &mut UserRepoT,
118        pin: &EncryptionPin,
119        mnemonic: &str,
120    ) -> Result<bool>;
121
122    /// Changes the password of the existing wallet by re-encrypting the backup share
123    async fn change_wallet_password(
124        &mut self,
125        config: &Config,
126        access_token: &Option<AccessToken>,
127        repo: &mut UserRepoT,
128        pin: &EncryptionPin,
129        new_password: &PlainPassword,
130    ) -> Result<()>;
131
132    /// Tries to instantiate a [`WalletUser`] object from shares and/or returns a mutable reference bound to
133    /// the lifetime of this object. The same instance may be reused across several calls to
134    /// `try_get`, hence the lifetime is bound to the lifetime of `self`.
135    async fn try_get<'a>(
136        &'a mut self,
137        config: &mut Config,
138        access_token: &Option<AccessToken>,
139        repo: &mut UserRepoT,
140        network: &ApiNetwork,
141        pin: &EncryptionPin,
142    ) -> Result<WalletBorrow<'a>>;
143}
144
145/// Implementation of [`WalletManager`] that uses the SSS schema to store and retrieve the mnemonic
146/// and create implementations of [`WalletUser`].
147#[derive(Debug)]
148pub struct WalletManagerImpl {
149    /// the name of the user we are creating wallets for
150    username: String,
151
152    /// The recovery share that the user should download
153    pub recovery_share: Option<Share>,
154}
155
156#[derive(Debug, PartialEq)]
157struct Status {
158    /// if the local share was used.
159    local: bool,
160    /// how the recovery share was used, or [`None`] if it was not used.
161    recovery: Option<RecoveryUsed>,
162    /// if the remote backup share was used.
163    backup: bool,
164}
165
166/// Which recovery share that was used.
167#[derive(Debug, PartialEq)]
168enum RecoveryUsed {
169    /// Recovery share stored locally was used
170    Local,
171    /// Recovery share stored remotely in the backend was used
172    Remote,
173}
174
175impl WalletManagerImpl {
176    /// Create a new [`WalletManagerImpl`] from a username.
177    pub fn new(username: impl Into<String>) -> Self {
178        Self {
179            username: username.into(),
180            recovery_share: None,
181        }
182    }
183
184    // fn for getting the mnemonic
185    async fn try_resemble_shares(
186        &mut self,
187        config: &Config,
188        access_token: &Option<AccessToken>,
189        repo: &mut UserRepoT,
190        pin: &EncryptionPin,
191    ) -> Result<(Mnemonic, Status)> {
192        info!("Initializing wallet for user from shares");
193
194        let username = &self.username;
195        let local_recovery_share = self.recovery_share.clone();
196
197        let user = repo.get(username)?;
198
199        // make sure the provided pin is valid (even though we currenctly do not need it unless
200        // decrypting some of the shares, but in the future we might want to encrypt the local
201        // share too, for example)
202        // let _ = user
203        //     .encrypted_password
204        //     .as_ref()
205        //     .ok_or(crate::Error::WalletNotInitialized(crate::WalletNotInitializedKind::MissingPassword))?
206        //     .decrypt(pin, &user.salt)?;
207
208        // check the availability of each share (in priority order of ease-of-use and availability)
209
210        let mut available_shares: Vec<Share> = Vec::new();
211        let mut recovery_share_available_with_user_action = false;
212        let mut password_required = false;
213
214        // in case of success we need to keep track of the share states
215        let mut local_used = false;
216        let mut recovery_used = None;
217        let mut backup_used = false;
218
219        if let Some(share) = user.local_share.map(|s| s.parse::<Share>()) {
220            available_shares.push(share?);
221            log::info!("Local storage share available");
222            local_used = true;
223        }
224
225        if let Some(share) = local_recovery_share {
226            available_shares.push(share);
227            log::info!("Local recovery share available");
228            recovery_used = Some(RecoveryUsed::Local);
229        } else {
230            log::info!("Local recovery share not available, checking if it can be downloaded");
231            recovery_share_available_with_user_action = true;
232
233            // try getting the oauth share (not encrypted)
234            if let Some(access_token) = &access_token {
235                match crate::backend::shares::download_recovery_share(config, access_token, username).await {
236                    Ok(Some(share)) => {
237                        available_shares.push(share);
238                        recovery_share_available_with_user_action = false; // this share is now available
239                        recovery_used = Some(RecoveryUsed::Remote);
240                        log::info!("Recovery share downloaded and available");
241                    }
242                    Ok(None) => log::info!("Recovery share not available"),
243                    Err(e) => log::warn!("Error fetching recovery share: {e}"),
244                }
245            } else {
246                log::info!("Access token not available, skipping recovery share download.");
247            }
248        }
249
250        // if we have less than two, we should try to get the backup share
251        // this is the last resort since it requires the password
252        if available_shares.len() < 2 {
253            if let Some(access_token) = &access_token {
254                // try to get it from the backend
255                match crate::backend::shares::download_backup_share(config, access_token, username).await {
256                    Ok(Some(share)) => {
257                        available_shares.push(share);
258                        password_required = true;
259                        backup_used = true;
260                        log::info!("Backup share (encrypted) downloaded and available");
261                    }
262                    Ok(None) => log::info!("Backup share (encrypted) not available"),
263                    Err(e) => log::warn!("Error fetching backup share: {e}"),
264                }
265            } else {
266                log::info!("Access token not available, skipping backup share download.");
267            }
268        }
269
270        // done, no need to leave the variables mutable anymore
271        let available_shares = available_shares;
272        let recovery_share_available_with_upload = recovery_share_available_with_user_action;
273        let password_required = password_required;
274
275        log::info!(
276            "Done collecting shares. Got {} shares, recovery_share_available_with_user_action = {}, password_required = {}",
277            available_shares.len(),
278            recovery_share_available_with_upload,
279            password_required
280        );
281
282        if available_shares.len() >= 2 {
283            // enough shares are available!
284
285            // if the password is required, we need to try to get it or return an error
286            let password = if password_required {
287                let password = user
288                    .encrypted_password
289                    .ok_or(WalletError::WalletNotInitialized(ErrorKind::MissingPassword))?
290                    .decrypt(pin, &user.salt)?;
291
292                Some(password)
293            } else {
294                None
295            };
296
297            let shares_ref = available_shares.iter().collect::<Vec<&Share>>();
298
299            // now we can finally try to recreate the mnemonic from the shares
300            let mnemonic = crate::share::reconstruct_mnemonic(
301                &shares_ref,
302                password.as_ref().map(PlainPassword::into_secret).as_ref(),
303            )?;
304
305            if !local_used {
306                log::info!("Local share not set, recreating shares and storing local share again");
307
308                // create the shares again, and just use a random password since we are not interested
309                // in the backup share anyways (which is the only reason this needs a password)
310                let shares = crate::share::create_shares_from_mnemonic(
311                    secrecy::ExposeSecret::expose_secret(&mnemonic).clone(),
312                    &SecretBox::new(String::from("dummy password").as_bytes().into()),
313                )?;
314
315                // ignore the error since we were still able to create a valid wallet
316                if let Err(e) = repo.set_local_share(username, Some(&shares.local)) {
317                    log::warn!("Error storing local share again: {e:#}");
318                } else {
319                    log::info!("Done storing local share again");
320                }
321            }
322
323            Ok((
324                mnemonic.expose_secret().clone(),
325                Status {
326                    local: local_used,
327                    recovery: recovery_used,
328                    backup: backup_used,
329                },
330            ))
331        } else if available_shares.len() == 1 && recovery_share_available_with_upload {
332            Err(WalletError::WalletNotInitialized(ErrorKind::SetRecoveryShare))
333        } else {
334            // there is no way to recover the shares
335            Err(WalletError::WalletNotInitialized(ErrorKind::UseMnemonic))
336        }
337    }
338
339    /// Creates shares from the provided mnemonic and stores the local share locally, uploads the other
340    /// shares to the backend and returns the recovery share for the user to download and save.
341    async fn create_and_upload_shares(
342        &mut self,
343        config: &Config,
344        access_token: &Option<AccessToken>,
345        repo: &mut UserRepoT,
346        pin: &EncryptionPin,
347        mnemonic: impl Into<Mnemonic>,
348    ) -> Result<()> {
349        log::info!("Creating and uploading shares");
350
351        // get the password from the repo
352        let user = repo.get(&self.username)?;
353        let Some(encrypted_password) = user.encrypted_password else {
354            return Err(WalletError::WalletNotInitialized(ErrorKind::MissingPassword));
355        };
356
357        let password = encrypted_password.decrypt(pin, &user.salt)?;
358        let shares = crate::share::create_shares_from_mnemonic(mnemonic, &password.into_secret())?;
359
360        log::info!("Shares created, storing local share");
361        repo.set_local_share(&user.username, Some(&shares.local))?;
362        self.recovery_share = Some(shares.recovery.clone());
363
364        if let Some(access_token) = access_token {
365            log::info!("Uploading shares");
366            crate::backend::shares::upload_shares(config, access_token, &shares.backup, &shares.recovery).await?;
367            log::info!("Done uploading shares");
368        } else {
369            log::info!("No access token, skipping uploading backup and recovery shares");
370        }
371        Ok(())
372    }
373}
374
375#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
376#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
377impl WalletManager for WalletManagerImpl {
378    fn get_recovery_share(&self) -> Option<Share> {
379        self.recovery_share.clone()
380    }
381    fn set_recovery_share(&mut self, share: Option<Share>) {
382        self.recovery_share = share;
383    }
384    /// Generate a new mnemonic and create shares. Returns the new mnemonic.
385    async fn create_wallet_from_new_mnemonic(
386        &mut self,
387        config: &Config,
388        access_token: &Option<AccessToken>,
389        repo: &mut UserRepoT,
390        pin: &EncryptionPin,
391    ) -> Result<String> {
392        let mnemonic = iota_sdk::client::Client::generate_mnemonic()?;
393        self.create_and_upload_shares(config, access_token, repo, pin, mnemonic.as_ref())
394            .await?;
395
396        Ok(mnemonic.to_string())
397    }
398
399    /// Create shares from a mnemonic
400    async fn create_wallet_from_existing_mnemonic(
401        &mut self,
402        config: &Config,
403        access_token: &Option<AccessToken>,
404        repo: &mut UserRepoT,
405        pin: &EncryptionPin,
406        mnemonic: &str,
407    ) -> Result<()> {
408        self.create_and_upload_shares(config, access_token, repo, pin, mnemonic)
409            .await
410    }
411
412    /// Create shares from a kdbx backup byte stream
413    async fn create_wallet_from_backup(
414        &mut self,
415        config: &Config,
416        access_token: &Option<AccessToken>,
417        repo: &mut UserRepoT,
418        pin: &EncryptionPin,
419        backup: &[u8],
420        backup_password: &PlainPassword,
421    ) -> Result<()> {
422        use secrecy::ExposeSecret;
423        let mnemonic = crate::kdbx::load_mnemonic(backup, &backup_password.into_secret_string())?;
424        self.create_and_upload_shares(config, access_token, repo, pin, mnemonic.expose_secret().clone())
425            .await
426    }
427
428    /// Create kdbx backup bytes from shares
429    async fn create_wallet_backup(
430        &mut self,
431        config: &Config,
432        access_token: &Option<AccessToken>,
433        repo: &mut UserRepoT,
434        pin: &EncryptionPin,
435        backup_password: &PlainPassword,
436    ) -> Result<Vec<u8>> {
437        let (mnemonic, _status) = self.try_resemble_shares(config, access_token, repo, pin).await?;
438
439        Ok(crate::kdbx::store_mnemonic(
440            &SecretBox::new(Box::new(mnemonic)),
441            &backup_password.into_secret_string(),
442        )?)
443    }
444
445    async fn delete_wallet(
446        &mut self,
447        config: &Config,
448        access_token: &Option<AccessToken>,
449        repo: &mut UserRepoT,
450    ) -> Result<()> {
451        // remove the wallet folder
452        #[cfg(not(target_arch = "wasm32"))]
453        {
454            let path = config.path_prefix.join("wallets").join(&self.username);
455            if let Err(e) = std::fs::remove_dir_all(path) {
456                warn!("Error removing wallet files: {e:?}");
457            }
458        }
459
460        // clear the local and recovery share
461        repo.set_local_share(&self.username, None)?;
462        self.recovery_share = None;
463
464        // call backend if access_token exists
465        if let Some(access_token) = access_token {
466            crate::backend::shares::delete_shares(config, access_token, &self.username).await?;
467        }
468
469        Ok(())
470    }
471
472    async fn check_mnemonic(
473        &mut self,
474        config: &Config,
475        access_token: &Option<AccessToken>,
476        repo: &mut UserRepoT,
477        pin: &EncryptionPin,
478        mnemonic: &str,
479    ) -> Result<bool> {
480        // first use the existing pin and stored (encrypted) password to resemble the shares into
481        // the mnemonic
482        let (existing_mnemonic, _status) = self.try_resemble_shares(config, access_token, repo, pin).await?;
483
484        // perform a str-str comparison
485        Ok(*mnemonic == **existing_mnemonic)
486    }
487
488    async fn change_wallet_password(
489        &mut self,
490        config: &Config,
491        access_token: &Option<AccessToken>,
492        repo: &mut UserRepoT,
493        pin: &EncryptionPin,
494        new_password: &PlainPassword,
495    ) -> Result<()> {
496        // first use the existing pin and stored (encrypted) password to try resemble the shares into the mnemonic
497        let result = self.try_resemble_shares(config, access_token, repo, pin).await;
498
499        // if there was a real error (not only a missing wallet), propagate it,
500        // otherwise we want to goahead and update the repo
501        match result {
502            Ok(_) => {}
503            Err(WalletError::WalletNotInitialized(ErrorKind::UseMnemonic)) => {
504                warn!("No wallet found (UseMnemonic error), continuing to change the password locally only")
505            }
506            Err(e) => return Err(e),
507        }
508
509        // now update the password in the repo (perhaps a bit hacky... xD)
510        let mut user = repo.get(&self.username)?;
511        let salt = EncryptionSalt::generate();
512        let encrypted_password = new_password.encrypt(pin, &salt)?;
513        user.salt = salt;
514        user.encrypted_password = Some(encrypted_password);
515        repo.update(&user)?;
516
517        // and if we need to reconstruct the shares, do it!
518        if let Ok((mnemonic, _status)) = result {
519            self.create_and_upload_shares(config, access_token, repo, pin, mnemonic.clone())
520                .await?;
521        }
522
523        Ok(())
524    }
525
526    async fn try_get<'a>(
527        &'a mut self,
528        config: &mut Config,
529        access_token: &Option<AccessToken>,
530        repo: &mut UserRepoT,
531        network: &ApiNetwork,
532        pin: &EncryptionPin,
533    ) -> Result<WalletBorrow<'a>> {
534        let (mnemonic, _status) = self.try_resemble_shares(config, access_token, repo, pin).await?;
535
536        // we have the mnemonic and can now instantiate the WalletImpl
537
538        let path = config
539            .path_prefix
540            .join("wallets")
541            .join(&self.username)
542            .join(network.clone().key);
543
544        let bo = match &network.protocol {
545            ApiProtocol::Evm { chain_id } => {
546                let wallet = WalletImplEvm::new(
547                    mnemonic,
548                    &network.node_urls,
549                    *chain_id,
550                    network.decimals,
551                    network.coin_type,
552                )?;
553                Box::new(wallet) as Box<dyn WalletUser + Sync + Send>
554            }
555            ApiProtocol::EvmERC20 {
556                chain_id,
557                contract_address,
558            } => {
559                let wallet = WalletImplEvmErc20::new(
560                    mnemonic,
561                    &network.node_urls,
562                    *chain_id,
563                    network.decimals,
564                    network.coin_type,
565                    contract_address,
566                )?;
567                Box::new(wallet) as Box<dyn WalletUser + Sync + Send>
568            }
569            ApiProtocol::Stardust {} => {
570                let wallet = WalletImplStardust::new(mnemonic, &path, network.coin_type, &network.node_urls).await?;
571                Box::new(wallet) as Box<dyn WalletUser + Sync + Send>
572            }
573        };
574
575        Ok(WalletBorrow {
576            inner: bo,
577            _lifetime: PhantomData,
578        })
579    }
580}
581
582#[cfg(test)]
583mod tests {
584    use super::*;
585    use crate::{
586        core::{Config, UserRepoT},
587        kdbx::KdbxStorageError,
588        testing_utils::{ENCRYPTED_WALLET_PASSWORD, IOTA_NETWORK_KEY, PIN, SALT, WALLET_PASSWORD, example_api_network},
589        types::{
590            newtypes::{AccessToken, EncryptionPin, EncryptionSalt, PlainPassword},
591            users::KycType,
592        },
593        user::{MockUserRepo, memory_storage::MemoryUserStorage, repository::UserRepoImpl},
594    };
595    use kdbx_rs::errors::UnlockError;
596    use rstest::rstest;
597    use std::sync::LazyLock;
598
599    const MNEMONIC: &str = "endorse answer radar about source reunion marriage tag sausage weekend frost daring base attack because joke dream slender leisure group reason prepare broken river";
600    const MNEMONIC_INCORRECT: &str = "answer radar about source reunion marriage tag sausage weekend frost daring base attack because joke dream slender leisure group reason prepare broken river";
601    const USERNAME: &str = "SuperAdmin";
602    static INVALID_BACKUP_PASSWORD: LazyLock<PlainPassword> =
603        LazyLock::new(|| PlainPassword::try_from_string("correcthorsebatterystapleaa").unwrap());
604
605    // share strings to use for testing the resemble-function
606    const SHARE_LOCAL: &str = "ME-RS-N-Mi0xLUNBRVFBaGdESXFBRStPVUZYZTJnMTdLRFY1L2pWRllQTHdtZ0dCWExJbitjTERReFRyRHArWGNVMG5yY3UyVmFONFEvZkVoeXNadm5qNFhmRDVIZXZ3eHB2bENTYnZIZTFtOTlXdjJwby8zVWl0d2VhMnVWOTZaejB5WmhEdHlkRDFYcEg1R0RIYXFvZDBpTHdpcDZ3d1k5T0VWdEJhZmtkUVRGaTNNM3gvY2dsK0FDWVQ5WG50TlJycnRtWFRTUGZ4MG54R1lVc0NWUnNKY3h5Q0JxSHBlRGVRekpSTlFxVldMNGpJU3JCZkFRcEpYMnJoT1o4OXM1V3VLaW5PWFd0YUZncTRnd2t1VzR0ZkJJZzVUMjFlaXpGNEpWNzlMcXFXSDZoY3N0Z1huYzZYWTJvZjRvaytlYnJWOFBmR1lOU1NxRWQ4VFpqUzlBL0h0clJGNThEbUdaL2Z2Nmp5MjJjS01hUWllK1ZqdFZ4OUJyblJjWThYYTgxWmNTWlF4YlFLbFQ3MC9tRk5aQlN4ZXNLTWVTU24vV2hycEs0OU80ZW4zRkZJVTJqd2lLcGwybHpHMk0vdThJTzRZSlNCL1B6aVp4cGczcVk5Z25PRHNQR2lDZGNyejErcTVhYUdoMDdXUGlISFg5K1VpbVJjRThZS1BBNXUwNTBkQ2l2eVM2a2VhZkpFalQ0UXkxcElPaFRUd3ZrMWxrR0ZmeWp3bVBqL3JMRGY4YUc3ZXZlVWQveGwxbzlKMnh5ckhvQW9heTNVNVpHYjFCZGJ2OGFGNHJLb2wwTkorUlZBSTZJSHJCUnE4OGxJeGtzSlFxTm9GQ3o5b051N011OTVkMUJpZ3ErNjZiYzBuTWcyWXZYQXdaMkh3RjAzS0xRWEFWYjZVekZ1Lzc0MjYraElNUlR1M01mZDZoa01vMllMVzlxSS9odlBsaWg4RG5qaUFTUG9Fbkx2cVFidVpXaVBnQ3h2c1F4eXFBQWdDazlzckhwaG51UnpTck95M1JmZzRYa0lndHhlb3ArUmZJOVgyaDBRcEVmcjgzYzExd0xhQkxDUmgwMlFXazA2Ty8yM2s2cWZNZHBxNVZ2b0ZnTkNJYlY1V01sSFpaV3RnVXFzaGtXRVJycjduZnVvd1BQQ0NUaHdxMC9tbUQ1NDVDb0VNWU16bUtQYlIyYmF4RkVTbUswTlRRT3VWR3A2Y3JqNWlYOGxzaU9kZ3FVNHhuSVpRcDRsT1lJcTlBOUhFS3NZZ1RuYysxRlRNazJEN05ydThlalh4UUR3amFqUTFNTmJ5cldBS0MvZ3RTWW9ONTFKY25FWFlUOWI1MVZOWWF4anArTE9oeDA0M3RNUW9TejNvN1kxbWtORlJUTmJMZWhDKzV0UkNKNjdQYk5Va29DbWxXbjFYODZxVlVsRFU0MjkxaXVLaE1YQ0lsVHlscGU2dw==";
607    const SHARE_BACKUP: &str = "ME-RS-AesGcm-K3vx+e6IF6BOUJ2DemvsdflQq2CbolcFqdazfapauZTHdY/Hovh5zC8s5Qmfb2tRmRaluRX1gxMZfGDP52rakFnZpOzOCNGyHiI/dsiFDFbty0fEheEw+p1LrOI4zNwy7NE7ZsK0C756ggVfrhCin2Yw0KA6pALFqfWnQokx5Q43pUFd6ZGD8fwathC4NGx/hVTi9lxA2L6ScNQY9V3bEie40MKdpLQ6ELsPq+38UVJtqIgE0wJs8fDKSIGJVEPvP6wbVa+oPB/uFl5h56YeuYB2UGHdMJ54DCEoUBSd5QGoeKwjIylrZ+wXzchPXhtAfaCmqlf0fmKi9f5FQGrFwH9drf5HFE5Z/JWQC1FMKJTeBZ2CgcFvCtHuVm8VNnhhes1fUc7gL8VNOqE25LHFFQp3fBfeHXRkCmX+PAU+1N8KU6SFX0XqDr5anKAMH6thViBdno2m6K9tzqyucUnfgHYgp/cc+XXo9Ffw7v6lVTW3ls9diZwdwcs9JYqoKhWAs9dVGPz0017glpeAz01moJDPSMkhZwQh9GGWvhyeTWE9T28NS1G3cOBkW0GbgmIDjKeDDXAOjDyN7Db0FFL3TRAXthFtRXjJyZD1Xu2quYyjz1ZG70ILp0rDzzDaikUPUt1TCsAz+8NfLwHKz+H4oPUGprdUqgBVSGOySH+lKZaUbN17qIXjEKg58jh686s6i4GTD7Ndf6Xqsdc00PRDlm+jHwK7bNvkqkcChQHockIaIi4ETHCz/jqrca7uY8RIABv9Ni46+Ix1CrNY4qCUhep9oYZBGSLy2fQWWNk2nZgbrkipwUbgoV1IJV/kWCQ6ycjGG005kv3AFb6sZyrnFbvT7sa/JCKlo8gcVtzXlrJJqiO/7Qb1nTfj9dLd+/4ihpmwpFwPmKHi6zrZjJ8FbaDGkXSg+a82RQqz/AsH10hBd/tSZeZ5chdwgxTouoGix99HZipTKXLiAqW7Mo0N93+atNb9EWeHPBfsVbwJ2shBT5030QrY2qQfhTb4GUl52vPQBvpjxCzjPlvCWzFMlO8wrCP1sJm5egEb0F6Fpa9H3blBdMcb2NuKJ2VfSQzuJrbLzirnX3X0Pbk93S2dE5vs/2xsL6fqV18EPkVXO1mQtqsM8sMF8o6G/PLILN268Ga7CwcCL3qnoaCvahN3sHbciy38UH6s5hRTDvV75nWDj4oIaByrYx+JdgSZ4sucAn/bEQJCDSTVQ3sYQbEJGxc+xImNWudEoxdCmKYZPDFhUEIfO8pQRHTX8ZHZST+m97kJMuPvWg49UlrGu2YE6KbkNBEz7cSoWOuWpbrNjv1I8XKf8Sd82dvRWn3ZDc/4GXXE5oscG8UHTlz3XIpWNNrpE+wmn+AvmU0+n5r4Nv0LOFrlqH8Z2DcfjGqAVJkQMWFriruEcsPOvRgvGUeUtjulxEwcqX/UVmE5871rx0C2aJhazTnLkzt9TDFTaAf7J7zkIkhvKx8AU2A=";
608    const SHARE_RECOVERY: &str = "ME-RS-N-Mi0yLUNBSVFBaGdESXFBRWk3b296TFVtbzNscG1jZEIxMXFESnVHbUpRUGYxUk9nOG5WQVNkR1NSTE5YQk41VytwV1dUcWVNSnVOakN3ZVd4eFk0Ylh6Y0NWTlhzYWFLNmM5UEUxWHVFT2lNVW9BeHdQNkRtM3BCT29EVklHWlFYbXZxQk9FOVYyU3FCZTlsblVRcUZBak9SQUllQjFVcWEvdlJMbXN3VWNqZlNJN0pVQWgvL3ZKZVBvbUxGUVFYcjZVSUE5dnpsaDgxVVNjaXZHUXlnOVQydWRNS202RTdveXQvcVpGam1DdUlYSFlkR1FCdkxrK01oaEErVmh2NlM2a0FkSU5veWRGUlZTdXpnU25zT25wcUJxb21oRWZaNkdmb1dsaHM5UUFadXRmeUgzdkxRT0hQeXc1TEZLbUE3dnpOTTJmMkc2dGZaZGR1Q2dnT2gydmZCUnh1ZmJSdStHN2VGSGtLdnVoOW16ekQ3YUF2Z3BRbldKakdrai9paGcyR1EyQVlBWWkrM200SXR5V3J3Q1ZRNDZlUjJCRGpmZllQK3BOTFNnL2xKNlNmbUswd204R2cvL2ZpbVBHODF4alcrckdIQkczUTV4U1JnWnlUcmt0TEFBWlY0VndJOENzdTlmSUNlT0tTYm9UVTVrOFJoWnZRS0pUNnhGS1d1K0l3OTVoWUlUZUdmVGxLa1NtdW93WmVXcFI1TjIyQUEyWENaVWVqVVBLdHlQWUYxOVNGTjRjUDlvTURuUng3bkxkY3B6cGE3QmJWQm1jRGNhTENZVW1PeXJKcDQwK1hoekJHVmlxVnBJTS9qNTJJQTg1TSt1TDVtM0xNUk12UFc4cEliNkpVYVlKV0FXSldWV3JKdlpzUGJrdmh3T0NlOXo5VWJUTXE1WUJzOC9OcFJnN0F4L2lmSTJ5ZHdxbDRacXd4N29MM0ZrK0daK1FDeHRyTWJSK2oxTzhROXFXOEJ5eEcveXFBQWdDazlzckhwaG51UnpTck95M1JmZzRYa0lndHhlb3ArUmZJOVgyaDBRcEVmcjgzYzExd0xhQkxDUmgwMlFXazA2Ty8yM2s2cWZNZHBxNVZ2b0ZnTkNJYlY1V01sSFpaV3RnVXFzaGtXRVJycjduZnVvd1BQQ0NUaHdxMC9tbUQ1NDVDb0VNWU16bUtQYlIyYmF4RkVTbUswTlRRT3VWR3A2Y3JqNWlYOGxzaU9kZ3FVNHhuSVpRcDRsT1lJcTlBOUhFS3NZZ1RuYysxRlRNazJEN05ydThlalh4UUR3amFqUTFNTmJ5cldBS0MvZ3RTWW9ONTFKY25FWFlUOWI1MVZOWWF4anArTE9oeDA0M3RNUW9TejNvN1kxbWtORlJUTmJMZWhDKzV0UkNKNjdQYk5Va29DbWxXbjFYODZxVlVsRFU0MjkxaXVLaE1YQ0lsVHlscGU2dw==";
609    const SHARE_PASSWORD: &str = "mnemonic share password";
610
611    fn get_user_repo() -> (&'static EncryptionPin, UserRepoT) {
612        let mut repo = Box::new(UserRepoImpl::new(MemoryUserStorage::new())) as UserRepoT;
613        repo.create(&crate::types::users::UserEntity {
614            user_id: None,
615            username: USERNAME.to_string(),
616            encrypted_password: Some(ENCRYPTED_WALLET_PASSWORD.clone()),
617            salt: SALT.into(),
618            is_kyc_verified: false,
619            kyc_type: KycType::Undefined,
620            viviswap_state: None,
621            local_share: None,
622            wallet_transactions: Vec::new(),
623        })
624        .unwrap();
625
626        (&PIN, repo)
627    }
628
629    #[rstest]
630    #[case(MNEMONIC, true)] // Valid mnemonic
631    #[case("", false)] // Empty mnemonic
632    #[case(MNEMONIC_INCORRECT, false)] // Incorrect mnemonic
633    #[tokio::test]
634    async fn test_create_wallet_from_mnemonic(#[case] mnemonic: &str, #[case] should_succeed: bool) {
635        // Arrange
636        let (config, _cleanup) = Config::new_test_with_cleanup();
637        let mut manager = WalletManagerImpl::new(USERNAME);
638        let (pin, mut repo) = get_user_repo();
639
640        // Act
641        let result = manager
642            .create_wallet_from_existing_mnemonic(&config, &None, &mut repo, pin, mnemonic)
643            .await;
644
645        // Assert
646        if should_succeed {
647            result.expect("should create wallet from valid mnemonic");
648            assert!(manager.recovery_share.is_some(), "Wallet was successfully created");
649        } else {
650            result.expect_err("should fail with invalid or empty mnemonic");
651        }
652    }
653
654    #[rstest]
655    #[case(&WALLET_PASSWORD, Ok(()))]
656    #[case(&INVALID_BACKUP_PASSWORD, Err(WalletError::KdbxStorage(KdbxStorageError::UnlockError(UnlockError::HmacInvalid))))]
657    #[tokio::test]
658    async fn test_backup_and_restore(#[case] password: &LazyLock<PlainPassword>, #[case] should_succeed: Result<()>) {
659        // Arrange
660        let (config, _cleanup) = Config::new_test_with_cleanup();
661        let mut manager = WalletManagerImpl::new(USERNAME);
662        let (pin, mut repo) = get_user_repo();
663
664        // Create wallet
665        manager
666            .create_wallet_from_new_mnemonic(&config, &None, &mut repo, pin)
667            .await
668            .expect("failed to create new wallet");
669
670        // Create backup
671        let backup = manager
672            .create_wallet_backup(&config, &None, &mut repo, pin, &WALLET_PASSWORD)
673            .await
674            .expect("failed to create backup");
675
676        // Backup restoration
677        let restore_result = manager
678            .create_wallet_from_backup(&config, &None, &mut repo, pin, &backup, password)
679            .await;
680
681        // Assert
682        match should_succeed {
683            Ok(_) => restore_result.unwrap(),
684            Err(ref expected_err) => {
685                assert_eq!(restore_result.err().unwrap().to_string(), expected_err.to_string());
686            }
687        }
688    }
689
690    #[tokio::test]
691    async fn test_change_password() {
692        //Arrange
693        let (mut config, _cleanup) = Config::new_test_with_cleanup();
694        let mut manager = WalletManagerImpl::new(USERNAME);
695        let (pin, mut repo) = get_user_repo();
696
697        // create a wallet
698        manager
699            .create_wallet_from_new_mnemonic(&config, &None, &mut repo, pin)
700            .await
701            .expect("should succeed to create new wallet");
702
703        let new_password = PlainPassword::try_from_string("new_correcthorsebatterystaple").unwrap();
704        manager
705            .change_wallet_password(&config, &None, &mut repo, pin, &new_password)
706            .await
707            .expect("should succeed to change wallet password");
708
709        let wallet = manager
710            .try_get(
711                &mut config,
712                &None,
713                &mut repo,
714                &example_api_network(IOTA_NETWORK_KEY.to_string()),
715                pin,
716            )
717            .await
718            .expect("should succeed to get wallet after password change");
719
720        wallet.get_address().await.expect("wallet should return an address");
721    }
722
723    #[tokio::test]
724    async fn delete_wallet_removes_files() {
725        //Arrange
726        let (mut config, _cleanup) = Config::new_test_with_cleanup();
727        let mut manager = WalletManagerImpl::new(USERNAME);
728        let (pin, mut repo) = get_user_repo();
729
730        // create a wallet
731        manager
732            .create_wallet_from_new_mnemonic(&config, &None, &mut repo, pin)
733            .await
734            .expect("should succeed to create new wallet");
735
736        // get the wallet instance to make sure any files are created
737        let _wallet = manager
738            .try_get(
739                &mut config,
740                &None,
741                &mut repo,
742                &example_api_network(IOTA_NETWORK_KEY.to_string()),
743                pin,
744            )
745            .await
746            .expect("should succeed to get wallet");
747
748        let file_count_before = walkdir::WalkDir::new(&config.path_prefix).into_iter().count();
749
750        //Act
751        manager
752            .delete_wallet(&config, &None, &mut repo)
753            .await
754            .expect("should delete wallet");
755
756        // Assert
757        let file_count_after = walkdir::WalkDir::new(&config.path_prefix).into_iter().count();
758        assert!(file_count_after < file_count_before, "should remove files");
759    }
760
761    // Note: all test cases assume that there is no password stored in the user database (since a wallet was never created before)
762    #[rstest::rstest]
763    // ############### Test cases without the local storage available ###############
764    // nothing available at all
765    #[case(
766        None,
767        None,
768        None,
769        None,
770        None,
771        Err(WalletError::WalletNotInitialized(ErrorKind::UseMnemonic))
772    )]
773    // only recovery share available
774    #[case(
775        None,
776        None,
777        Some(SHARE_RECOVERY),
778        None,
779        None,
780        Err(WalletError::WalletNotInitialized(ErrorKind::UseMnemonic))
781    )]
782    // only backup share available
783    #[case(
784        None,
785        None,
786        None,
787        Some(SHARE_RECOVERY),
788        None,
789        Err(WalletError::WalletNotInitialized(ErrorKind::SetRecoveryShare))
790    )]
791    // local recovery and backup share available, but no password
792    #[case(
793        None,
794        Some(SHARE_RECOVERY),
795        None,
796        Some(SHARE_BACKUP),
797        None,
798        Err(WalletError::WalletNotInitialized(ErrorKind::MissingPassword))
799    )]
800    // local recovery and backup share available, and password
801    #[case(
802        None,
803        Some(SHARE_RECOVERY),
804        None,
805        Some(SHARE_BACKUP),
806        Some(SHARE_PASSWORD),
807        Ok(Status {local: false, recovery: Some(RecoveryUsed::Local), backup: true })
808    )]
809    // recovery and backup available but no password provided
810    #[case(
811        None,
812        None,
813        Some(SHARE_RECOVERY),
814        Some(SHARE_BACKUP),
815        None,
816        Err(WalletError::WalletNotInitialized(ErrorKind::MissingPassword))
817    )]
818    // recovery and backup available and password provided
819    #[case(
820        None,
821        None,
822        Some(SHARE_RECOVERY),
823        Some(SHARE_BACKUP),
824        Some(SHARE_PASSWORD),
825        Ok(Status{local: false, recovery: Some(RecoveryUsed::Remote), backup: true })
826    )]
827    // ############### Test cases with the local storage available ###############
828    #[case(
829        Some(SHARE_LOCAL),
830        None,
831        None,
832        None,
833        None,
834        Err(WalletError::WalletNotInitialized(ErrorKind::SetRecoveryShare))
835    )]
836    #[case(Some(SHARE_LOCAL), Some(SHARE_RECOVERY), None, None, None, Ok(Status{local: true, recovery: Some(RecoveryUsed::Local), backup: false }))]
837    #[case(Some(SHARE_LOCAL), None, Some(SHARE_RECOVERY), None, None, Ok(Status{local: true, recovery: Some(RecoveryUsed::Remote), backup: false}))]
838    #[case(
839        Some(SHARE_LOCAL),
840        None,
841        None,
842        Some(SHARE_RECOVERY),
843        None,
844        Err(WalletError::WalletNotInitialized(ErrorKind::MissingPassword))
845    )]
846    #[case(
847        Some(SHARE_LOCAL),
848        None,
849        None,
850        Some(SHARE_RECOVERY),
851        Some(SHARE_PASSWORD),
852        Ok(Status{local: true, recovery: None, backup: true})
853    )]
854    #[tokio::test]
855    async fn test_resemble_shares(
856        #[case] local_share: Option<&str>,
857        #[case] local_recovery_share: Option<&str>,
858        #[case] remote_recovery_share: Option<&str>,
859        #[case] remote_backup_share: Option<&str>,
860        #[case] password: Option<&str>,
861        #[case] expected_result: Result<Status>,
862    ) {
863        use crate::{
864            share::Share,
865            wallet_manager::{WalletManager, WalletManagerImpl},
866        };
867
868        // setup the sdk
869        let mut srv = mockito::Server::new_async().await;
870        let url = format!("{}/api", srv.url());
871
872        let m1_body = remote_recovery_share.map(|s| format!("{{ \"share\":\"{s}\" }}"));
873        let m1 = srv
874            .mock("GET", "/api/user/shares/recovery")
875            .expect(if local_recovery_share.is_none() {
876                1 + if expected_result.is_ok() { 1 } else { 0 } // if Ok we try to create wallet
877            } else {
878                0
879            }) // do not download the recovery share if it exists locally
880            .with_status(if remote_recovery_share.is_some() { 200 } else { 404 })
881            .with_header("content-type", "application/json")
882            .with_body(m1_body.unwrap_or_default())
883            .create();
884
885        let m2_body = remote_backup_share.map(|s| format!("{{ \"share\":\"{s}\" }}"));
886        let m2 = srv
887            .mock("GET", "/api/user/shares/backup")
888            .expect(
889                // only expect a call if we don't have enough shares from the others
890                if local_share.is_some() as usize
891                    + local_recovery_share.is_some() as usize
892                    + remote_recovery_share.is_some() as usize
893                    >= 2
894                {
895                    0
896                } else {
897                    1 + if expected_result.is_ok() { 1 } else { 0 } // if Ok we try to create wallet
898                },
899            )
900            .with_status(if remote_backup_share.is_some() { 200 } else { 404 })
901            .with_header("content-type", "application/json")
902            .with_body(m2_body.unwrap_or_default())
903            .create();
904
905        // Initialize your Sdk instance with necessary parameters
906
907        let (mut config, _cleanup) = Config::new_test_with_cleanup_url(&url);
908
909        let access_token = Some(AccessToken::try_from_string("a fake token").unwrap());
910        let mut repo = MockUserRepo::new();
911
912        let mut manager = WalletManagerImpl::new("share_user");
913
914        // setup shares if provided
915
916        let salt = EncryptionSalt::generate();
917        let pin = EncryptionPin::try_from_string("123456").unwrap();
918        let encrypted_password =
919            password.map(|s| PlainPassword::try_from_string(s).unwrap().encrypt(&pin, &salt).unwrap());
920        let user = crate::types::users::UserEntity {
921            user_id: None,
922            username: "share_user".to_string(),
923            encrypted_password,
924            salt,
925            is_kyc_verified: true,
926            kyc_type: KycType::Undefined,
927            viviswap_state: None,
928            local_share: local_share.map(|s| s.to_string()),
929            wallet_transactions: Vec::new(),
930        };
931
932        repo.expect_get().returning(move |_| Ok(user.clone()));
933
934        manager.recovery_share = local_recovery_share.map(|s| s.parse::<Share>().unwrap());
935
936        // If the expected result is OK and there was no local share from the beginning, we expect
937        // the local share to be set to a valid share.
938        if expected_result.is_ok() && local_share.is_none() {
939            repo.expect_set_local_share().returning(|_, _| Ok(()));
940        }
941        let mut repo = Box::new(repo) as UserRepoT;
942
943        // Function you want to test
944        let result = manager
945            .try_resemble_shares(&config, &access_token, &mut repo, &pin)
946            .await
947            .map(|(_mnemonic, status)| status);
948
949        // Assert
950        match (&result, expected_result) {
951            (Ok(s), Ok(s2)) => assert_eq!(s, &s2),
952            (Err(WalletError::WalletNotInitialized(k)), Err(WalletError::WalletNotInitialized(k2))) => {
953                assert_eq!(*k, k2)
954            }
955            (other, other2) => panic!("Expected {other2:?} but got {other:?}"),
956        }
957
958        // if the result is Ok, make sure we have access to a valid wallet
959        if result.is_ok() {
960            tokio::time::sleep(std::time::Duration::from_secs(1)).await;
961
962            let wallet = manager
963                .try_get(
964                    &mut config,
965                    &access_token,
966                    &mut repo,
967                    &example_api_network(IOTA_NETWORK_KEY.to_string()),
968                    &pin,
969                )
970                .await
971                .unwrap();
972
973            let address = wallet.get_address().await.unwrap();
974            let balance = wallet.get_balance().await.unwrap();
975            println!("Recevier address: {address}, balance = {balance:?}");
976
977            // This check is nice, but does not play nice with nextest running the tests in parallel
978            // from the CI, commented out for now.
979            // assert!(balance.base_coin() > 0);
980        }
981
982        m1.assert();
983        m2.assert();
984    }
985}