refactor: open for different encryption methods

puzzle_generator/create_puzzle.py

@@ -1,4 +1,5 @@
 import pathlib
+import hashlib
 import inspect
 import sys
 
@@ -7,54 +8,53 @@
 from . import bytestr_utils as bu
 
 
-def _run_puzzle(in_puzzle, in_decrypt_str):
+def _run_puzzle(in_puzzle, in_decryptor):
     print(in_puzzle["str"])
     if "rest" in in_puzzle:
         this_pass = input()
         new_puzzle = decrypt_data(
-            in_puzzle["rest"], in_puzzle["hash"], this_pass, in_decrypt_str
+            in_puzzle["rest"], in_puzzle["hash"], this_pass, in_decryptor
         )
         if new_puzzle is None:
             print("This is a wrong answer. Try again!")
             sys.exit(1)
         else:
-            _run_puzzle(new_puzzle, in_decrypt_str)
+            _run_puzzle(new_puzzle, in_decryptor)
 
 
-def _create_str(in_modules, in_functions, in_encrypted_puzzle):
+def _create_str(in_modules, in_objects, in_encrypted_puzzle):
     advertisement = """# generated with puzzle-generator
 #
 # https://pypi.org/project/puzzle-generator/
 # https://github.com/vil02/puzzle_generator
 """
     modules_str = "\n".join("import " + _ for _ in in_modules) + "\n"
-    functions_str = "\n".join(inspect.getsource(_) for _ in in_functions)
+    objects_str = "\n".join(inspect.getsource(_) for _ in in_objects)
     puzzle_data_str = f"_PUZZLE = {in_encrypted_puzzle}"
-    call_str = "_run_puzzle(_PUZZLE, decrypt_str)"
+    call_str = "_run_puzzle(_PUZZLE, Decryptor(hashlib.sha512, hashlib.sha3_512))"
 
     return (
-        "\n".join(
-            [advertisement, modules_str, functions_str, puzzle_data_str, call_str]
-        )
+        "\n".join([advertisement, modules_str, objects_str, puzzle_data_str, call_str])
         + "\n"
     )
 
 
 def create(in_puzzle, output_path: pathlib.Path) -> None:
-    encrypted_puzzle = encrypt_data(in_puzzle, seu.encrypt_str)
+    encrypted_puzzle = encrypt_data(
+        in_puzzle, seu.Encryptor(hashlib.sha512, hashlib.sha3_512)
+    )
 
     needed_modules = ["hashlib", "itertools", "base64", "json", "sys"]
 
-    needed_functions = [
+    needed_objects = [
         seu.hash_bytes,
         seu.int_to_bytes,
         seu.proc_bytes,
-        seu.decrypt_bytes,
         bu.bytestr_to_bytes,
-        seu.decrypt_str,
+        seu.Decryptor,
         decrypt_data,
         _run_puzzle,
     ]
 
     with open(output_path, "w", encoding="utf-8") as res_file:
-        res_file.write(_create_str(needed_modules, needed_functions, encrypted_puzzle))
+        res_file.write(_create_str(needed_modules, needed_objects, encrypted_puzzle))

puzzle_generator/puzzle_data_encryption.py

@@ -1,16 +1,16 @@
 import json
 
 
-def encrypt_data(in_data, in_encrypt):
+def encrypt_data(in_data, in_encryptor):
     if list(in_data.keys()) == ["str"]:
         return {"str": in_data["str"]}
-    rest_str = json.dumps(encrypt_data(in_data["rest"], in_encrypt))
-    encrypted_str, hash_str = in_encrypt(rest_str, in_data["pass"])
+    rest_str = json.dumps(encrypt_data(in_data["rest"], in_encryptor))
+    encrypted_str, hash_str = in_encryptor(rest_str, in_data["pass"])
     return {"str": in_data["str"], "rest": encrypted_str, "hash": hash_str}
 
 
-def decrypt_data(in_rest, in_hash, in_pass, in_decrypt):
-    rest_str = in_decrypt(in_rest, in_pass, in_hash)
+def decrypt_data(in_rest, in_hash, in_pass, in_decryptor):
+    rest_str = in_decryptor(in_rest, in_pass, in_hash)
     if rest_str is None:
         return None
     return json.loads(rest_str)

puzzle_generator/simple_encryption_utils.py

@@ -1,47 +1,61 @@
 import itertools
 import typing
-import hashlib
 
 from .bytestr_utils import bytes_to_bytestr, bytestr_to_bytes
 
 
-def hash_bytes(in_bytes: bytes) -> str:
-    return hashlib.sha512(in_bytes).hexdigest()
+def hash_bytes(in_bytes: bytes, in_hasher) -> str:
+    return in_hasher(in_bytes).hexdigest()
 
 
 def int_to_bytes(in_val: int) -> bytes:
     return in_val.to_bytes((in_val.bit_length() + 7) // 8, "big")
 
 
-def proc_bytes(in_bytes: bytes, in_key: bytes) -> bytes:
+def proc_bytes(in_bytes: bytes, in_key: bytes, in_hasher) -> bytes:
     """xors the in_bytes with a sequence of bytes generated with in_key"""
     key_bytes = itertools.chain.from_iterable(
-        hashlib.sha512(in_key + int_to_bytes(block_num)).digest()
+        in_hasher(in_key + int_to_bytes(block_num)).digest()
         for block_num in itertools.count(0)
     )
     return bytes(_d ^ _k for (_d, _k) in zip(in_bytes, key_bytes))
 
 
-def encrypt_bytes(in_bytes: bytes, in_pass: bytes) -> typing.Tuple[bytes, str]:
-    return proc_bytes(in_bytes, in_pass), hash_bytes(in_bytes)
+class Encryptor:
+    def __init__(self, proc_hasher, signature_hasher):
+        self._proc_hasher = proc_hasher
+        self._signature_hasher = signature_hasher
 
+    def __call__(self, in_str: str, in_pass: str) -> typing.Tuple[str, str]:
+        encrypted_bytes, hash_str = self._encrypt_bytes(
+            in_str.encode(), in_pass.encode()
+        )
+        return bytes_to_bytestr(encrypted_bytes), hash_str
 
-def decrypt_bytes(in_bytes: bytes, in_pass: bytes, in_hash: str) -> bytes | None:
-    res = proc_bytes(in_bytes, in_pass)
+    def _encrypt_bytes(
+        self, in_bytes: bytes, in_pass: bytes
+    ) -> typing.Tuple[bytes, str]:
+        return proc_bytes(in_bytes, in_pass, self._proc_hasher), hash_bytes(
+            in_bytes, self._signature_hasher
+        )
 
-    if hash_bytes(res) == in_hash:
-        return res
-    return None
 
+class Decryptor:
+    def __init__(self, proc_hasher, signature_hasher):
+        self._proc_hasher = proc_hasher
+        self._signature_hasher = signature_hasher
 
-def encrypt_str(in_str: str, in_pass: str) -> typing.Tuple[str, str]:
-    """encrypts in_str using in_pass"""
-    encrypted_bytes, hash_str = encrypt_bytes(in_str.encode(), in_pass.encode())
-    return bytes_to_bytestr(encrypted_bytes), hash_str
+    def __call__(self, in_str: str, in_pass: str, in_hash: str) -> str | None:
+        res = self._decrypt_bytes(bytestr_to_bytes(in_str), in_pass.encode(), in_hash)
+        if res is not None:
+            return res.decode()
+        return res
 
+    def _decrypt_bytes(
+        self, in_bytes: bytes, in_pass: bytes, in_hash: str
+    ) -> bytes | None:
+        res = proc_bytes(in_bytes, in_pass, self._proc_hasher)
 
-def decrypt_str(in_str: str, in_pass: str, in_hash: str) -> str | None:
-    res = decrypt_bytes(bytestr_to_bytes(in_str), in_pass.encode(), in_hash)
-    if res is not None:
-        return res.decode()
-    return res
+        if hash_bytes(res, self._signature_hasher) == in_hash:
+            return res
+        return None

tests/test_encryption_utils.py

@@ -1,36 +1,52 @@
 import string
+import hashlib
+import itertools
 import pytest
 
 import puzzle_generator.simple_encryption_utils as seu
 
 
 _STRS = [
     "",
-    "some_str",
-    "other_strstr!?#",
+    "some_STR?!",
     string.printable,
-    string.ascii_uppercase,
     string.whitespace,
     "ąęćśłń󿟥ĘĆŚŁŃÓŻŹ",
     "some_msg_with 🔨 and 🛷!",
-    "a🎄b",
-    "🎮🎈🥅🐾",
+    "🎮🎈🥅🐾 a🎄b",
     "🏀",
 ]
 
+_SOME_HASHES = [
+    hashlib.md5,
+    hashlib.sha512,
+    hashlib.sha3_512,
+    hashlib.blake2s,
+]
+
+
+def _get_encryptor_decryptor_pair(proc_hasher, signature_hasher):
+    return seu.Encryptor(proc_hasher, signature_hasher), seu.Decryptor(
+        proc_hasher, signature_hasher
+    )
+
 
 @pytest.mark.parametrize("in_str", _STRS)
 @pytest.mark.parametrize("in_pass", _STRS)
 @pytest.mark.parametrize(
-    ("in_encrypt_str", "in_decrypt_str"), [(seu.encrypt_str, seu.decrypt_str)]
+    ("encryptor", "decryptor"),
+    [
+        _get_encryptor_decryptor_pair(*_)
+        for _ in itertools.product(_SOME_HASHES, repeat=2)
+    ],
 )
-def test_seu(in_str, in_pass, in_encrypt_str, in_decrypt_str):
-    encrypted, reshash = in_encrypt_str(in_str, in_pass)
+def test_seu(in_str, in_pass, encryptor, decryptor):
+    encrypted, reshash = encryptor(in_str, in_pass)
     if in_str:
         assert encrypted != in_str
     else:
         assert not encrypted
-    decrypted = in_decrypt_str(encrypted, in_pass, reshash)
+    decrypted = decryptor(encrypted, in_pass, reshash)
     assert decrypted == in_str
     if in_str:
-        assert in_decrypt_str(encrypted, in_pass + "?", reshash) is None
+        assert decryptor(encrypted, in_pass + "?", reshash) is None

tests/test_puzzle_data_encryption.py

@@ -1,8 +1,21 @@
+import hashlib
+import itertools
 import pytest
 
 import puzzle_generator.puzzle_data_encryption as pde
 import puzzle_generator.simple_encryption_utils as seu
 
+_SOME_HASHES = [
+    hashlib.sha1,
+    hashlib.sha256,
+]
+
+
+def _get_encryptor_decryptor_pair(proc_hasher, signature_hasher):
+    return seu.Encryptor(proc_hasher, signature_hasher), seu.Decryptor(
+        proc_hasher, signature_hasher
+    )
+
 
 @pytest.mark.parametrize(
     "in_puzzle",
@@ -46,8 +59,15 @@
         },
     ],
 )
-def test_pde(in_puzzle):
-    encrypted_puzzle = pde.encrypt_data(in_puzzle, seu.encrypt_str)
+@pytest.mark.parametrize(
+    ("encryptor", "decryptor"),
+    [
+        _get_encryptor_decryptor_pair(*_)
+        for _ in itertools.product(_SOME_HASHES, repeat=2)
+    ],
+)
+def test_pde(in_puzzle, encryptor, decryptor):
+    encrypted_puzzle = pde.encrypt_data(in_puzzle, encryptor)
     tmp_puzzle_data = in_puzzle
     while "rest" in encrypted_puzzle:
         cur_pass = tmp_puzzle_data["pass"]
@@ -57,15 +77,15 @@ def test_pde(in_puzzle):
                 encrypted_puzzle["rest"],
                 encrypted_puzzle["hash"],
                 cur_pass + "!",
-                seu.decrypt_str,
+                decryptor,
             )
             is None
         )
         encrypted_puzzle = pde.decrypt_data(
             encrypted_puzzle["rest"],
             encrypted_puzzle["hash"],
             cur_pass,
-            seu.decrypt_str,
+            decryptor,
         )
         tmp_puzzle_data = tmp_puzzle_data["rest"]
     assert encrypted_puzzle == tmp_puzzle_data