minor refactor

This commit is contained in:
Lex Berezhny 2022-10-17 23:59:26 -04:00
parent 94deaf55df
commit 734f0651a4
3 changed files with 30 additions and 38 deletions

View file

@ -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=<password>] [--wallet_id=<wallet_id>]
@ -1342,21 +1342,21 @@ class Daemon(metaclass=JSONRPCServerType):
--wallet_id=<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> | --data=<data>) [<password> | --password=<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):

View file

@ -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'])

View file

@ -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})