pycrypt -> cryptography + remove manual padding, use lib

This commit is contained in:
Victor Shyba 2017-07-24 04:05:40 -03:00 committed by Jack Robison
parent d2fc1daf26
commit e69ba64707
No known key found for this signature in database
GPG key ID: 284699E7404E3CFF

View file

@ -1,11 +1,15 @@
import binascii import binascii
import logging import logging
from Crypto.Cipher import AES from cryptography.hazmat.primitives.ciphers import Cipher, modes
from cryptography.hazmat.primitives.ciphers.algorithms import AES
from cryptography.hazmat.primitives.padding import PKCS7
from cryptography.hazmat.backends import default_backend
from lbrynet import conf from lbrynet import conf
from lbrynet.core.BlobInfo import BlobInfo from lbrynet.core.BlobInfo import BlobInfo
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
backend = default_backend()
class CryptBlobInfo(BlobInfo): class CryptBlobInfo(BlobInfo):
@ -31,7 +35,9 @@ class StreamBlobDecryptor(object):
self.length = length self.length = length
self.buff = b'' self.buff = b''
self.len_read = 0 self.len_read = 0
self.cipher = AES.new(self.key, AES.MODE_CBC, self.iv) cipher = Cipher(AES(self.key), modes.CBC(self.iv), backend=backend)
self.unpadder = PKCS7(AES.block_size).unpadder()
self.cipher = cipher.decryptor()
def decrypt(self, write_func): def decrypt(self, write_func):
""" """
@ -42,22 +48,19 @@ class StreamBlobDecryptor(object):
""" """
def remove_padding(data): def remove_padding(data):
pad_len = ord(data[-1]) return self.unpadder.update(data) + self.unpadder.finalize()
data, padding = data[:-1 * pad_len], data[-1 * pad_len:]
for c in padding:
assert ord(c) == pad_len
return data
def write_bytes(): def write_bytes():
if self.len_read < self.length: if self.len_read < self.length:
num_bytes_to_decrypt = greatest_multiple(len(self.buff), self.cipher.block_size) num_bytes_to_decrypt = greatest_multiple(len(self.buff), (AES.block_size / 8))
data_to_decrypt, self.buff = split(self.buff, num_bytes_to_decrypt) data_to_decrypt, self.buff = split(self.buff, num_bytes_to_decrypt)
write_func(self.cipher.decrypt(data_to_decrypt)) write_func(self.cipher.update(data_to_decrypt))
def finish_decrypt(): def finish_decrypt():
assert len(self.buff) % self.cipher.block_size == 0 assert len(self.buff) % (AES.block_size / 8) == 0
data_to_decrypt, self.buff = self.buff, b'' data_to_decrypt, self.buff = self.buff, b''
write_func(remove_padding(self.cipher.decrypt(data_to_decrypt))) last_chunk = self.cipher.update(data_to_decrypt) + self.cipher.finalize()
write_func(remove_padding(last_chunk))
def decrypt_bytes(data): def decrypt_bytes(data):
self.buff += data self.buff += data
@ -84,8 +87,9 @@ class CryptStreamBlobMaker(object):
self.iv = iv self.iv = iv
self.blob_num = blob_num self.blob_num = blob_num
self.blob = blob self.blob = blob
self.cipher = AES.new(self.key, AES.MODE_CBC, self.iv) cipher = Cipher(AES(self.key), modes.CBC(self.iv), backend=backend)
self.buff = b'' self.padder = PKCS7(AES.block_size).padder()
self.cipher = cipher.encryptor()
self.length = 0 self.length = 0
def write(self, data): def write(self, data):
@ -104,10 +108,11 @@ class CryptStreamBlobMaker(object):
done = True done = True
else: else:
num_bytes_to_write = len(data) num_bytes_to_write = len(data)
self.length += num_bytes_to_write
data_to_write = data[:num_bytes_to_write] data_to_write = data[:num_bytes_to_write]
self.buff += data_to_write self.length += len(data_to_write)
self._write_buffer() padded_data = self.padder.update(data_to_write)
encrypted_data = self.cipher.update(padded_data)
self.blob.write(encrypted_data)
return done, num_bytes_to_write return done, num_bytes_to_write
def close(self): def close(self):
@ -119,20 +124,10 @@ class CryptStreamBlobMaker(object):
log.debug("called the finished_callback from CryptStreamBlobMaker.close") log.debug("called the finished_callback from CryptStreamBlobMaker.close")
return d return d
def _write_buffer(self):
num_bytes_to_encrypt = (len(self.buff) // AES.block_size) * AES.block_size
data_to_encrypt, self.buff = split(self.buff, num_bytes_to_encrypt)
encrypted_data = self.cipher.encrypt(data_to_encrypt)
self.blob.write(encrypted_data)
def _close_buffer(self): def _close_buffer(self):
data_to_encrypt, self.buff = self.buff, b'' self.length += (AES.block_size / 8) - (self.length % (AES.block_size / 8))
assert len(data_to_encrypt) < AES.block_size padded_data = self.padder.finalize()
pad_len = AES.block_size - len(data_to_encrypt) encrypted_data = self.cipher.update(padded_data) + self.cipher.finalize()
padded_data = data_to_encrypt + chr(pad_len) * pad_len
self.length += pad_len
assert len(padded_data) == AES.block_size
encrypted_data = self.cipher.encrypt(padded_data)
self.blob.write(encrypted_data) self.blob.write(encrypted_data)
def _return_info(self, blob_hash): def _return_info(self, blob_hash):