Merge pull request #61 from vil02/use_hashlib_scrypt

feat: use `hashlib.scrypt` to derive key

puzzle_generator/configurators/simple/common.py

@@ -1 +1,23 @@
-MODULES = ["hashlib", "itertools", "base64", "sys", "typing"]
+import secrets
+
+from ... import bytestr_utils as bu
+
+MODULES = ["hashlib", "base64", "sys", "typing"]
+
+
+def scrypt_params(**kwargs):
+    print(kwargs)
+    default = {"salt": secrets.token_bytes(16), "n": 2**14, "r": 8, "p": 1}
+    if "maxmem" in kwargs:
+        default["maxmem"] = kwargs["maxmem"]
+    res = {_k: kwargs.get(_k, _v) for _k, _v in default.items()}
+    if "maxmem" not in res:
+        res["maxmem"] = (128 + 100) * res["n"] * res["r"] * res["p"]
+    return res
+
+
+def scrypt_params_to_code_str(**kwargs) -> str:
+    pieces = [f'"{_k}":{_v}' for _k, _v in kwargs.items() if _k != "salt"]
+    salt_str = '"' + bu.bytes_to_bytestr(kwargs["salt"]) + '"'
+    pieces.append(f'"salt":bytestr_to_bytes({salt_str})')
+    return f"_SCRYPT_PARAMS = {{{','.join(pieces)}}}"

puzzle_generator/configurators/simple/simple.py

@@ -5,28 +5,28 @@
 from ...encryption_algorithms.simple import simple as se
 from ...puzzle_data_encryption import decrypt_data
 from .. import common as cc
-from .common import MODULES
+from .common import MODULES, scrypt_params, scrypt_params_to_code_str
 from ...run_puzzle import run_puzzle
 from ...bytes_utils import bytes_to_int, split
 
 
 class Simple:
     def __init__(self, **kwargs):
-        self._proc_hasher = kwargs.get("proc_hasher", cc.DefaultHasher)
+        self._scrypt_params = scrypt_params(**kwargs)
         self._signature_hasher = kwargs.get("signature_hasher", cc.DefaultHasher)
 
     def get_modules(self) -> typing.List[str]:
         return MODULES
 
     def get_encrypt(self):
