Merge pull request #77 from vil02/use_hmac

feat: use `hmac` to create signatures
vil02 · Aug 19, 2024 · 97f2133 · 97f2133
2 parents 99e7f7f + 1db6aa2
commit 97f2133
Show file tree

Hide file tree

Showing 8 changed files with 37 additions and 67 deletions.
diff --git a/puzzle_generator/configurators/simple/common.py b/puzzle_generator/configurators/simple/common.py
@@ -6,10 +6,10 @@
 from ... import puzzle_data_encryption as pde
 from ... import run_puzzle as rp
 
-MODULES = ["hashlib", "base64", "sys", "typing"]
+MODULES = ["hmac", "hashlib", "base64", "sys", "typing"]
 
 OBJECTS = [
-    common.hash_bytes,
+    common.sign_bytes,
     common.derive_key,
     common.split_data_and_signature,
     common.digest_size,
@@ -34,10 +34,7 @@ def scrypt_params(**kwargs):
 
 
 def signature_params(**kwargs):
-    res = kwargs.get("signature_params", {"hasher": {"name": "sha512"}, "digest": {}})
-    if "digest" not in res:
-        res["digest"] = {}
-    return res
+    return kwargs.get("signature_params", {"digest": "sha512"})
 
 
 def scrypt_params_to_code_str(**kwargs) -> str:

diff --git a/puzzle_generator/encryption_algorithms/simple/common.py b/puzzle_generator/encryption_algorithms/simple/common.py
@@ -1,20 +1,15 @@
+import hmac
 import hashlib
 import typing
 
 
 def digest_size(params) -> int:
-    hasher = hashlib.new(**params["hasher"])
-    res = hasher.digest_size
-    if res > 0:
-        return res
-    return params["digest"]["length"]
-
-
-def hash_bytes(in_bytes: bytes, params) -> bytes:
-    hasher = hashlib.new(**params["hasher"])
-    hasher.update(in_bytes)
-    res = hasher.digest(**params["digest"])
-    return res
+    hasher = hashlib.new(params["digest"])
+    return hasher.digest_size
+
+
+def sign_bytes(in_bytes: bytes, in_key: bytes, params) -> bytes:
+    return hmac.digest(msg=in_bytes, key=in_key, **params)
 
 
 def derive_key(**kwargs):

diff --git a/puzzle_generator/encryption_algorithms/simple/simple.py b/puzzle_generator/encryption_algorithms/simple/simple.py
@@ -4,7 +4,7 @@
 from .common import (
     derive_key,
     xor_bytes,
-    hash_bytes,
+    sign_bytes,
     digest_size,
     merge_data_and_signature,
     split_data_and_signature,
@@ -15,7 +15,7 @@ def get_encrypt(
     scrypt_params, signature_params
 ) -> typing.Callable[[bytes, bytes], bytes]:
     def _encrypt(in_bytes: bytes, in_pass: bytes) -> bytes:
-        signature = hash_bytes(in_bytes, signature_params)
+        signature = sign_bytes(in_bytes, in_pass, signature_params)
         merged = merge_data_and_signature(in_bytes, signature)
         key = derive_key(password=in_pass, dklen=len(merged), **scrypt_params)
         return xor_bytes(merged, key)
@@ -33,7 +33,7 @@ def _decrypt(in_bytes: bytes, in_pass: bytes) -> bytes | None:
             data, digest_size(signature_params)
         )
 
-        if hash_bytes(decrypted, signature_params) == signature:
+        if sign_bytes(decrypted, in_pass, signature_params) == signature:
             return decrypted
         return None
 

diff --git a/puzzle_generator/encryption_algorithms/simple/spiced.py b/puzzle_generator/encryption_algorithms/simple/spiced.py
@@ -4,7 +4,7 @@
 from .common import (
     derive_key,
     xor_bytes,
-    hash_bytes,
+    sign_bytes,
     digest_size,
     merge_data_and_signature,
     split_data_and_signature,
@@ -22,7 +22,7 @@ def get_encrypt(
 
     def _encrypt(in_bytes: bytes, in_pass: bytes) -> bytes:
         signature_spice = secrets.choice(signature_spices)
-        signature = hash_bytes(in_bytes + signature_spice, signature_params)
+        signature = sign_bytes(in_bytes + signature_spice, in_pass, signature_params)
         merged = merge_data_and_signature(in_bytes, signature)
         proc_spice = secrets.choice(proc_spices)
         key = derive_key(
@@ -53,7 +53,7 @@ def _decrypt(in_bytes: bytes, in_pass: bytes) -> bytes | None:
             )
 
             if any(
-                hash_bytes(decrypted + _, signature_params) == signature
+                sign_bytes(decrypted + _, in_pass, signature_params) == signature
                 for _ in signature_spices
             ):
                 return decrypted

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "puzzle-generator"
-version = "0.11.0"
+version = "0.12.0"
 description = "Generates python code representing a puzzle"
 authors = ["piotr.idzik <[email protected]>"]
 readme = "./puzzle_generator/README.md"

diff --git a/tests/encryption_algorithms/test_common.py b/tests/encryption_algorithms/test_common.py
@@ -4,7 +4,7 @@
 from .. import utils
 
 
-@pytest.mark.parametrize("in_hash_params", utils.SOME_SIGNATURE_PARAMS)
-def test_digest_size(in_hash_params):
-    some_hash = eac.hash_bytes(b"some_msg", in_hash_params)
-    assert eac.digest_size(in_hash_params) == len(some_hash)
+@pytest.mark.parametrize("in_signature_params", utils.SOME_SIGNATURE_PARAMS)
+def test_digest_size(in_signature_params):
+    some_hash = eac.sign_bytes(b"some_msg", b"some_key", in_signature_params)
+    assert eac.digest_size(in_signature_params) == len(some_hash)
diff --git a/tests/test_create_puzzle.py b/tests/test_create_puzzle.py
@@ -86,38 +86,19 @@ def _run_puzzle_str(
     {"encryption": "simple"},
     {"encryption": "spiced"},
     {"scrypt_params": {"n": 2**4, "p": 2, "maxmem": 200000}},
-    {
-        "signature_params": {
-            "hasher": {"name": "sha3_384"},
-        }
-    },
+    {"signature_params": {"digest": "sha3_384"}},
     {"encryption": "simple", "scrypt_params": {"n": 2**3, "maxmem": 100000}},
-    {
-        "encryption": "simple",
-        "signature_params": {"hasher": {"name": "blake2b", "digest_size": 17}},
-    },
-    {
-        "encryption": "simple",
-        "signature_params": {
-            "hasher": {"name": "shake256", "data": b"init"},
-            "digest": {"length": 91},
-        },
-    },
+    {"encryption": "simple", "signature_params": {"digest": "blake2b"}},
+    {"encryption": "simple", "signature_params": {"digest": "blake2s"}},
     {
         "encryption": "spiced",
         "proc_spices": [b"\1"],
-        "signature_params": {
-            "hasher": {"name": "shake128"},
-            "digest": {"length": 5},
-        },
+        "signature_params": {"digest": "sha3_512"},
     },
     {
         "encryption": "spiced",
         "signature_spices": [b"\0", b"\10"],
-        "signature_params": {
-            "hasher": {"name": "sha3_256", "data": b"00000"},
-            "digest": {},
-        },
+        "signature_params": {"digest": "sha3_256"},
         "scrypt_params": {"n": 2**5, "r": 16, "salt": b"testSalt!!!"},
     },
 ]

diff --git a/tests/utils.py b/tests/utils.py
@@ -26,22 +26,19 @@
 PROC_SPICES = [b"a", b"bb", b"ccc", b"dddd"]
 SIGNATURE_SPICES = [b"XXX", b"YY", b"Z"]
 
-SOME_SIGNATURE_PARAMS = [
-    {"hasher": {"name": "sha512"}, "digest": {}},
-    {"hasher": {"name": "sha3_512", "data": b"initial_data"}, "digest": {}},
-    {"hasher": {"name": "blake2b", "digest_size": 63}, "digest": {}},
-    {"hasher": {"name": "blake2s", "digest_size": 10, "person": b"tmp?"}, "digest": {}},
-    {"hasher": {"name": "shake_256"}, "digest": {"length": 999}},
-    {
-        "hasher": {
-            "name": "shake_128",
-            "data": b"some_initial_data",
-            "usedforsecurity": False,
-        },
-        "digest": {"length": 10},
-    },
+_SOME_HASHES = [
+    "sha256",
+    "sha384",
+    "sha512",
+    "sha3_256",
+    "sha3_384",
+    "sha3_512",
+    "blake2b",
+    "blake2s",
 ]
 
+SOME_SIGNATURE_PARAMS = [{"digest": _} for _ in _SOME_HASHES]
+
 
 def _get_simple_encrypt_decrypt_pair(*args):
     return se.get_encrypt(*args), se.get_decrypt(*args)