Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: use hashlib.scrypt to derive key #61

Merged
merged 2 commits into from
Aug 2, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 23 additions & 1 deletion puzzle_generator/configurators/simple/common.py
Original file line number Diff line number Diff line change
@@ -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)}}}"
18 changes: 10 additions & 8 deletions puzzle_generator/configurators/simple/simple.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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])
17 changes: 9 additions & 8 deletions puzzle_generator/configurators/simple/spiced.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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())
Expand All @@ -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,
Expand All @@ -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])
16 changes: 6 additions & 10 deletions puzzle_generator/encryption_algorithms/simple/common.py
Original file line number Diff line number Diff line change
@@ -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:
Expand Down
13 changes: 8 additions & 5 deletions puzzle_generator/encryption_algorithms/simple/simple.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand Down
17 changes: 12 additions & 5 deletions puzzle_generator/encryption_algorithms/simple/spiced.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
)
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -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 <[email protected]>"]
readme = "./puzzle_generator/README.md"
Expand Down
12 changes: 8 additions & 4 deletions tests/encryption_algorithms/test_encryption_algorithms.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand All @@ -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):
Expand Down
26 changes: 17 additions & 9 deletions tests/test_create_puzzle.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
]

Expand Down Expand Up @@ -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)
]


Expand Down
14 changes: 10 additions & 4 deletions tests/test_puzzle_data_encryption.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
[
Expand Down Expand Up @@ -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):
Expand Down