lbry-sdk/lbrynet/cryptstream/CryptStreamCreator.py

156 lines
5.2 KiB
Python
Raw Normal View History

2015-08-20 17:27:15 +02:00
"""
Utility for creating Crypt Streams, which are encrypted blobs and associated metadata.
"""
2018-05-05 04:07:05 +02:00
import os
2015-08-20 17:27:15 +02:00
import logging
2018-05-05 04:07:05 +02:00
from cryptography.hazmat.primitives.ciphers.algorithms import AES
2018-07-05 05:16:52 +02:00
from twisted.internet import defer
2015-08-20 17:27:15 +02:00
from lbrynet.cryptstream.CryptBlob import CryptStreamBlobMaker
log = logging.getLogger(__name__)
2017-09-07 21:38:47 +02:00
class CryptStreamCreator(object):
"""
Create a new stream with blobs encrypted by a symmetric cipher.
2015-08-20 17:27:15 +02:00
2016-11-04 17:39:52 +01:00
Each blob is encrypted with the same key, but each blob has its
own initialization vector which is associated with the blob when
the blob is associated with the stream.
"""
2017-09-07 21:38:47 +02:00
#implements(interfaces.IConsumer)
2017-09-07 21:38:47 +02:00
2015-08-20 17:27:15 +02:00
def __init__(self, blob_manager, name=None, key=None, iv_generator=None):
2016-11-04 17:39:52 +01:00
"""@param blob_manager: Object that stores and provides access to blobs.
2015-08-20 17:27:15 +02:00
@type blob_manager: BlobManager
@param name: the name of the stream, which will be presented to the user
@type name: string
2016-11-04 17:39:52 +01:00
@param key: the raw AES key which will be used to encrypt the
blobs. If None, a random key will be generated.
2015-08-20 17:27:15 +02:00
@type key: string
2016-11-04 17:39:52 +01:00
@param iv_generator: a generator which yields initialization
vectors for the blobs. Will be called once for each blob.
2015-08-20 17:27:15 +02:00
@type iv_generator: a generator function which yields strings
@return: None
"""
self.blob_manager = blob_manager
2017-09-07 21:38:47 +02:00
self.name = name
2015-08-20 17:27:15 +02:00
self.key = key
if iv_generator is None:
self.iv_generator = self.random_iv_generator()
else:
self.iv_generator = iv_generator
2017-09-07 21:38:47 +02:00
self.stopped = True
self.producer = None
self.streaming = None
self.blob_count = -1
self.current_blob = None
self.finished_deferreds = []
def registerProducer(self, producer, streaming):
from twisted.internet import reactor
self.producer = producer
self.streaming = streaming
self.stopped = False
if streaming is False:
reactor.callLater(0, self.producer.resumeProducing)
def unregisterProducer(self):
self.stopped = True
self.producer = None
def _close_current_blob(self):
# close the blob that was being written to
# and save it to blob manager
should_announce = self.blob_count == 0
d = self.current_blob.close()
d.addCallback(self._blob_finished)
d.addCallback(lambda blob_info: self.blob_manager.creator_finished(blob_info,
should_announce))
self.finished_deferreds.append(d)
self.current_blob = None
2017-09-07 21:38:47 +02:00
def stop(self):
"""Stop creating the stream. Create the terminating zero-length blob."""
log.debug("stop has been called for StreamCreator")
self.stopped = True
if self.current_blob is not None:
self._close_current_blob()
2018-02-21 22:37:00 +01:00
d = self._finalize()
d.addCallback(lambda _: self._finished())
return d
2017-09-07 21:38:47 +02:00
# TODO: move the stream creation process to its own thread and
# remove the reactor from this process.
def write(self, data):
from twisted.internet import reactor
self._write(data)
if self.stopped is False and self.streaming is False:
reactor.callLater(0, self.producer.resumeProducing)
2015-08-20 17:27:15 +02:00
@staticmethod
def random_iv_generator():
while 1:
yield os.urandom(AES.block_size // 8)
2015-08-20 17:27:15 +02:00
def get_next_iv(self):
iv = next(self.iv_generator)
if not isinstance(iv, bytes):
return iv.encode()
return iv
2015-08-20 17:27:15 +02:00
def setup(self):
"""Create the symmetric key if it wasn't provided"""
if self.key is None:
self.key = os.urandom(AES.block_size // 8)
2015-08-20 17:27:15 +02:00
return defer.succeed(True)
2018-02-21 22:37:00 +01:00
@defer.inlineCallbacks
2015-08-20 17:27:15 +02:00
def _finalize(self):
"""
Finalize a stream by adding an empty
blob at the end, this is to indicate that
the stream has ended. This empty blob is not
saved to the blob manager
"""
2018-02-21 22:37:00 +01:00
yield defer.DeferredList(self.finished_deferreds)
2015-08-20 17:27:15 +02:00
self.blob_count += 1
iv = self.get_next_iv()
2018-02-21 22:37:00 +01:00
final_blob = self._get_blob_maker(iv, self.blob_manager.get_blob_creator())
stream_terminator = yield final_blob.close()
terminator_info = yield self._blob_finished(stream_terminator)
defer.returnValue(terminator_info)
2015-08-20 17:27:15 +02:00
def _write(self, data):
while len(data) > 0:
if self.current_blob is None:
self.next_blob_creator = self.blob_manager.get_blob_creator()
2015-08-20 17:27:15 +02:00
self.blob_count += 1
iv = self.get_next_iv()
self.current_blob = self._get_blob_maker(iv, self.next_blob_creator)
2015-08-20 17:27:15 +02:00
done, num_bytes_written = self.current_blob.write(data)
data = data[num_bytes_written:]
if done is True:
self._close_current_blob()
2015-08-20 17:27:15 +02:00
def _get_blob_maker(self, iv, blob_creator):
2016-11-04 17:39:52 +01:00
return CryptStreamBlobMaker(self.key, iv, self.blob_count, blob_creator)
2017-09-07 21:38:47 +02:00
def _finished(self):
raise NotImplementedError()
def _blob_finished(self, blob_info):
raise NotImplementedError()