lbry-sdk/lbrynet/cryptstream/CryptStreamCreator.py

"""
Utility for creating Crypt Streams, which are encrypted blobs and associated metadata.
"""

import logging

from Crypto import Random
from Crypto.Cipher import AES

from twisted.internet import defer
from lbrynet.core.StreamCreator import StreamCreator
from lbrynet.cryptstream.CryptBlob import CryptStreamBlobMaker


log = logging.getLogger(__name__)


class CryptStreamCreator(StreamCreator):
    """Create a new stream with blobs encrypted by a symmetric cipher.

    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."""
    def __init__(self, blob_manager, name=None, key=None, iv_generator=None):
        """
        @param blob_manager: Object that stores and provides access to blobs.
        @type blob_manager: BlobManager

        @param name: the name of the stream, which will be presented to the user
        @type name: string

        @param key: the raw AES key which will be used to encrypt the blobs. If None, a random key will
            be generated.
        @type key: string

        @param iv_generator: a generator which yields initialization vectors for the blobs. Will be called
            once for each blob.
        @type iv_generator: a generator function which yields strings

        @return: None
        """
        StreamCreator.__init__(self, name)
        self.blob_manager = blob_manager
        self.key = key
        if iv_generator is None:
            self.iv_generator = self.random_iv_generator()
        else:
            self.iv_generator = iv_generator

    @staticmethod
    def random_iv_generator():
        while 1:
            yield Random.new().read(AES.block_size)

    def setup(self):
        """Create the symmetric key if it wasn't provided"""

        if self.key is None:
            self.key = Random.new().read(AES.block_size)

        return defer.succeed(True)

    def _finalize(self):
        log.debug("_finalize has been called")
        self.blob_count += 1
        iv = self.iv_generator.next()
        final_blob_creator = self.blob_manager.get_blob_creator()
        log.debug("Created the finished_deferred")
        final_blob = self._get_blob_maker(iv, final_blob_creator)
        log.debug("Created the final blob")
        log.debug("Calling close on final blob")
        d = final_blob.close()
        d.addCallback(self._blob_finished)
        self.finished_deferreds.append(d)
        log.debug("called close on final blob, returning from make_final_blob")
        return d

    def _write(self, data):

        def close_blob(blob):
            d = blob.close()
            d.addCallback(self._blob_finished)
            self.finished_deferreds.append(d)

        while len(data) > 0:
            if self.current_blob is None:
                next_blob_creator = self.blob_manager.get_blob_creator()
                self.blob_count += 1
                iv = self.iv_generator.next()
                self.current_blob = self._get_blob_maker(iv, next_blob_creator)
            done, num_bytes_written = self.current_blob.write(data)
            data = data[num_bytes_written:]
            if done is True:
                close_blob(self.current_blob)
                self.current_blob = None

    def _get_blob_maker(self, iv, blob_creator):
        return CryptStreamBlobMaker(self.key, iv, self.blob_count, blob_creator)