-        return se.get_encrypt(self._proc_hasher, self._signature_hasher)
+        return se.get_encrypt(self._signature_hasher, self._scrypt_params)
 
     def get_needed_objects(self):
         return [
             common.hash_bytes,
-            common.int_to_bytes,
             common.split_data_and_signature,
-            common.proc_bytes,
+            common.derive_key,
+            common.xor_bytes,
             bu.bytestr_to_bytes,
             se.get_decrypt,
             bytes_to_int,
@@ -38,8 +38,10 @@ def get_needed_objects(self):
     def get_constants_str(
         self,
     ) -> str:
-        return (
+        _scrypt_params = scrypt_params_to_code_str(**self._scrypt_params)
+        decrypt: str = (
             "_DECRYPT = get_decrypt("
-            f"{cc.get_hasher_name(self._proc_hasher)}, "
-            f"{cc.get_hasher_name(self._signature_hasher)})"
+            f"{cc.get_hasher_name(self._signature_hasher)}, "
+            "_SCRYPT_PARAMS)"
         )
+        return "\n".join([_scrypt_params, decrypt])

puzzle_generator/configurators/simple/spiced.py

@@ -6,7 +6,7 @@
 from ...encryption_algorithms.simple import spiced as sse
 from ...puzzle_data_encryption import decrypt_data
 from .. import common as cc
-from .common import MODULES
+from .common import MODULES, scrypt_params, scrypt_params_to_code_str
 from ...run_puzzle import run_puzzle
 from ...bytes_utils import bytes_to_int, split
 
@@ -26,7 +26,7 @@ def _list_of_bytes_to_codestr(in_list: typing.List[bytes]) -> str:
 
 class Spiced:
     def __init__(self, **kwargs):
-        self._proc_hasher = kwargs.get("proc_hasher", cc.DefaultHasher)
+        self._scrypt_params = scrypt_params(**kwargs)
         self._signature_hasher = kwargs.get("signature_hasher", cc.DefaultHasher)
         self._proc_spices = kwargs.get("proc_spices", _get_some_spices())
         self._signature_spices = kwargs.get("signature_spices", _get_some_spices())
@@ -36,18 +36,18 @@ def get_modules(self) -> typing.List[str]:
 
     def get_encrypt(self):
         return sse.get_encrypt(
-            self._proc_hasher,
             self._signature_hasher,
             self._proc_spices,
             self._signature_spices,
+            self._scrypt_params,
         )
 
     def get_needed_objects(self):
         return [
             common.hash_bytes,
-            common.int_to_bytes,
+            common.derive_key,
             common.split_data_and_signature,
-            common.proc_bytes,
+            common.xor_bytes,
             bu.bytestr_to_bytes,
             sse.get_decrypt,
             bytes_to_int,
@@ -65,11 +65,12 @@ def get_constants_str(
         signature_spices: str = (
             f"_SIGNATURE_SPICES = {_list_of_bytes_to_codestr(self._signature_spices)}"
         )
+        _scrypt_params = scrypt_params_to_code_str(**self._scrypt_params)
         decrypt: str = (
             "_DECRYPT = get_decrypt("
-            f"{cc.get_hasher_name(self._proc_hasher)}, "
             f"{cc.get_hasher_name(self._signature_hasher)}, "
             "_PROC_SPICES, "
-            "_SIGNATURE_SPICES)"
+            "_SIGNATURE_SPICES, "
+            "_SCRYPT_PARAMS)"
         )
-        return "\n".join([proc_spices, signature_spices, decrypt])
+        return "\n".join([proc_spices, signature_spices, _scrypt_params, decrypt])

puzzle_generator/encryption_algorithms/simple/common.py

@@ -1,22 +1,18 @@
-import itertools
+import hashlib
 import typing
 
 
 def hash_bytes(in_bytes: bytes, in_hasher) -> bytes:
     return in_hasher(in_bytes).digest()
 
 
-def int_to_bytes(in_val: int) -> bytes:
-    return in_val.to_bytes((in_val.bit_length() + 7) // 8, "big")
+def derive_key(**kwargs):
+    return hashlib.scrypt(**kwargs)
 
 
-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(
-        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 xor_bytes(in_bytes: bytes, in_key: bytes) -> bytes:
+    """xors the in_bytes with a sequence of bytes with in_key"""
+    return bytes(_d ^ _k for (_d, _k) in zip(in_bytes, in_key, strict=True))
 
 
 def merge_data_and_signature(in_data: bytes, in_signature: bytes) -> bytes:

puzzle_generator/encryption_algorithms/simple/simple.py

@@ -2,29 +2,32 @@
 
 
 from .common import (
-    proc_bytes,
+    derive_key,
+    xor_bytes,
     hash_bytes,
     merge_data_and_signature,
     split_data_and_signature,
 )
 
 
 def get_encrypt(
-    proc_hasher, signature_hasher
+    signature_hasher, scrypt_params
 ) -> typing.Callable[[bytes, bytes], bytes]:
     def _encrypt(in_bytes: bytes, in_pass: bytes) -> bytes:
         signature = hash_bytes(in_bytes, signature_hasher)
         merged = merge_data_and_signature(in_bytes, signature)
-        return proc_bytes(merged, in_pass, proc_hasher)
+        key = derive_key(password=in_pass, dklen=len(merged), **scrypt_params)
+        return xor_bytes(merged, key)
 
     return _encrypt
 
 
 def get_decrypt(
-    proc_hasher, signature_hasher
+    signature_hasher, scrypt_params
 ) -> typing.Callable[[bytes, bytes], bytes | None]:
     def _decrypt(in_bytes: bytes, in_pass: bytes) -> bytes | None:
-        data = proc_bytes(in_bytes, in_pass, proc_hasher)
+        key = derive_key(password=in_pass, dklen=len(in_bytes), **scrypt_params)
+        data = xor_bytes(in_bytes, key)
         decrypted, signature = split_data_and_signature(
             data, signature_hasher().digest_size
         )

puzzle_generator/encryption_algorithms/simple/spiced.py

@@ -2,18 +2,19 @@
 import secrets
 
 from .common import (
-    proc_bytes,
+    derive_key,
+    xor_bytes,
     hash_bytes,
     merge_data_and_signature,
     split_data_and_signature,
 )
 
 
 def get_encrypt(
-    proc_hasher,
     signature_hasher,
     proc_spices: typing.List[bytes],
     signature_spices: typing.List[bytes],
+    scrypt_params,
 ) -> typing.Callable[[bytes, bytes], bytes]:
     assert proc_spices  # nosec B101
     assert signature_spices  # nosec B101
@@ -23,23 +24,29 @@ def _encrypt(in_bytes: bytes, in_pass: bytes) -> bytes:
         signature = hash_bytes(in_bytes + signature_spice, signature_hasher)
         merged = merge_data_and_signature(in_bytes, signature)
         proc_spice = secrets.choice(proc_spices)
-        return proc_bytes(merged, in_pass + proc_spice, proc_hasher)
+        key = derive_key(
+            password=in_pass + proc_spice, dklen=len(merged), **scrypt_params
+        )
+        return xor_bytes(merged, key)
 
     return _encrypt
 
 
 def get_decrypt(
-    proc_hasher,
     signature_hasher,
     proc_spices: typing.List[bytes],
     signature_spices: typing.List[bytes],
+    scrypt_params,
 ) -> typing.Callable[[bytes, bytes], bytes | None]:
     assert proc_spices  # nosec B101
     assert signature_spices  # nosec B101
 
     def _decrypt(in_bytes: bytes, in_pass: bytes) -> bytes | None:
         for proc_spice in proc_spices:
-            data = proc_bytes(in_bytes, in_pass + proc_spice, proc_hasher)
+            key = derive_key(
+                password=in_pass + proc_spice, dklen=len(in_bytes), **scrypt_params
+            )
+            data = xor_bytes(in_bytes, key)
             decrypted, signature = split_data_and_signature(
                 data, signature_hasher().digest_size
             )

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "puzzle-generator"
-version = "0.8.0"
+version = "0.9.0"
 description = "Generates python code representing a puzzle"
 authors = ["piotr.idzik <vil02_puzzle_generator@10g.pl>"]
 readme = "./puzzle_generator/README.md"

tests/encryption_algorithms/test_encryption_algorithms.py

@@ -12,6 +12,10 @@
     hashlib.blake2s,
 ]
 
+_SOME_SCRYPT_PARAMS = [
+    {"salt": b"some_bad_salt", "n": 8, "r": 10, "p": 1},
+    {"salt": b"some_other_bad_salt", "n": 16, "r": 20, "p": 1},
+]
 
 _PROC_SPICES = [b"a", b"bb", b"ccc", b"dddd"]
 _SIGNATURE_SPICES = [b"XXX", b"YY", b"Z"]
@@ -22,14 +26,14 @@
 @pytest.mark.parametrize(
     ("encrypt", "decrypt"),
     [
-        utils.get_simple_encrypt_decrypt_pair(*_)
-        for _ in itertools.product(_SOME_HASHES, repeat=2)
+        utils.get_simple_encrypt_decrypt_pair(hash, scrypt_params)
+        for hash, scrypt_params in itertools.product(_SOME_HASHES, _SOME_SCRYPT_PARAMS)
     ]
     + [
         utils.get_spiced_simple_encrypt_decrypt_pair(
-            *_, _PROC_SPICES, _SIGNATURE_SPICES
+            hash, _PROC_SPICES, _SIGNATURE_SPICES, scrypt_params
         )
-        for _ in itertools.product(_SOME_HASHES, repeat=2)
+        for hash, scrypt_params in itertools.product(_SOME_HASHES, _SOME_SCRYPT_PARAMS)
     ],
 )
 def test_encryption_decryption(in_bytes, in_pass, encrypt, decrypt):

tests/test_create_puzzle.py

@@ -85,22 +85,23 @@ def _run_puzzle_str(
 
 _CONFIGURATIONS = [
     {},
-    {"proc_hasher": hashlib.md5},
-    {"proc_hasher": hashlib.sha1, "signature_hasher": hashlib.sha3_224},
-    {"signature_hasher": hashlib.blake2b},
+    {"signature_hasher": hashlib.sha3_224, "n": 2**2},
+    {"signature_hasher": hashlib.blake2b, "salt": b"very_bad_salt"},
     {"encryption": "simple", "signature_hasher": hashlib.blake2b},
     {
         "encryption": "spiced",
-        "proc_hasher": hashlib.sha224,
         "signature_hasher": hashlib.sha3_224,
         "proc_spices": [b"\0"],
+        "n": 2**3,
+        "maxmem": 0,
     },
     {
         "encryption": "spiced",
-        "proc_hasher": hashlib.sha3_224,
         "signature_hasher": hashlib.sha224,
         "proc_spices": [b"\1"],
         "signature_spices": [b"\2"],
+        "salt": b"even_worse_salt!",
+        "r": 16,
     },
 ]
 
@@ -153,12 +154,19 @@ def _input_simulator() -> str:
 _PROC_SPICES = [b"11", b"22"]
 _SIGNATURE_SPICES = [b"27", b"07", b"2024"]
 
+_SOME_SCRYPT_PARAMS = [
+    {"salt": b"salt_1", "n": 8, "r": 5, "p": 1},
+    {"salt": b"salt_two", "n": 4, "r": 2, "p": 1},
+]
+
 _ENCRYPT_DECRYPT_PAIRS = [
-    utils.get_simple_encrypt_decrypt_pair(*_)
-    for _ in itertools.product(_SOME_HASHES, repeat=2)
+    utils.get_simple_encrypt_decrypt_pair(hash, scrypt_params)
+    for hash, scrypt_params in itertools.product(_SOME_HASHES, _SOME_SCRYPT_PARAMS)
 ] + [
-    utils.get_spiced_simple_encrypt_decrypt_pair(*_, _PROC_SPICES, _SIGNATURE_SPICES)
-    for _ in itertools.product(_SOME_HASHES, repeat=2)
+    utils.get_spiced_simple_encrypt_decrypt_pair(
+        hash, _PROC_SPICES, _SIGNATURE_SPICES, scrypt_params
+    )
+    for hash, scrypt_params in itertools.product(_SOME_HASHES, _SOME_SCRYPT_PARAMS)
 ]
 
 

tests/test_puzzle_data_encryption.py

@@ -15,6 +15,12 @@
 _SIGNATURE_SPICES = [b"1", b"12"]
 
 
+_SOME_SCRYPT_PARAMS = [
+    {"salt": b"some_bad_salt_0", "n": 8, "r": 5, "p": 1},
+    {"salt": b"some_other_bad_salt_1", "n": 4, "r": 2, "p": 1},
+]
+
+
 @pytest.mark.parametrize(
     "in_puzzle",
     [
@@ -60,14 +66,14 @@
 @pytest.mark.parametrize(
     ("encrypt", "decrypt"),
     [
-        utils.get_simple_encrypt_decrypt_pair(*_)
-        for _ in itertools.product(_SOME_HASHES, repeat=2)
+        utils.get_simple_encrypt_decrypt_pair(hash, scrypt_params)
+        for hash, scrypt_params in itertools.product(_SOME_HASHES, _SOME_SCRYPT_PARAMS)
     ]
     + [
         utils.get_spiced_simple_encrypt_decrypt_pair(
-            *_, _PROC_SPICES, _SIGNATURE_SPICES
+            hash, _PROC_SPICES, _SIGNATURE_SPICES, scrypt_params
         )
-        for _ in itertools.product(_SOME_HASHES, repeat=2)
+        for hash, scrypt_params in itertools.product(_SOME_HASHES, _SOME_SCRYPT_PARAMS)
     ],
 )
 def test_pde(in_puzzle, encrypt, decrypt):