From 734f0651a47b6bb3f4bfbcdc0386526d928fb2a2 Mon Sep 17 00:00:00 2001 From: Lex Berezhny Date: Mon, 17 Oct 2022 23:59:26 -0400 Subject: [PATCH] minor refactor --- lbry/extras/daemon/daemon.py | 24 +++++++++---------- lbry/wallet/wallet.py | 22 +++++++++-------- .../blockchain/test_wallet_commands.py | 22 ++++++----------- 3 files changed, 30 insertions(+), 38 deletions(-) diff --git a/lbry/extras/daemon/daemon.py b/lbry/extras/daemon/daemon.py index b257fd2a0..d11d7bab8 100644 --- a/lbry/extras/daemon/daemon.py +++ b/lbry/extras/daemon/daemon.py @@ -1330,9 +1330,9 @@ class Daemon(metaclass=JSONRPCServerType): @requires("wallet") async def jsonrpc_wallet_export(self, password=None, wallet_id=None): """ - Export wallet data + Exports encrypted wallet data if password is supplied; otherwise plain JSON. - Wallet must be unlocked to perform this operation. Exports JSON if password is not supplied. + Wallet must be unlocked to perform this operation. Usage: wallet_export [--password=] [--wallet_id=] @@ -1342,21 +1342,21 @@ class Daemon(metaclass=JSONRPCServerType): --wallet_id= : (str) wallet being exported Returns: - (str) data: Either base64-encoded encrypted wallet, or cleartext JSON + (str) data: base64-encoded encrypted wallet, or cleartext JSON """ wallet = self.wallet_manager.get_wallet_or_default(wallet_id) - if password is None: - return wallet.pack() - encrypted = wallet.pack(password) - return encrypted.decode() + if password: + return wallet.pack(password).decode() + return wallet.to_json() @requires("wallet") async def jsonrpc_wallet_import(self, data, password=None, wallet_id=None, blocking=False): """ - Import wallet data and merge accounts and preferences. + Import wallet data and merge accounts and preferences. Data is expected to be JSON if + password is not supplied. - Wallet must be unlocked to perform this operation. Data is expected to be JSON if password is not supplied. + Wallet must be unlocked to perform this operation. Usage: wallet_import ( | --data=) [ | --password=] @@ -1369,9 +1369,8 @@ class Daemon(metaclass=JSONRPCServerType): --blocking : (bool) wait until any new accounts have merged Returns: - (str) Either base64-encoded encrypted wallet, or cleartext JSON + (str) base64-encoded encrypted wallet, or cleartext JSON """ - wallet = self.wallet_manager.get_wallet_or_default(wallet_id) added_accounts, merged_accounts = wallet.merge(self.wallet_manager, password, data) for new_account in itertools.chain(added_accounts, merged_accounts): @@ -1385,8 +1384,7 @@ class Daemon(metaclass=JSONRPCServerType): for new_account in added_accounts: asyncio.create_task(self.ledger.subscribe_account(new_account)) wallet.save() - encrypted = wallet.pack(password) - return encrypted.decode() + return await self.jsonrpc_wallet_export(password=password, wallet_id=wallet_id) @requires("wallet") async def jsonrpc_wallet_add(self, wallet_id): diff --git a/lbry/wallet/wallet.py b/lbry/wallet/wallet.py index 6c9ecb7b1..956b7c18f 100644 --- a/lbry/wallet/wallet.py +++ b/lbry/wallet/wallet.py @@ -139,6 +139,10 @@ class Wallet: 'accounts': [a.to_dict(encrypt_password) for a in self.accounts] } + def to_json(self): + assert not self.is_locked, "Cannot serialize a wallet with locked/encrypted accounts." + return json.dumps(self.to_dict()) + def save(self): if self.preferences.get(ENCRYPT_ON_DISK, False): if self.encryption_password is not None: @@ -163,19 +167,14 @@ class Wallet: h.update(account.hash) return h.digest() - def pack(self, password=None): + def pack(self, password): assert not self.is_locked, "Cannot pack a wallet with locked/encrypted accounts." - new_data = json.dumps(self.to_dict()) - if password is None: - return new_data.encode() - new_data_compressed = zlib.compress(new_data.encode()) + new_data_compressed = zlib.compress(self.to_json().encode()) return better_aes_encrypt(password, new_data_compressed) @classmethod - def unpack(cls, password, data): - if password is None: - return json.loads(data) - decrypted = better_aes_decrypt(password, data) + def unpack(cls, password, encrypted): + decrypted = better_aes_decrypt(password, encrypted) try: decompressed = zlib.decompress(decrypted) except zlib.error as e: @@ -190,7 +189,10 @@ class Wallet: password: str, data: str) -> (List['Account'], List['Account']): assert not self.is_locked, "Cannot sync apply on a locked wallet." added_accounts, merged_accounts = [], [] - decrypted_data = self.unpack(password, data) + if password: + decrypted_data = self.unpack(password, data) + else: + decrypted_data = json.loads(data) self.preferences.merge(decrypted_data.get('preferences', {})) for account_dict in decrypted_data['accounts']: ledger = manager.get_or_create_ledger(account_dict['ledger']) diff --git a/tests/integration/blockchain/test_wallet_commands.py b/tests/integration/blockchain/test_wallet_commands.py index 192fa9200..6ead8e0fe 100644 --- a/tests/integration/blockchain/test_wallet_commands.py +++ b/tests/integration/blockchain/test_wallet_commands.py @@ -493,19 +493,11 @@ class WalletEncryptionAndSynchronization(CommandTestCase): ) # test without passwords - daemon2.jsonrpc_preference_set("three", "3") - jsondata = await daemon2.jsonrpc_wallet_export() - await daemon.jsonrpc_wallet_import(data=jsondata, blocking=True) - self.assertDictEqual( - # "two" key added and "conflict" value changed to "2" - daemon.jsonrpc_preference_get(), - { - "one": "1", - "two": "2", - "three": "3", - "conflict": "2", - "another": "B", - "fruit": ["peach", "apricot"] - } - ) + data = await daemon2.jsonrpc_wallet_export() + json_data = json.loads(data) + self.assertEqual(json_data["name"], "Wallet") + self.assertNotIn("four", json_data["preferences"]) + json_data["preferences"]["four"] = {"value": 4, "ts": 0} + await daemon.jsonrpc_wallet_import(data=json.dumps(json_data), blocking=True) + self.assertEqual(daemon.jsonrpc_preference_get("four"), {"four": 4})