From 46fd777b10aa622da77ae69a92bea807d6a75d07 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 | 28 +++++------ puzzle_generator/puzzle_data_encryption.py | 10 ++-- puzzle_generator/simple_encryption_utils.py | 56 +++++++++++++-------- tests/test_encryption_utils.py | 36 +++++++++---- tests/test_puzzle_data_encryption.py | 28 +++++++++-- 5 files changed, 104 insertions(+), 54 deletions(-) diff --git a/puzzle_generator/create_puzzle.py b/puzzle_generator/create_puzzle.py index 734dbc7..7b2290d 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,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)) diff --git a/puzzle_generator/puzzle_data_encryption.py b/puzzle_generator/puzzle_data_encryption.py index 99d217d..cb18c6e 100644 --- a/puzzle_generator/puzzle_data_encryption.py +++ b/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) diff --git a/puzzle_generator/simple_encryption_utils.py b/puzzle_generator/simple_encryption_utils.py index dfde71a..6c4c762 100644 --- a/puzzle_generator/simple_encryption_utils.py +++ b/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 diff --git a/tests/test_encryption_utils.py b/tests/test_encryption_utils.py index 6f8cdc4..1c7487d 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,45 @@ _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 diff --git a/tests/test_puzzle_data_encryption.py b/tests/test_puzzle_data_encryption.py index 422e428..2d8c7ed 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_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,7 +77,7 @@ def test_pde(in_puzzle): encrypted_puzzle["rest"], encrypted_puzzle["hash"], cur_pass + "!", - seu.decrypt_str, + decryptor, ) is None ) @@ -65,7 +85,7 @@ def test_pde(in_puzzle): 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