From c0e0b4b74586c4f89e4536a5918b68a19bd0a08d Mon Sep 17 00:00:00 2001 From: Jack Robison Date: Fri, 21 Sep 2018 14:49:16 -0400 Subject: [PATCH] review feedback and pylint --- torba/baseaccount.py | 39 +++++++++++++++++++++++---------------- torba/hash.py | 24 ++++++++++++------------ 2 files changed, 35 insertions(+), 28 deletions(-) diff --git a/torba/baseaccount.py b/torba/baseaccount.py index ee3c012b5..eb649e411 100644 --- a/torba/baseaccount.py +++ b/torba/baseaccount.py @@ -1,10 +1,10 @@ import typing -from typing import Dict, Tuple, Type, Optional, Any +from typing import Dict, Tuple, Type, Optional, Any, Union from twisted.internet import defer from torba.mnemonic import Mnemonic from torba.bip32 import PrivateKey, PubKey, from_extended_key_string -from torba.hash import double_sha256, aes_encrypt, aes_decrypt +from torba.hash import aes_encrypt, aes_decrypt from torba.constants import COIN if typing.TYPE_CHECKING: @@ -199,16 +199,18 @@ class BaseAccount: def __init__(self, ledger: 'baseledger.BaseLedger', wallet: 'basewallet.Wallet', name: str, seed: str, encrypted: bool, private_key: PrivateKey, public_key: PubKey, - address_generator: dict) -> None: + address_generator: dict, password: str = None) -> None: self.ledger = ledger self.wallet = wallet self.id = public_key.address self.name = name self.seed = seed - self.password = None + self.password = password + self.encryption_init_vector = None self.encrypted = encrypted - self.encrypted_on_disk = encrypted - self.private_key = private_key + self.serialize_encrypted = encrypted + + self.private_key: Union[PrivateKey, str] = private_key self.public_key = public_key generator_name = address_generator.get('name', HierarchicalDeterministic.name) self.address_generator = self.address_generators[generator_name] @@ -265,9 +267,9 @@ class BaseAccount: private_key, seed = self.private_key, self.seed if not self.encrypted and self.private_key: private_key = self.private_key.extended_key_string() - if not self.encrypted and self.encrypted_on_disk: - private_key = aes_encrypt(self.password, private_key) - seed = aes_encrypt(self.password, self.seed) + if not self.encrypted and self.serialize_encrypted: + private_key = aes_encrypt(self.password, private_key, init_vector=self.encryption_init_vector) + seed = aes_encrypt(self.password, self.seed, init_vector=self.encryption_init_vector) return { 'ledger': self.ledger.get_id(), @@ -295,20 +297,25 @@ class BaseAccount: details['seed'] = self.seed return details - def decrypt(self, password): + def decrypt(self, password: str) -> None: assert self.encrypted, "Key is not encrypted." - self.seed = aes_decrypt(password, self.seed.encode()).decode() + self.seed = aes_decrypt(password, self.seed) + p_k: Union[PrivateKey, str] = self.private_key + assert isinstance(p_k, str) self.private_key = from_extended_key_string( - self.ledger, aes_decrypt(password, self.private_key.encode()).decode() + self.ledger, aes_decrypt(password, str(p_k)) ) self.password = password self.encrypted = False - def encrypt(self, password): + def encrypt(self, password: str) -> None: assert not self.encrypted, "Key is already encrypted." - self.seed = aes_encrypt(password, self.seed.encode()).decode() - private_key: PrivateKey = self.private_key - self.private_key = aes_encrypt(password, private_key.extended_key_string().encode()).decode() + assert isinstance(self.private_key, PrivateKey) + self.seed = aes_encrypt(password, self.seed, init_vector=self.encryption_init_vector) + p_k: PrivateKey = self.private_key # this is because the type is changing from PrivateKey <-> str + extended: str = p_k.extended_key_string() + self.private_key = aes_encrypt(password, extended, + init_vector=self.encryption_init_vector) self.password = None self.encrypted = True diff --git a/torba/hash.py b/torba/hash.py index e6315fd2d..656dec2a6 100644 --- a/torba/hash.py +++ b/torba/hash.py @@ -106,27 +106,27 @@ def hex_str_to_hash(x): return reversed(unhexlify(x)) -def aes_encrypt(secret, value, iv=None): - if iv: - assert len(iv) == 16 +def aes_encrypt(secret: str, value: str, init_vector: bytes = None) -> str: + if init_vector is not None: + assert len(init_vector) == 16 else: - iv = os.urandom(16) - key = double_sha256(secret) - encryptor = Cipher(AES(key), modes.CBC(iv), default_backend()).encryptor() + 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() + padded_data = padder.update(value.encode()) + padder.finalize() encrypted_data = encryptor.update(padded_data) + encryptor.finalize() - return base64.b64encode(iv + encrypted_data) + return base64.b64encode(init_vector + encrypted_data).decode() -def aes_decrypt(secret, value): - data = base64.b64decode(value) - key = double_sha256(secret) +def aes_decrypt(secret: str, value: str) -> str: + data = base64.b64decode(value.encode()) + 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() result = unpadder.update(decryptor.update(data)) + unpadder.finalize() - return result + return result.decode() class Base58Error(Exception):