From f04eb3c14bffb473e04c44287eb069986468ab63 Mon Sep 17 00:00:00 2001
From: Alexandr Mansurov <alex@eghuro.cz>
Date: Sat, 9 Oct 2021 14:53:28 +0200
Subject: [PATCH] Docs, tests, typing

---
 src/cryptbase/cryptbase.py | 55 ++++++++++++++++++++++++++++++++------
 tests/test_cryptbase.py    | 19 +++++++++++++
 2 files changed, 66 insertions(+), 8 deletions(-)

diff --git a/src/cryptbase/cryptbase.py b/src/cryptbase/cryptbase.py
index 1fb5832..af4ff62 100644
--- a/src/cryptbase/cryptbase.py
+++ b/src/cryptbase/cryptbase.py
@@ -1,38 +1,69 @@
 import base64
 import secrets
+from abc import ABC
+from abc import abstractmethod
+from typing import Tuple
 
 from cryptography.hazmat.primitives.ciphers import Cipher
 from cryptography.hazmat.primitives.ciphers import algorithms
 from cryptography.hazmat.primitives.ciphers import modes
 
 
-class AES256CTREncryptor:
+class Encryptor(ABC):
+    """Abstract base class defining interface for encryptors."""
 
-    def __init__(self, key):
+    @abstractmethod
+    def encrypt(self, plaintext: bytes) -> Tuple[bytes, bytes]:
+        """Transform plaintext into a pair of ciphertext and IV."""
+        pass
+
+    @abstractmethod
+    def decrypt(self, ciphertext: bytes, iv: bytes) -> bytes:
+        """Transform ciphertext and IV into plaintext."""
+        pass
+
+
+class AES256CTREncryptor(Encryptor):
+    """Encryptor using AES with 256-bit long keys using CTR mode."""
+
+    def __init__(self, key: bytes):
         if not isinstance(key, bytes):
             raise ValueError()
         if not len(key) == 32:
             raise ValueError()
         self.__key = key
 
-    def encrypt(self, plaintext):
+    def encrypt(self, plaintext: bytes) -> Tuple[bytes, bytes]:
         iv = secrets.token_bytes(16)
         cipher = Cipher(algorithms.AES(self.__key), modes.CTR(iv))
         encryptor = cipher.encryptor()
         ciphertext = encryptor.update(plaintext) + encryptor.finalize()
         return ciphertext, iv
 
-    def decrypt(self, ciphertext, iv):
+    def decrypt(self, ciphertext: bytes, iv: bytes) -> bytes:
         cipher = Cipher(algorithms.AES(self.__key), modes.CTR(iv))
         decryptor = cipher.decryptor()
         return decryptor.update(ciphertext) + decryptor.finalize()
 
 
 class CryptoContainer:
+    """
+    CryptoContainer is providing unified interface to encryptors.
+    Enter the encryptor instance and either plaintext string or
+    ciphertext and iv bytes in constructor. Then access plaintext,
+    ciphertext and IV strings in properties.
+    Ciphertext and IV strings properties yield base64 encoded strings.
+    Converting this container to string yields a concatenation of
+    ciphertext and IV in a single string. Calling from_encrypted
+    static method will construct CryptoContainer from such concatenation,
+    allowing access to plaintext.
+    """
 
     def __init__(self, **kwargs):
         if 'encryptor' not in kwargs:
             raise ValueError()
+        if not isinstance(kwargs['encryptor'], Encryptor):
+            raise ValueError()
 
         if 'plaintext' in kwargs:
             self.__plaintext = kwargs['plaintext']
@@ -45,22 +76,30 @@ class CryptoContainer:
             raise ValueError()
 
     @property
-    def plaintext(self):
+    def plaintext(self) -> str:
+        """Provide plaintext."""
         return self.__plaintext
 
     @property
-    def ciphertext(self):
+    def ciphertext(self) -> str:
+        """Provide base64 encoded ciphertext as string."""
         return base64.b64encode(self.__ciphertext).decode('ascii')
 
     @property
-    def iv(self):
+    def iv(self) -> str:
+        """Provide base64 encoded IV as string."""
         return base64.b64encode(self.__iv).decode('ascii')
 
     def __str__(self):
+        """Concatenate ciphertext and IV to store in the database."""
         return '@'.join([self.ciphertext, self.iv])
 
     @staticmethod
-    def from_encrypted(container_str, encryptor):
+    def from_encrypted(container_str: str, encryptor: Encryptor):
+        """
+        Builds a CryptoContainer from a string as provided by __str__,
+        allowing access to plaintext.
+        """
         ciphertext, iv = container_str.split('@')
         return CryptoContainer(ciphertext=base64.b64decode(ciphertext),
                                iv=base64.b64decode(iv),
diff --git a/tests/test_cryptbase.py b/tests/test_cryptbase.py
index 72e22db..32fbc16 100644
--- a/tests/test_cryptbase.py
+++ b/tests/test_cryptbase.py
@@ -94,6 +94,19 @@ def test_container_init():
     try:
         cryptbase.cryptbase.CryptoContainer(encryptor=None, plaintext=None)
         assert False
+    except ValueError:
+        pass
+
+    class DummyEncryptor(cryptbase.cryptbase.Encryptor):
+        def encrypt(self, plaintext):
+            return None, None
+
+        def decrypt(self, ct, iv):
+            return None
+
+    try:
+        cryptbase.cryptbase.CryptoContainer(encryptor=DummyEncryptor(), plaintext=None)
+        assert False
     except AttributeError:
         pass
 
@@ -112,5 +125,11 @@ def test_container_init():
     try:
         cryptbase.cryptbase.CryptoContainer(encryptor=None, ciphertext=None, iv=None)
         assert False
+    except ValueError:
+        pass
+
+    try:
+        cryptbase.cryptbase.CryptoContainer(encryptor=DummyEncryptor(), ciphertext=None, iv=None)
+        assert False
     except AttributeError:
         pass