From 414f5ebc1b975619bc81dd4ac8641d54e0c1e6ae Mon Sep 17 00:00:00 2001 From: vil02 <65706193+vil02@users.noreply.github.com> Date: Wed, 24 Jul 2024 22:51:59 +0200 Subject: [PATCH] refactor: open for different encryption methods --- puzzle_generator/create_puzzle.py | 51 +++++++++++++++------ puzzle_generator/simple_encryption_utils.py | 50 +++++++++++--------- tests/test_create_puzzle.py | 26 ++++++++--- tests/test_encryption_utils.py | 33 +++++++++---- tests/test_puzzle_data_encryption.py | 25 ++++++++-- 5 files changed, 129 insertions(+), 56 deletions(-) diff --git a/puzzle_generator/create_puzzle.py b/puzzle_generator/create_puzzle.py index 734dbc7..44748fd 100644 --- a/puzzle_generator/create_puzzle.py +++ b/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) + ) diff --git a/puzzle_generator/simple_encryption_utils.py b/puzzle_generator/simple_encryption_utils.py index dfde71a..83ed320 100644 --- a/puzzle_generator/simple_encryption_utils.py +++ b/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 diff --git a/tests/test_create_puzzle.py b/tests/test_create_puzzle.py index 3395b2c..d92b3b8 100644 --- a/tests/test_create_puzzle.py +++ b/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,8 +40,17 @@ 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 @@ -48,8 +58,9 @@ def test_all_good_answers(puzzle, puzzle_path: pathlib.Path) -> None: 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" diff --git a/tests/test_encryption_utils.py b/tests/test_encryption_utils.py index 6f8cdc4..54494fb 100644 --- a/tests/test_encryption_utils.py +++ b/tests/test_encryption_utils.py @@ -1,4 +1,6 @@ import string +import hashlib +import itertools import pytest import puzzle_generator.simple_encryption_utils as seu @@ -6,31 +8,42 @@ _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 diff --git a/tests/test_puzzle_data_encryption.py b/tests/test_puzzle_data_encryption.py index 422e428..3e17e33 100644 --- a/tests/test_puzzle_data_encryption.py +++ b/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,7 +74,7 @@ def test_pde(in_puzzle): encrypted_puzzle["rest"], encrypted_puzzle["hash"], cur_pass + "!", - seu.decrypt_str, + decrypt, ) is None ) @@ -65,7 +82,7 @@ def test_pde(in_puzzle): 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