add account encryption tests, update aes tests

This commit is contained in:
Jack Robison 2018-09-24 13:57:17 -04:00 committed by Lex Berezhny
parent c0e0b4b745
commit c45c792657
2 changed files with 113 additions and 10 deletions

View file

@ -323,3 +323,95 @@ class TestSingleKeyAccount(unittest.TestCase):
self.maxDiff = None
account_data['ledger'] = 'btc_mainnet'
self.assertDictEqual(account_data, account.to_dict())
class AccountEncryptionTests(unittest.TestCase):
password = "password"
init_vector = b'0000000000000000'
unencrypted_account = {
'name': 'My Account',
'seed':
"carbon smart garage balance margin twelve chest sword toast envelope bottom stomac"
"h absent",
'encrypted': False,
'private_key':
'xprv9s21ZrQH143K3TsAz5efNV8K93g3Ms3FXcjaWB9fVUsMwAoE3ZT4vYymkp'
'5BxKKfnpz8J6sHDFriX1SnpvjNkzcks8XBnxjGLS83BTyfpna',
'public_key':
'xpub661MyMwAqRbcFwwe67Bfjd53h5WXmKm6tqfBJZZH3pQLoy8Nb6mKUMJFc7'
'UbpVNzmwFPN2evn3YHnig1pkKVYcvCV8owTd2yAcEkJfCX53g',
'address_generator': {'name': 'single-address'}
}
encrypted_account = {
'name': 'My Account',
'seed':
"MDAwMDAwMDAwMDAwMDAwMJ4e4W4pE6nQtPiD6MujNIQ7aFPhUBl63GwPziAgGN"
"MBTMoaSjZfyyvw7ELMCqAYTWJ61aV7K4lmd2hR11g9dpdnnpCb9f9j3zLZHRv7+"
"bIkZ//trah9AIkmrc/ZvNkC0Q==",
'encrypted': True,
'private_key':
'MDAwMDAwMDAwMDAwMDAwMLkWikOLScA/ZxlFSGU7dl//7Q/1gS9h7vqQyrd8DX+'
'jwcp7SwlJ1mkMwuraUaWLq9/LxiaGmqJBUZ50p77YVZbDycaCN1unBr1/i1q6RP'
'Ob2MNCaG8nyjxZhQai+V/2JmJ+UnFMp3nHany7F8/Hr0g=',
'public_key':
'xpub661MyMwAqRbcFwwe67Bfjd53h5WXmKm6tqfBJZZH3pQLoy8Nb6mKUMJFc7'
'UbpVNzmwFPN2evn3YHnig1pkKVYcvCV8owTd2yAcEkJfCX53g',
'address_generator': {'name': 'single-address'}
}
def setUp(self):
self.ledger = ledger_class({
'db': ledger_class.database_class(':memory:'),
'headers': ledger_class.headers_class(':memory:'),
})
def test_encrypt_wallet(self):
account = self.ledger.account_class.from_dict(self.ledger, Wallet(), self.unencrypted_account)
account.encryption_init_vector = self.init_vector
self.assertFalse(account.serialize_encrypted)
self.assertFalse(account.encrypted)
account.encrypt(self.password)
self.assertFalse(account.serialize_encrypted)
self.assertTrue(account.encrypted)
self.assertEqual(account.seed, self.encrypted_account['seed'])
self.assertEqual(account.private_key, self.encrypted_account['private_key'])
self.assertEqual(account.to_dict()['seed'], self.encrypted_account['seed'])
self.assertEqual(account.to_dict()['private_key'], self.encrypted_account['private_key'])
account.serialize_encrypted = True
account.decrypt(self.password)
self.assertEqual(account.seed, self.unencrypted_account['seed'])
self.assertEqual(account.private_key.extended_key_string(), self.unencrypted_account['private_key'])
self.assertEqual(account.to_dict()['seed'], self.encrypted_account['seed'])
self.assertEqual(account.to_dict()['private_key'], self.encrypted_account['private_key'])
account.encryption_init_vector = None
self.assertNotEqual(account.to_dict()['seed'], self.encrypted_account['seed'])
self.assertNotEqual(account.to_dict()['private_key'], self.encrypted_account['private_key'])
self.assertFalse(account.encrypted)
self.assertTrue(account.serialize_encrypted)
def test_decrypt_wallet(self):
account = self.ledger.account_class.from_dict(self.ledger, Wallet(), self.encrypted_account)
account.encryption_init_vector = self.init_vector
self.assertTrue(account.encrypted)
self.assertTrue(account.serialize_encrypted)
account.decrypt(self.password)
self.assertFalse(account.encrypted)
self.assertTrue(account.serialize_encrypted)
self.assertEqual(account.seed, self.unencrypted_account['seed'])
self.assertEqual(account.private_key.extended_key_string(), self.unencrypted_account['private_key'])
self.assertEqual(account.to_dict()['seed'], self.encrypted_account['seed'])
self.assertEqual(account.to_dict()['private_key'], self.encrypted_account['private_key'])
account.serialize_encrypted = False
self.assertEqual(account.to_dict()['seed'], self.unencrypted_account['seed'])
self.assertEqual(account.to_dict()['private_key'], self.unencrypted_account['private_key'])

