lbry-sdk/lbrynet/cryptstream/CryptBlob.py

106 lines
3.8 KiB
Python
Raw Normal View History

2015-08-20 17:27:15 +02:00
import binascii
import logging
from Crypto.Cipher import AES
from lbrynet.conf import BLOB_SIZE
from lbrynet.core.BlobInfo import BlobInfo
class CryptBlobInfo(BlobInfo):
def __init__(self, blob_hash, blob_num, length, iv):
BlobInfo.__init__(self, blob_hash, blob_num, length)
self.iv = iv
class StreamBlobDecryptor(object):
def __init__(self, blob, key, iv, length):
self.blob = blob
self.key = key
self.iv = iv
self.length = length
self.buff = b''
self.len_read = 0
self.cipher = AES.new(self.key, AES.MODE_CBC, self.iv)
def decrypt(self, write_func):
def remove_padding(data):
pad_len = ord(data[-1])
data, padding = data[:-1 * pad_len], data[-1 * pad_len:]
for c in padding:
assert ord(c) == pad_len
return data
def write_bytes():
if self.len_read < self.length:
num_bytes_to_decrypt = (len(self.buff) // self.cipher.block_size) * self.cipher.block_size
data_to_decrypt, self.buff = self.buff[:num_bytes_to_decrypt], self.buff[num_bytes_to_decrypt:]
write_func(self.cipher.decrypt(data_to_decrypt))
def finish_decrypt():
assert len(self.buff) % self.cipher.block_size == 0
data_to_decrypt, self.buff = self.buff, b''
write_func(remove_padding(self.cipher.decrypt(data_to_decrypt)))
def decrypt_bytes(data):
self.buff += data
self.len_read += len(data)
write_bytes()
#write_func(remove_padding(self.cipher.decrypt(self.buff)))
d = self.blob.read(decrypt_bytes)
d.addCallback(lambda _: finish_decrypt())
return d
class CryptStreamBlobMaker(object):
"""This class encrypts data and writes it to a new blob"""
def __init__(self, key, iv, blob_num, blob):
self.key = key
self.iv = iv
self.blob_num = blob_num
self.blob = blob
self.cipher = AES.new(self.key, AES.MODE_CBC, self.iv)
self.buff = b''
self.length = 0
def write(self, data):
max_bytes_to_write = BLOB_SIZE - self.length - 1
done = False
if max_bytes_to_write <= len(data):
num_bytes_to_write = max_bytes_to_write
done = True
else:
num_bytes_to_write = len(data)
self.length += num_bytes_to_write
data_to_write = data[:num_bytes_to_write]
self.buff += data_to_write
self._write_buffer()
return done, num_bytes_to_write
def close(self):
logging.debug("closing blob %s with plaintext len %s", str(self.blob_num), str(self.length))
if self.length != 0:
self._close_buffer()
d = self.blob.close()
d.addCallback(self._return_info)
logging.debug("called the finished_callback from CryptStreamBlobMaker.close")
return d
def _write_buffer(self):
num_bytes_to_encrypt = (len(self.buff) // AES.block_size) * AES.block_size
data_to_encrypt, self.buff = self.buff[:num_bytes_to_encrypt], self.buff[num_bytes_to_encrypt:]
encrypted_data = self.cipher.encrypt(data_to_encrypt)
self.blob.write(encrypted_data)
def _close_buffer(self):
data_to_encrypt, self.buff = self.buff, b''
assert len(data_to_encrypt) < AES.block_size
pad_len = AES.block_size - len(data_to_encrypt)
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)
def _return_info(self, blob_hash):
return CryptBlobInfo(blob_hash, self.blob_num, self.length, binascii.hexlify(self.iv))