From 46fdda871f3ee5edf841d86de385de4d2e828db9 Mon Sep 17 00:00:00 2001
From: Alexandr Mansurov <alex@eghuro.cz>
Date: Sat, 9 Oct 2021 11:00:46 +0200
Subject: [PATCH] Cryptbase tests

---
 src/cryptbase/cryptbase.py |   8 ++-
 tests/test_cryptbase.py    | 110 +++++++++++++++++++++++++++++++++++++
 2 files changed, 116 insertions(+), 2 deletions(-)

diff --git a/src/cryptbase/cryptbase.py b/src/cryptbase/cryptbase.py
index 5973143..1fb5832 100644
--- a/src/cryptbase/cryptbase.py
+++ b/src/cryptbase/cryptbase.py
@@ -9,6 +9,10 @@ from cryptography.hazmat.primitives.ciphers import modes
 class AES256CTREncryptor:
 
     def __init__(self, key):
+        if not isinstance(key, bytes):
+            raise ValueError()
+        if not len(key) == 32:
+            raise ValueError()
         self.__key = key
 
     def encrypt(self, plaintext):
@@ -36,13 +40,13 @@ class CryptoContainer:
         elif 'ciphertext' in kwargs and 'iv' in kwargs:
             self.__ciphertext = kwargs['ciphertext']
             self.__iv = kwargs['iv']
-            self.__plaintext = kwargs['encryptor'].decrypt(self.__ciphertext, self.__iv)
+            self.__plaintext = kwargs['encryptor'].decrypt(self.__ciphertext, self.__iv).decode('utf-8')
         else:
             raise ValueError()
 
     @property
     def plaintext(self):
-        return self.__plaintext.decode('utf-8')
+        return self.__plaintext
 
     @property
     def ciphertext(self):
diff --git a/tests/test_cryptbase.py b/tests/test_cryptbase.py
index 1a1030a..72e22db 100644
--- a/tests/test_cryptbase.py
+++ b/tests/test_cryptbase.py
@@ -1,6 +1,116 @@
 
+import base64
+import random
+import secrets
+import string
+
 import cryptbase
 
 
 def test_main():
     assert cryptbase.__version__ == '0.0.0'
+
+
+def test_crypto():
+    aes_n = 256
+    key_bytes = int(aes_n / 8)
+    population = string.ascii_letters + string.digits + string.punctuation
+    length = 2048
+
+    for _ in range(16):
+        key = secrets.token_bytes(key_bytes)
+        encryptor = cryptbase.cryptbase.AES256CTREncryptor(key)
+        for _ in range(32):
+            plaintext = ''.join(random.choices(population, k=length)).encode('ascii')
+            assert plaintext == encryptor.decrypt(*encryptor.encrypt(plaintext))
+
+
+def test_encryptor():
+    try:
+        cryptbase.cryptbase.AES256CTREncryptor(None)
+        assert False
+    except ValueError:
+        pass
+
+    try:
+        cryptbase.cryptbase.AES256CTREncryptor(bytes([]))
+        assert False
+    except ValueError:
+        pass
+
+    try:
+        key = secrets.token_bytes(1)
+        cryptbase.cryptbase.AES256CTREncryptor(key)
+        assert False
+    except ValueError:
+        pass
+
+
+def test_containter_encrypt():
+    aes_n = 256
+    key_bytes = int(aes_n / 8)
+    population = string.ascii_letters + string.digits + string.punctuation
+    length = 2048
+
+    for _ in range(16):
+        key = secrets.token_bytes(key_bytes)
+        encryptor = cryptbase.cryptbase.AES256CTREncryptor(key)
+        for _ in range(32):
+            plaintext = ''.join(random.choices(population, k=length))
+            container = cryptbase.cryptbase.CryptoContainer(encryptor=encryptor, plaintext=plaintext)
+
+            assert container.plaintext == plaintext
+            assert str(container) == '@'.join([container.ciphertext, container.iv])
+
+            cipher_b64 = container.ciphertext
+            iv_b64 = container.iv
+
+            ciphertext = base64.b64decode(cipher_b64)
+            iv = base64.b64decode(iv_b64)
+            decrypted = encryptor.decrypt(ciphertext, iv).decode('utf-8')
+            assert type(plaintext) == type(decrypted)
+            assert len(plaintext) == len(decrypted)
+            assert plaintext == decrypted
+
+            decontainer = cryptbase.cryptbase.CryptoContainer.from_encrypted(str(container), encryptor)
+            assert decontainer.plaintext == plaintext
+            assert decontainer.iv == iv_b64
+            assert decontainer.ciphertext == cipher_b64
+
+
+def test_container_init():
+    try:
+        cryptbase.cryptbase.CryptoContainer()
+        assert False
+    except ValueError:
+        pass
+
+    try:
+        cryptbase.cryptbase.CryptoContainer(encryptor=None)
+        assert False
+    except ValueError:
+        pass
+
+    try:
+        cryptbase.cryptbase.CryptoContainer(encryptor=None, plaintext=None)
+        assert False
+    except AttributeError:
+        pass
+
+    try:
+        cryptbase.cryptbase.CryptoContainer(encryptor=None, ciphertext=None)
+        assert False
+    except ValueError:
+        pass
+
+    try:
+        cryptbase.cryptbase.CryptoContainer(encryptor=None, iv=None)
+        assert False
+    except ValueError:
+        pass
+
+    try:
+        cryptbase.cryptbase.CryptoContainer(encryptor=None, ciphertext=None, iv=None)
+        assert False
+    except AttributeError:
+        pass