View file

@ -8,18 +8,29 @@ except ImportError:
class TestAESEncryptDecrypt(TestCase):
message = 'The Times 03/Jan/2009 Chancellor on brink of second bailout for banks'
expected = 'ZmZmZmZmZmZmZmZmZmZmZjlrKptoKD+MFwDxcg3XtCD9qz8UWhEhq/TVJT5+Mtp2a8sE' \
'CaO6WQj7fYsWGu2Hvbc0qYqxdN0HeTsiO+cZRo3eJISgr3F+rXFYi5oSBlD2'
password = 'bubblegum'
@mock.patch('os.urandom', side_effect=lambda i: b'd'*i)
def test_encrypt_iv_f(self, _):
self.assertEqual(
aes_encrypt(self.password, self.message),
'ZGRkZGRkZGRkZGRkZGRkZKBP/4pR+47hLHbHyvDJm9aRKDuoBdTG8SrFvHqfagK6Co1VrHUOd'
'oF+6PGSxru3+VR63ybkXLNM75s/qVw+dnKVAkI8OfoVnJvGRSc49e38'
)
@mock.patch('os.urandom', side_effect=lambda i: b'f'*i)
def test_encrypt(self, _):
self.assertEqual(aes_encrypt(
b'bubblegum', b'The Times 03/Jan/2009 Chancellor on brink of second bailout for banks'),
b'OWsqm2goP4wXAPFyDde0IP2rPxRaESGr9NUlPn4y2nZrywQJo7pZCPt9ixYa7Ye9tzSpirF03Qd5OyI75xlGjd'
b'4khKCvcX6tcViLmhIGUPY='
def test_encrypt_iv_d(self, _):
self.assertEqual(
aes_encrypt(self.password, self.message),
'ZmZmZmZmZmZmZmZmZmZmZjlrKptoKD+MFwDxcg3XtCD9qz8UWhEhq/TVJT5+Mtp2a8sE'
'CaO6WQj7fYsWGu2Hvbc0qYqxdN0HeTsiO+cZRo3eJISgr3F+rXFYi5oSBlD2'
)
def test_decrypt(self):
self.assertEqual(aes_decrypt(
b'bubblegum', b'WeW99mQgRExAEzPjJOAC/MdTJaHgz3hT+kazFbvVQqF/KFva48ulVMOewU7JWD0ufWJIxtAIQ'
b'bGtlbvbq5w74bsCCJLrtNTHBhenkms8XccJXTr/UF/ZYTF1Prz8b0AQ'),
b'The Times 03/Jan/2009 Chancellor on brink of second bailout for banks'
def test_encrypt_decrypt(self):
self.assertEqual(
aes_decrypt('bubblegum', aes_encrypt('bubblegum', self.message)),
self.message
)