lbry-sdk/tests/unit/cryptstream/test_cryptblob.py

from twisted.trial import unittest
from twisted.internet import defer
from lbrynet.cryptstream import CryptBlob
from lbrynet.blob.blob_file import MAX_BLOB_SIZE

from tests.mocks import mock_conf_settings

from cryptography.hazmat.primitives.ciphers.algorithms import AES
import random
import string
from six import BytesIO
import os

AES_BLOCK_SIZE_BYTES = int(AES.block_size / 8)

class MocBlob(object):
    def __init__(self):
        self.data = b''

    def read(self, write_func):
        data = self.data
        write_func(data)
        return defer.succeed(True)

    def open_for_reading(self):
        return BytesIO(self.data)

    def write(self, data):
        if not isinstance(data, bytes):
            data = data.encode()
        self.data += data

    def close(self):
        return defer.succeed(True)


def random_string(length):
    return ''.join(random.choice(string.ascii_lowercase) for i in range(length))


class TestCryptBlob(unittest.TestCase):
    def setUp(self):
        mock_conf_settings(self)


    @defer.inlineCallbacks
    def _test_encrypt_decrypt(self, size_of_data):
        # max blob size is 2*2**20 -1 ( -1 due to required padding in the end )
        blob = MocBlob()
        blob_num = 0
        key = os.urandom(AES_BLOCK_SIZE_BYTES)
        iv = os.urandom(AES_BLOCK_SIZE_BYTES)
        maker = CryptBlob.CryptStreamBlobMaker(key, iv, blob_num, blob)
        write_size = size_of_data
        string_to_encrypt = random_string(size_of_data).encode()

        # encrypt string
        done, num_bytes = maker.write(string_to_encrypt)
        yield maker.close()
        self.assertEqual(size_of_data, num_bytes)
        expected_encrypted_blob_size = int((size_of_data / AES_BLOCK_SIZE_BYTES) + 1) * AES_BLOCK_SIZE_BYTES
        self.assertEqual(expected_encrypted_blob_size, len(blob.data))

        if size_of_data < MAX_BLOB_SIZE-1:
            self.assertFalse(done)
        else:
            self.assertTrue(done)
        self.data_buf = b''

        def write_func(data):
            self.data_buf += data

        # decrypt string
        decryptor = CryptBlob.StreamBlobDecryptor(blob, key, iv, size_of_data)
        yield decryptor.decrypt(write_func)
        self.assertEqual(self.data_buf, string_to_encrypt)

    @defer.inlineCallbacks
    def test_encrypt_decrypt(self):
        yield self._test_encrypt_decrypt(1)
        yield self._test_encrypt_decrypt(16*2)
        yield self._test_encrypt_decrypt(2000)
        yield self._test_encrypt_decrypt(2*2**20-1)