diff --git a/clive/__private/cli/commands/unlock.py b/clive/__private/cli/commands/unlock.py index 453c1788e7a864c139752dc2d0939d3025ac56f5..7538c211345a92e375e4992a7083cf58906d3bb5 100644 --- a/clive/__private/cli/commands/unlock.py +++ b/clive/__private/cli/commands/unlock.py @@ -4,7 +4,7 @@ import sys from dataclasses import dataclass from datetime import timedelta from getpass import getpass -from typing import TYPE_CHECKING, Final +from typing import Final import typer @@ -14,19 +14,11 @@ from clive.__private.cli.exceptions import ( CLIInvalidSelectionError, CLIProfileDoesNotExistsError, ) -from clive.__private.cli.notify import notify from clive.__private.core.constants.cli import UNLOCK_CREATE_PROFILE_HELP, UNLOCK_CREATE_PROFILE_SELECT -from clive.__private.core.constants.wallet_recovery import ( - USER_WALLET_RECOVERED_MESSAGE, - USER_WALLET_RECOVERED_NOTIFICATION_LEVEL, -) from clive.__private.core.error_handlers.abc.error_notificator import CannotNotifyError from clive.__private.core.error_handlers.general_error_notificator import INVALID_PASSWORD_MESSAGE from clive.__private.core.profile import Profile -if TYPE_CHECKING: - from clive.__private.core.commands.recover_wallets import RecoverWalletsStatus - PASSWORD_SELECTION_ATTEMPTS: Final[int] = 3 PROFILE_SELECTION_ATTEMPTS: Final[int] = 3 ProfileSelectionOptions = dict[int, str] @@ -81,19 +73,16 @@ class Unlock(WorldBasedCommand): async def _unlock_profile(self, profile_name: str, password: str) -> None: try: - result = ( - await self.world.commands.unlock( - profile_name=profile_name, - password=password, - time=self._duration, - permanent=self._is_unlock_permanent, - ) - ).result_or_raise + await self.world.commands.unlock( + profile_name=profile_name, + password=password, + time=self._duration, + permanent=self._is_unlock_permanent, + ) except CannotNotifyError as error: if INVALID_PASSWORD_MESSAGE in error.reason: raise CLIInvalidPasswordError(profile_name) from error raise - self._display_wallet_recovery_status(result) def _prompt_for_profile_name(self) -> str | None: options = self._generate_profile_options() @@ -166,7 +155,3 @@ class Unlock(WorldBasedCommand): def _display_create_profile_help_info(self) -> None: typer.echo(UNLOCK_CREATE_PROFILE_HELP) - - def _display_wallet_recovery_status(self, status: RecoverWalletsStatus) -> None: - if status == "user_wallet_recovered": - notify(USER_WALLET_RECOVERED_MESSAGE, level=USER_WALLET_RECOVERED_NOTIFICATION_LEVEL) diff --git a/clive/__private/core/commands/commands.py b/clive/__private/core/commands/commands.py index c1c543e2038693f5a046a56318fa4afc37515ed0..1249cd4bd52239f9faf4bd7dc24f6962c9b4b138 100644 --- a/clive/__private/core/commands/commands.py +++ b/clive/__private/core/commands/commands.py @@ -43,6 +43,7 @@ from clive.__private.core.commands.is_wallet_unlocked import IsWalletUnlocked from clive.__private.core.commands.load_profile import LoadProfile from clive.__private.core.commands.load_transaction import LoadTransaction from clive.__private.core.commands.lock import Lock +from clive.__private.core.commands.migrate_profile import MigrateProfile from clive.__private.core.commands.perform_actions_on_transaction import PerformActionsOnTransaction from clive.__private.core.commands.remove_key import RemoveKey from clive.__private.core.commands.save_profile import SaveProfile @@ -51,7 +52,7 @@ from clive.__private.core.commands.set_timeout import SetTimeout from clive.__private.core.commands.sign import ALREADY_SIGNED_MODE_DEFAULT, AlreadySignedMode, Sign from clive.__private.core.commands.sync_data_with_beekeeper import SyncDataWithBeekeeper from clive.__private.core.commands.sync_state_with_beekeeper import SyncStateWithBeekeeper -from clive.__private.core.commands.unlock import Unlock +from clive.__private.core.commands.unlock import Unlock, UnlockWalletStatus from clive.__private.core.commands.unsign import UnSign from clive.__private.core.commands.update_transaction_metadata import ( UpdateTransactionMetadata, @@ -80,14 +81,13 @@ if TYPE_CHECKING: from clive.__private.core.accounts.accounts import TrackedAccount from clive.__private.core.app_state import LockSource from clive.__private.core.commands.abc.command import Command - from clive.__private.core.commands.recover_wallets import RecoverWalletsStatus from clive.__private.core.ensure_transaction import TransactionConvertibleType from clive.__private.core.error_handlers.abc.error_handler_context_manager import ( AnyErrorHandlerContextManager, ) from clive.__private.core.keys import PrivateKeyAliased, PublicKey, PublicKeyAliased from clive.__private.core.profile import Profile - from clive.__private.core.types import NotifyLevel + from clive.__private.core.types import MigrationStatus, NotifyLevel from clive.__private.core.world import CLIWorld, TUIWorld, World from clive.__private.models import Transaction from clive.__private.models.schemas import ( @@ -172,7 +172,7 @@ class Commands[WorldT: World]: async def unlock( self, *, profile_name: str | None = None, password: str, time: timedelta | None = None, permanent: bool = True - ) -> CommandWithResultWrapper[RecoverWalletsStatus]: + ) -> CommandWithResultWrapper[UnlockWalletStatus]: """ Return a CommandWrapper instance to unlock the profile-related wallets (user keys and encryption key). @@ -184,12 +184,13 @@ class Commands[WorldT: World]: permanently. permanent: Whether to unlock the wallet permanently. Will take precedence when `time` is also set. """ + profile_to_unlock = profile_name or self._world.profile.name wrapper = await self.__surround_with_exception_handlers( Unlock( password=password, app_state=self._world.app_state, session=self._world.beekeeper_manager.session, - profile_name=profile_name or self._world.profile.name, + profile_name=profile_to_unlock, time=time, permanent=permanent, ) @@ -197,8 +198,10 @@ class Commands[WorldT: World]: if wrapper.success: result = wrapper.result_or_raise - if result == "user_wallet_recovered": + if result.recovery_status == "user_wallet_recovered": self._notify(USER_WALLET_RECOVERED_MESSAGE, level=USER_WALLET_RECOVERED_NOTIFICATION_LEVEL) + if result.migration_status == "migrated": + self._notify(f"Migration of profile `{profile_to_unlock}` was performed") return wrapper async def lock(self) -> CommandWrapper: @@ -524,6 +527,15 @@ class Commands[WorldT: World]: ) ) + async def migrate_profile(self, *, profile_name: str | None = None) -> CommandWithResultWrapper[MigrationStatus]: + return await self.__surround_with_exception_handlers( + MigrateProfile( + profile_name=profile_name or self._world.profile.name, + unlocked_wallet=self._world.beekeeper_manager.user_wallet, + unlocked_encryption_wallet=self._world.beekeeper_manager.encryption_wallet, + ) + ) + async def collect_account_authorities( self, *, account_name: str ) -> CommandWithResultWrapper[WaxAccountAuthorityInfo]: diff --git a/clive/__private/core/commands/migrate_profile.py b/clive/__private/core/commands/migrate_profile.py new file mode 100644 index 0000000000000000000000000000000000000000..ba233780d119e3b6b407017d7205088af29b618b --- /dev/null +++ b/clive/__private/core/commands/migrate_profile.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from clive.__private.core.commands.abc.command_encryption import CommandEncryption +from clive.__private.core.commands.abc.command_with_result import CommandWithResult +from clive.__private.core.encryption import EncryptionService +from clive.__private.core.types import MigrationStatus +from clive.__private.core.wallet_container import WalletContainer +from clive.__private.storage.service import PersistentStorageService + + +@dataclass(kw_only=True) +class MigrateProfile(CommandEncryption, CommandWithResult[MigrationStatus]): + profile_name: str + + async def _execute(self) -> None: + encryption_service = EncryptionService(WalletContainer(self.unlocked_wallet, self.unlocked_encryption_wallet)) + self._result = await PersistentStorageService(encryption_service).migrate(self.profile_name) diff --git a/clive/__private/core/commands/unlock.py b/clive/__private/core/commands/unlock.py index 443108b0aa7495275df0722556642516560812c8..24d449fec00a127f9a4640a185aec50962661ecb 100644 --- a/clive/__private/core/commands/unlock.py +++ b/clive/__private/core/commands/unlock.py @@ -7,6 +7,7 @@ from beekeepy.exceptions import NoWalletWithSuchNameError from clive.__private.core.commands.abc.command_secured import CommandPasswordSecured from clive.__private.core.commands.abc.command_with_result import CommandWithResult +from clive.__private.core.commands.migrate_profile import MigrateProfile from clive.__private.core.commands.recover_wallets import RecoverWallets, RecoverWalletsStatus from clive.__private.core.commands.set_timeout import SetTimeout from clive.__private.core.encryption import EncryptionService @@ -18,10 +19,17 @@ if TYPE_CHECKING: from beekeepy import AsyncSession, AsyncUnlockedWallet from clive.__private.core.app_state import AppState + from clive.__private.core.types import MigrationStatus + + +@dataclass +class UnlockWalletStatus: + recovery_status: RecoverWalletsStatus + migration_status: MigrationStatus @dataclass(kw_only=True) -class Unlock(CommandPasswordSecured, CommandWithResult[RecoverWalletsStatus]): +class Unlock(CommandPasswordSecured, CommandWithResult[UnlockWalletStatus]): """Unlock the profile-related wallets (user keys and encryption key) managed by the beekeeper.""" profile_name: str @@ -47,11 +55,11 @@ class Unlock(CommandPasswordSecured, CommandWithResult[RecoverWalletsStatus]): should_recover_user_wallet=is_user_wallet_missing, ).execute_with_result() - self._result = recover_wallets_result.status + recovery_status = recover_wallets_result.status - if self._result == "encryption_wallet_recovered": + if recovery_status == "encryption_wallet_recovered": encryption_wallet = recover_wallets_result.recovered_wallet - elif self._result == "user_wallet_recovered": + elif recovery_status == "user_wallet_recovered": user_wallet = recover_wallets_result.recovered_wallet assert user_wallet is not None, "User wallet should be created at this point" @@ -60,6 +68,14 @@ class Unlock(CommandPasswordSecured, CommandWithResult[RecoverWalletsStatus]): if self.app_state is not None: await self.app_state.unlock(WalletContainer(user_wallet, encryption_wallet)) + migration_result = await MigrateProfile( + profile_name=self.profile_name, + unlocked_wallet=user_wallet, + unlocked_encryption_wallet=encryption_wallet, + ).execute_with_result() + + self._result = UnlockWalletStatus(recovery_status=recovery_status, migration_status=migration_result) + async def _unlock_wallet(self, name: str) -> AsyncUnlockedWallet | None: try: wallet = await self.session.open_wallet(name=name) diff --git a/clive/__private/core/types.py b/clive/__private/core/types.py index bd777121fe54fd21f5ab50d72df284d1f5354ddc..d3214dd5a4acfc2e55c9a9ff6ccbb911440f392a 100644 --- a/clive/__private/core/types.py +++ b/clive/__private/core/types.py @@ -2,4 +2,5 @@ from __future__ import annotations from typing import Literal +MigrationStatus = Literal["migrated", "already_newest"] NotifyLevel = Literal["info", "warning", "error"] diff --git a/clive/__private/storage/service.py b/clive/__private/storage/service.py index cbc20de68d9621f8b834b1990621453b7a4db352..7db57ae8fc22c507e7984f8119999fe6c5cd77e5 100644 --- a/clive/__private/storage/service.py +++ b/clive/__private/storage/service.py @@ -22,6 +22,7 @@ if TYPE_CHECKING: from clive.__private.core.encryption import EncryptionService from clive.__private.core.profile import Profile + from clive.__private.core.types import MigrationStatus from clive.__private.storage import ProfileStorageBase, ProfileStorageModel @@ -69,6 +70,12 @@ class ProfileEncryptionError(PersistentStorageServiceError): super().__init__(self.MESSAGE) +@dataclass +class _MigrationResult: + model: ProfileStorageModel + status: MigrationStatus + + class PersistentStorageService: BACKUP_FILENAME_SUFFIX: Final[str] = ".backup" PROFILE_FILENAME_SUFFIX: Final[str] = ".profile" @@ -133,12 +140,17 @@ class PersistentStorageService: or communication with beekeeper failed. """ self._raise_if_profile_not_stored(profile_name) - profile_storage_model = await self._get_latest_stored_profile_model(profile_name) + result = await self._load_and_migrate_latest_profile_model(profile_name) + profile_storage_model = result.model profile = StorageToRuntimeConverter(profile_storage_model).create_profile() profile._update_hash_of_stored_profile() return profile + async def migrate(self, profile_name: str) -> MigrationStatus: + result = await self._load_and_migrate_latest_profile_model(profile_name) + return result.status + @classmethod def delete_profile(cls, profile_name: str, *, force: bool = False) -> None: """ @@ -323,7 +335,7 @@ class PersistentStorageService: filepath = profile_directory / self.get_current_version_profile_filename() filepath.write_text(encrypted_profile) - async def _get_latest_stored_profile_model(self, profile_name: str) -> ProfileStorageModel: + async def _load_and_migrate_latest_profile_model(self, profile_name: str) -> _MigrationResult: """ Find current version of profile storage model by name in the clive data directory or migrate older version. @@ -365,13 +377,15 @@ class PersistentStorageService: async def _migrate_profile_model( self, profile_model: ProfileStorageBase, profile_filepath: Path - ) -> ProfileStorageModel: - """Migrate profile model to the latest version.""" + ) -> _MigrationResult: + """Migrate profile model and return current version of model even it there was no migration needed.""" + was_migrated = False model_migrated = StorageHistory.apply_all_migrations(profile_model) if profile_model.get_this_version() != StorageHistory.get_latest_version(): await self._save_profile_model(model_migrated) self._move_profile_to_backup(profile_filepath) - return model_migrated + was_migrated = True + return _MigrationResult(model_migrated, status="migrated" if was_migrated else "already_newest") @classmethod def _move_profile_to_backup(cls, path: Path) -> None: diff --git a/tests/functional/cli/test_storage_migration.py b/tests/functional/cli/test_storage_migration.py index ff040d1a66c22c5b6dff381729a622ccc0279705..2095588ddc85c41c359365950e8428619a43fd40 100644 --- a/tests/functional/cli/test_storage_migration.py +++ b/tests/functional/cli/test_storage_migration.py @@ -35,15 +35,14 @@ def _get_profile_file_version(profile_name: str) -> int: async def test_unlock_and_migrate(cli_tester_locked: CLITester) -> None: # ARRANGE - cli_tester_locked.unlock(profile_name=PROFILE_NAME, password_stdin=PROFILE_PASSWORD) - cli_tester = cli_tester_locked assert _get_profile_file_version(PROFILE_NAME) == v0.ProfileStorageModel.get_this_version(), ( f"before migration profile should be stored in version {v0.ProfileStorageModel.get_this_version()}" ) - # ACT & ASSERT - result = cli_tester.show_profile() # remove after resolving issue https://gitlab.syncad.com/hive/clive/-/issues/413 - assert f"Profile name: {PROFILE_NAME}" in result.output, "profile is not loaded" + # ACT + cli_tester_locked.unlock(profile_name=PROFILE_NAME, password_stdin=PROFILE_PASSWORD) + + # ASSERT assert _get_profile_file_version(PROFILE_NAME) == ProfileStorageModel.get_this_version(), ( f"after migration profile should be stored in version {ProfileStorageModel.get_this_version()}" )