apply scrypt as kdf, add tests

This commit is contained in:
Victor Shyba 2019-03-11 17:55:56 -03:00 committed by Lex Berezhny
parent fcac4f1488
commit 1286417aec
2 changed files with 18 additions and 4 deletions

View file

@ -1,5 +1,5 @@
from unittest import TestCase, mock from unittest import TestCase, mock
from torba.client.hash import aes_decrypt, aes_encrypt from torba.client.hash import aes_decrypt, aes_encrypt, better_aes_decrypt, better_aes_encrypt
class TestAESEncryptDecrypt(TestCase): class TestAESEncryptDecrypt(TestCase):
@ -33,3 +33,10 @@ class TestAESEncryptDecrypt(TestCase):
aes_decrypt('bubblegum', aes_encrypt('bubblegum', self.message))[0], aes_decrypt('bubblegum', aes_encrypt('bubblegum', self.message))[0],
self.message self.message
) )
def test_better_encrypt_decrypt(self):
self.assertEqual(
b'valuable value',
better_aes_decrypt(
'super secret',
better_aes_encrypt('super secret', b'valuable value')))

View file

@ -14,6 +14,7 @@ import hashlib
import hmac import hmac
import typing import typing
from binascii import hexlify, unhexlify from binascii import hexlify, unhexlify
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
from cryptography.hazmat.primitives.ciphers import Cipher, modes from cryptography.hazmat.primitives.ciphers import Cipher, modes
from cryptography.hazmat.primitives.ciphers.algorithms import AES from cryptography.hazmat.primitives.ciphers.algorithms import AES
from cryptography.hazmat.primitives.padding import PKCS7 from cryptography.hazmat.primitives.padding import PKCS7
@ -146,23 +147,29 @@ def aes_decrypt(secret: str, value: str) -> typing.Tuple[str, bytes]:
def better_aes_encrypt(secret: str, value: bytes) -> bytes: def better_aes_encrypt(secret: str, value: bytes) -> bytes:
init_vector = os.urandom(16) init_vector = os.urandom(16)
key = double_sha256(secret.encode()) key = scrypt(secret.encode(), salt=init_vector)
encryptor = Cipher(AES(key), modes.CBC(init_vector), default_backend()).encryptor() encryptor = Cipher(AES(key), modes.CBC(init_vector), default_backend()).encryptor()
padder = PKCS7(AES.block_size).padder() padder = PKCS7(AES.block_size).padder()
padded_data = padder.update(value) + padder.finalize() padded_data = padder.update(value) + padder.finalize()
encrypted_data = encryptor.update(padded_data) + encryptor.finalize() encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
return base64.b64encode(init_vector + encrypted_data) return base64.b64encode(b's:8192:16:1:' + init_vector + encrypted_data)
def better_aes_decrypt(secret: str, value: bytes) -> bytes: def better_aes_decrypt(secret: str, value: bytes) -> bytes:
data = base64.b64decode(value) data = base64.b64decode(value)
key = double_sha256(secret.encode()) type, n, r, p, data = data.split(b':', maxsplit=4)
init_vector, data = data[:16], data[16:] init_vector, data = data[:16], data[16:]
key = scrypt(secret.encode(), salt=init_vector, n=int(n), r=int(r), p=int(p))
decryptor = Cipher(AES(key), modes.CBC(init_vector), default_backend()).decryptor() decryptor = Cipher(AES(key), modes.CBC(init_vector), default_backend()).decryptor()
unpadder = PKCS7(AES.block_size).unpadder() unpadder = PKCS7(AES.block_size).unpadder()
return unpadder.update(decryptor.update(data)) + unpadder.finalize() return unpadder.update(decryptor.update(data)) + unpadder.finalize()
def scrypt(passphrase, salt, n=1<<13, r=16, p=1):
kdf = Scrypt(salt, length=32, n=n, r=r, p=p, backend=default_backend())
return kdf.derive(passphrase)
class Base58Error(Exception): class Base58Error(Exception):
""" Exception used for Base58 errors. """ """ Exception used for Base58 errors. """