diff --git a/torba/client/baseaccount.py b/torba/client/baseaccount.py index 86d1d0d7d..c7cab4d55 100644 --- a/torba/client/baseaccount.py +++ b/torba/client/baseaccount.py @@ -1,3 +1,4 @@ +import json import asyncio import random import typing @@ -5,7 +6,7 @@ from typing import Dict, Tuple, Type, Optional, Any, List from torba.client.mnemonic import Mnemonic from torba.client.bip32 import PrivateKey, PubKey, from_extended_key_string -from torba.client.hash import aes_encrypt, aes_decrypt +from torba.client.hash import aes_encrypt, aes_decrypt, sha256 from torba.client.constants import COIN if typing.TYPE_CHECKING: @@ -235,7 +236,8 @@ class BaseAccount: ) @classmethod - def from_dict(cls, ledger: 'baseledger.BaseLedger', wallet: 'basewallet.Wallet', d: dict): + def keys_from_dict(cls, ledger: 'baseledger.BaseLedger', d: dict) \ + -> Tuple[str, Optional[PrivateKey], PubKey]: seed = d.get('seed', '') private_key_string = d.get('private_key', '') private_key = None @@ -250,6 +252,11 @@ class BaseAccount: public_key = private_key.public_key if public_key is None: public_key = from_extended_key_string(ledger, d['public_key']) + return seed, private_key, public_key + + @classmethod + def from_dict(cls, ledger: 'baseledger.BaseLedger', wallet: 'basewallet.Wallet', d: dict): + seed, private_key, public_key = cls.keys_from_dict(ledger, d) name = d.get('name') if not name: name = 'Account #{}'.format(public_key.address) @@ -258,8 +265,8 @@ class BaseAccount: wallet=wallet, name=name, seed=seed, - private_key_string=private_key_string, - encrypted=encrypted, + private_key_string=d.get('private_key', ''), + encrypted=d.get('encrypted', False), private_key=private_key, public_key=public_key, address_generator=d.get('address_generator', {}) @@ -273,8 +280,8 @@ class BaseAccount: assert None not in [self.seed_encryption_init_vector, self.private_key_encryption_init_vector] private_key_string = aes_encrypt( self.password, private_key_string, self.private_key_encryption_init_vector - ) - seed = aes_encrypt(self.password, self.seed, self.seed_encryption_init_vector) + )[0] + seed = aes_encrypt(self.password, self.seed, self.seed_encryption_init_vector)[0] return { 'ledger': self.ledger.get_id(), 'name': self.name, @@ -285,6 +292,10 @@ class BaseAccount: 'address_generator': self.address_generator.to_dict(self.receiving, self.change) } + @property + def hash(self) -> bytes: + return sha256(json.dumps(self.to_dict()).encode()) + async def get_details(self, show_seed=False, **kwargs): satoshis = await self.get_balance(**kwargs) details = { @@ -329,10 +340,10 @@ class BaseAccount: assert not self.encrypted, "Key is already encrypted." assert isinstance(self.private_key, PrivateKey) - self.seed = aes_encrypt(password, self.seed, self.seed_encryption_init_vector) + self.seed = aes_encrypt(password, self.seed, self.seed_encryption_init_vector)[0] self.private_key_string = aes_encrypt( password, self.private_key.extended_key_string(), self.private_key_encryption_init_vector - ) + )[0] self.private_key = None self.password = None self.encrypted = True diff --git a/torba/client/hash.py b/torba/client/hash.py index dcd61a080..bb18ee57a 100644 --- a/torba/client/hash.py +++ b/torba/client/hash.py @@ -121,7 +121,7 @@ def hex_str_to_hash(x): return reversed(unhexlify(x)) -def aes_encrypt(secret: str, value: str, init_vector: bytes = None) -> str: +def aes_encrypt(secret: str, value: str, init_vector: bytes = None) -> typing.Tuple[str, bytes]: if init_vector is not None: assert len(init_vector) == 16 else: @@ -144,6 +144,25 @@ def aes_decrypt(secret: str, value: str) -> typing.Tuple[str, bytes]: return result.decode(), init_vector +def better_aes_encrypt(secret: str, value: bytes) -> bytes: + init_vector = os.urandom(16) + key = double_sha256(secret.encode()) + encryptor = Cipher(AES(key), modes.CBC(init_vector), default_backend()).encryptor() + padder = PKCS7(AES.block_size).padder() + padded_data = padder.update(value) + padder.finalize() + encrypted_data = encryptor.update(padded_data) + encryptor.finalize() + return base64.b64encode(init_vector + encrypted_data) + + +def better_aes_decrypt(secret: str, value: bytes) -> bytes: + data = base64.b64decode(value) + key = double_sha256(secret.encode()) + init_vector, data = data[:16], data[16:] + decryptor = Cipher(AES(key), modes.CBC(init_vector), default_backend()).decryptor() + unpadder = PKCS7(AES.block_size).unpadder() + return unpadder.update(decryptor.update(data)) + unpadder.finalize() + + class Base58Error(Exception): """ Exception used for Base58 errors. """ diff --git a/torba/client/wallet.py b/torba/client/wallet.py index 66dc23252..166891899 100644 --- a/torba/client/wallet.py +++ b/torba/client/wallet.py @@ -1,8 +1,12 @@ +import os import stat import json -import os +import zlib import typing +from hashlib import sha256 from typing import Sequence, MutableSequence +from torba.client.hash import better_aes_encrypt, better_aes_decrypt +from operator import attrgetter if typing.TYPE_CHECKING: from torba.client import basemanager, baseaccount, baseledger @@ -55,6 +59,24 @@ class Wallet: for account in self.accounts: return account + @property + def hash(self) -> str: + h = sha256() + for account in sorted(self.accounts, key=attrgetter('id')): + h.update(account.hash) + return h.digest() + + def pack(self, password): + new_data = json.dumps(self.to_dict()) + new_data_compressed = zlib.compress(new_data.encode()) + return better_aes_encrypt(password, new_data_compressed) + + @classmethod + def unpack(cls, password, encrypted): + decrypted = better_aes_decrypt(password, encrypted) + decompressed = zlib.decompress(decrypted) + return json.loads(decompressed) + class WalletStorage: