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,76 @@
 from . import bytestr_utils as bu
 
 
-def _run_puzzle(in_puzzle, in_decrypt_str):
+def _run_puzzle(in_puzzle, in_decrypt):
     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_decrypt
         )
         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_decrypt)
 
 
-def _create_str(in_modules, in_functions, in_encrypted_puzzle):
+def _create_str(in_modules, in_objects, in_encrypted_puzzle, decrypt_str: str) -> str:
     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, _DECRYPT)"
 
     return (
         "\n".join(
-            [advertisement, modules_str, functions_str, puzzle_data_str, call_str]
+            [
+                advertisement,
+                modules_str,
+                objects_str,
+                puzzle_data_str,
+                decrypt_str,
+                call_str,
+            ]
         )
         + "\n"
     )
 
 
-def create(in_puzzle, output_path: pathlib.Path) -> None:
-    encrypted_puzzle = encrypt_data(in_puzzle, seu.encrypt_str)
+def _get_hasher_name(in_hasher) -> str:
+    return "hashlib." + in_hasher().name
 
-    needed_modules = ["hashlib", "itertools", "base64", "json", "sys"]
 
-    needed_functions = [
+def create(in_puzzle, output_path: pathlib.Path, **kwargs) -> None:
+    proc_hasher = kwargs.get("proc_hasher", hashlib.sha512)
+    signature_hasher = kwargs.get("signature_hasher", hashlib.sha512)
+    encrypted_puzzle = encrypt_data(
+        in_puzzle, seu.get_encrypt(proc_hasher, signature_hasher)
+    )
+
+    needed_modules = ["hashlib", "itertools", "base64", "json", "sys", "typing"]
+
+    needed_objects = [
         seu.hash_bytes,
         seu.int_to_bytes,
         seu.proc_bytes,
-        seu.decrypt_bytes,
         bu.bytestr_to_bytes,
-        seu.decrypt_str,
+        seu.get_decrypt,
         decrypt_data,
         _run_puzzle,
     ]
 
+    decrypt_str = (
+        "_DECRYPT = get_decrypt("
+        f"{_get_hasher_name(proc_hasher)}, "
+        f"{_get_hasher_name(signature_hasher)})"
+    )
+
     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, decrypt_str)
+        )

puzzle_generator/simple_encryption_utils.py

@@ -1,47 +1,55 @@
 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 str(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)
+def get_encrypt(
+    proc_hasher, signature_hasher
+) -> typing.Callable[[str, str], typing.Tuple[str, str]]:
+    def _encrypt_bytes(in_bytes: bytes, in_pass: bytes) -> typing.Tuple[bytes, str]:
+        return proc_bytes(in_bytes, in_pass, proc_hasher), hash_bytes(
+            in_bytes, signature_hasher
+        )
 
+    def _encrypt(in_str: str, in_pass: str) -> typing.Tuple[str, str]:
+        encrypted_bytes, hash_str = _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)
+    return _encrypt
 
-    if hash_bytes(res) == in_hash:
-        return res
-    return None
 
+def get_decrypt(
+    proc_hasher, signature_hasher
+) -> typing.Callable[[str, str, str], str | None]:
+    def _decrypt_bytes(in_bytes: bytes, in_pass: bytes, in_hash: str) -> bytes | None:
+        res = proc_bytes(in_bytes, in_pass, proc_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
+        if hash_bytes(res, signature_hasher) == in_hash:
+            return res
+        return None
 
+    def _decrypt(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
 
-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
+    return _decrypt

tests/test_create_puzzle.py

@@ -1,4 +1,5 @@
 import pathlib
+import hashlib
 import subprocess
 import pytest
 
@@ -19,7 +20,7 @@ def fixture_puzzle():
 
 
 @pytest.fixture(name="puzzle_path")
-def fixture_puzzle_path(tmp_path) -> pathlib.Path:
+def fixture_puzzle_path(tmp_path: pathlib.Path) -> pathlib.Path:
     return tmp_path / "puzzle.py"
 
 
@@ -39,17 +40,27 @@ def _run_puzzle_file(
     )
 
 
-def test_all_good_answers(puzzle, puzzle_path: pathlib.Path) -> None:
-    cp.create(puzzle, puzzle_path)
+_CONFIGURATIONS = [
+    {},
+    {"proc_hasher": hashlib.md5},
+    {"proc_hasher": hashlib.sha1, "signature_hasher": hashlib.sha3_224},
+    {"signature_hasher": hashlib.blake2b},
+]
+
+
+@pytest.mark.parametrize("configuration", _CONFIGURATIONS)
+def test_all_good_answers(puzzle, puzzle_path: pathlib.Path, configuration) -> None:
+    cp.create(puzzle, puzzle_path, **configuration)
     res = _run_puzzle_file(puzzle_path, ["Answer 1", "Is this the final answer?"])
 
     assert res.returncode == 0
     assert res.stdout == "Question 1?\nQuestion 2?\nCongratulations!\n"
     assert not res.stderr
 
 
-def test_second_answer_wrong(puzzle, puzzle_path: pathlib.Path) -> None:
-    cp.create(puzzle, puzzle_path)
+@pytest.mark.parametrize("configuration", _CONFIGURATIONS)
+def test_second_answer_wrong(puzzle, puzzle_path: pathlib.Path, configuration) -> None:
+    cp.create(puzzle, puzzle_path, **configuration)
     res = _run_puzzle_file(puzzle_path, ["Answer 1", "This is a wrong answer"])
     assert res.returncode == 1
     assert (
@@ -58,8 +69,9 @@ def test_second_answer_wrong(puzzle, puzzle_path: pathlib.Path) -> None:
     assert not res.stderr
 
 
-def test_first_answer_wrong(puzzle, puzzle_path: pathlib.Path) -> None:
-    cp.create(puzzle, puzzle_path)
+@pytest.mark.parametrize("configuration", _CONFIGURATIONS)
+def test_first_answer_wrong(puzzle, puzzle_path: pathlib.Path, configuration) -> None:
+    cp.create(puzzle, puzzle_path, **configuration)
     res = _run_puzzle_file(puzzle_path, ["This is a wrong answer."])
     assert res.returncode == 1
     assert res.stdout == "Question 1?\nThis is a wrong answer. Try again!\n"

tests/test_encryption_utils.py

@@ -1,36 +1,49 @@
 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_encrypt_decrypt_pair(proc_hasher, signature_hasher):
+    return seu.get_encrypt(proc_hasher, signature_hasher), seu.get_decrypt(
+        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)]
+    ("encrypt", "decrypt"),
+    [_get_encrypt_decrypt_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, encrypt, decrypt):
+    encrypted, reshash = encrypt(in_str, in_pass)
     if in_str:
         assert encrypted != in_str
     else:
         assert not encrypted
-    decrypted = in_decrypt_str(encrypted, in_pass, reshash)
+    decrypted = decrypt(encrypted, in_pass, reshash)
     assert decrypted == in_str
     if in_str:
-        assert in_decrypt_str(encrypted, in_pass + "?", reshash) is None
+        assert decrypt(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_encrypt_decrypt_pair(proc_hasher, signature_hasher):
+    return seu.get_encrypt(proc_hasher, signature_hasher), seu.get_decrypt(
+        proc_hasher, signature_hasher
+    )
+
 
 @pytest.mark.parametrize(
     "in_puzzle",
@@ -46,8 +59,12 @@
         },
     ],
 )
-def test_pde(in_puzzle):
-    encrypted_puzzle = pde.encrypt_data(in_puzzle, seu.encrypt_str)
+@pytest.mark.parametrize(
+    ("encrypt", "decrypt"),
+    [_get_encrypt_decrypt_pair(*_) for _ in itertools.product(_SOME_HASHES, repeat=2)],
+)
+def test_pde(in_puzzle, encrypt, decrypt):
+    encrypted_puzzle = pde.encrypt_data(in_puzzle, encrypt)
     tmp_puzzle_data = in_puzzle
     while "rest" in encrypted_puzzle:
         cur_pass = tmp_puzzle_data["pass"]
@@ -57,15 +74,15 @@ def test_pde(in_puzzle):
                 encrypted_puzzle["rest"],
                 encrypted_puzzle["hash"],
                 cur_pass + "!",
-                seu.decrypt_str,
+                decrypt,
             )
             is None
         )
         encrypted_puzzle = pde.decrypt_data(
             encrypted_puzzle["rest"],
             encrypted_puzzle["hash"],
             cur_pass,
-            seu.decrypt_str,
+            decrypt,
         )
         tmp_puzzle_data = tmp_puzzle_data["rest"]
     assert encrypted_puzzle == tmp_puzzle_data