Skip to content

Commit

Permalink
feat: open for all hashes from hashlib
Browse files Browse the repository at this point in the history
  • Loading branch information
vil02 committed Aug 2, 2024
1 parent 227703e commit 995110c
Show file tree
Hide file tree
Showing 12 changed files with 170 additions and 187 deletions.
8 changes: 0 additions & 8 deletions puzzle_generator/configurators/common.py

This file was deleted.

36 changes: 30 additions & 6 deletions puzzle_generator/configurators/simple/common.py
Original file line number Diff line number Diff line change
@@ -1,23 +1,47 @@
import secrets

from ... import bytestr_utils as bu
from ...encryption_algorithms.simple import common
from ... import bytestr_utils as bsu
from ... import bytes_utils as bu
from ... import puzzle_data_encryption as pde
from ... import run_puzzle as rp

MODULES = ["hashlib", "base64", "sys", "typing"]

OBJECTS = [
common.hash_bytes,
common.derive_key,
common.split_data_and_signature,
common.digest_size,
common.xor_bytes,
bsu.bytestr_to_bytes,
bu.bytes_to_int,
bu.split,
pde.decrypt_data,
rp.run_puzzle,
]


def scrypt_params(**kwargs):
print(kwargs)
in_params = kwargs.get("scrypt_params", {})
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" in in_params:
default["maxmem"] = in_params["maxmem"]
res = {_k: in_params.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 signature_params(**kwargs):
res = kwargs.get("signature_params", {"hasher": {"name": "sha512"}, "digest": {}})
if "digest" not in res:
res["digest"] = {}
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"]) + '"'
salt_str = '"' + bsu.bytes_to_bytestr(kwargs["salt"]) + '"'
pieces.append(f'"salt":bytestr_to_bytes({salt_str})')
return f"_SCRYPT_PARAMS = {{{','.join(pieces)}}}"
38 changes: 10 additions & 28 deletions puzzle_generator/configurators/simple/simple.py
Original file line number Diff line number Diff line change
@@ -1,47 +1,29 @@
import typing

from ... import bytestr_utils as bu
from ...encryption_algorithms.simple import common
from ...encryption_algorithms.simple import simple as se
from ...puzzle_data_encryption import decrypt_data
from .. import common as cc
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
from . import common as csc


class Simple:
def __init__(self, **kwargs):
self._scrypt_params = scrypt_params(**kwargs)
self._signature_hasher = kwargs.get("signature_hasher", cc.DefaultHasher)
self._scrypt_params = csc.scrypt_params(**kwargs)
self._signature_params = csc.signature_params(**kwargs)

def get_modules(self) -> typing.List[str]:
return MODULES
return csc.MODULES

def get_encrypt(self):
return se.get_encrypt(self._signature_hasher, self._scrypt_params)
return se.get_encrypt(self._scrypt_params, self._signature_params)

def get_needed_objects(self):
return [
common.hash_bytes,
common.split_data_and_signature,
common.derive_key,
common.xor_bytes,
bu.bytestr_to_bytes,
return csc.OBJECTS + [
se.get_decrypt,
bytes_to_int,
split,
decrypt_data,
run_puzzle,
]

def get_constants_str(
self,
) -> str:
_scrypt_params = scrypt_params_to_code_str(**self._scrypt_params)
decrypt: str = (
"_DECRYPT = get_decrypt("
f"{cc.get_hasher_name(self._signature_hasher)}, "
"_SCRYPT_PARAMS)"
)
return "\n".join([_scrypt_params, decrypt])
_scrypt_params = csc.scrypt_params_to_code_str(**self._scrypt_params)
_signature_params = "_SIGNATURE_PARAMS = " + repr(self._signature_params)
decrypt: str = "_DECRYPT = get_decrypt(_SCRYPT_PARAMS, _SIGNATURE_PARAMS)"
return "\n".join([_scrypt_params, _signature_params, decrypt])
36 changes: 12 additions & 24 deletions puzzle_generator/configurators/simple/spiced.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,8 @@
import typing

from ... import bytestr_utils as bu
from ...encryption_algorithms.simple import common
from ...encryption_algorithms.simple import spiced as sse
from ...puzzle_data_encryption import decrypt_data
from .. import common as cc
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
from . import common as csc


def _get_some_spices():
Expand All @@ -26,34 +21,25 @@ def _list_of_bytes_to_codestr(in_list: typing.List[bytes]) -> str:

class Spiced:
def __init__(self, **kwargs):
self._scrypt_params = scrypt_params(**kwargs)
self._signature_hasher = kwargs.get("signature_hasher", cc.DefaultHasher)
self._scrypt_params = csc.scrypt_params(**kwargs)
self._signature_params = csc.signature_params(**kwargs)
self._proc_spices = kwargs.get("proc_spices", _get_some_spices())
self._signature_spices = kwargs.get("signature_spices", _get_some_spices())

def get_modules(self) -> typing.List[str]:
return MODULES
return csc.MODULES

def get_encrypt(self):
return sse.get_encrypt(
self._signature_hasher,
self._proc_spices,
self._signature_spices,
self._scrypt_params,
self._signature_params,
)

def get_needed_objects(self):
return [
common.hash_bytes,
common.derive_key,
common.split_data_and_signature,
common.xor_bytes,
bu.bytestr_to_bytes,
return csc.OBJECTS + [
sse.get_decrypt,
bytes_to_int,
split,
decrypt_data,
run_puzzle,
]

def get_constants_str(
Expand All @@ -65,12 +51,14 @@ 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)
_scrypt_params = csc.scrypt_params_to_code_str(**self._scrypt_params)
_signature_params = "_SIGNATURE_PARAMS = " + repr(self._signature_params)
decrypt: str = (
"_DECRYPT = get_decrypt("
f"{cc.get_hasher_name(self._signature_hasher)}, "
"_PROC_SPICES, "
"_SIGNATURE_SPICES, "
"_SCRYPT_PARAMS)"
"_SCRYPT_PARAMS, _SIGNATURE_PARAMS)"
)
return "\n".join(
[proc_spices, signature_spices, _scrypt_params, _signature_params, decrypt]
)
return "\n".join([proc_spices, signature_spices, _scrypt_params, decrypt])
15 changes: 13 additions & 2 deletions puzzle_generator/encryption_algorithms/simple/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,19 @@
import typing


def hash_bytes(in_bytes: bytes, in_hasher) -> bytes:
return in_hasher(in_bytes).digest()
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


def derive_key(**kwargs):
Expand Down
11 changes: 6 additions & 5 deletions puzzle_generator/encryption_algorithms/simple/simple.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,16 +5,17 @@
derive_key,
xor_bytes,
hash_bytes,
digest_size,
merge_data_and_signature,
split_data_and_signature,
)


def get_encrypt(
signature_hasher, scrypt_params
scrypt_params, signature_params
) -> typing.Callable[[bytes, bytes], bytes]:
def _encrypt(in_bytes: bytes, in_pass: bytes) -> bytes:
signature = hash_bytes(in_bytes, signature_hasher)
signature = hash_bytes(in_bytes, 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)
Expand All @@ -23,16 +24,16 @@ def _encrypt(in_bytes: bytes, in_pass: bytes) -> bytes:


def get_decrypt(
signature_hasher, scrypt_params
scrypt_params, signature_params
) -> typing.Callable[[bytes, bytes], bytes | None]:
def _decrypt(in_bytes: bytes, in_pass: bytes) -> bytes | None:
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
data, digest_size(signature_params)
)

if hash_bytes(decrypted, signature_hasher) == signature:
if hash_bytes(decrypted, signature_params) == signature:
return decrypted
return None

Expand Down
11 changes: 6 additions & 5 deletions puzzle_generator/encryption_algorithms/simple/spiced.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,23 +5,24 @@
derive_key,
xor_bytes,
hash_bytes,
digest_size,
merge_data_and_signature,
split_data_and_signature,
)


def get_encrypt(
signature_hasher,
proc_spices: typing.List[bytes],
signature_spices: typing.List[bytes],
scrypt_params,
signature_params,
) -> typing.Callable[[bytes, bytes], bytes]:
assert proc_spices # nosec B101
assert signature_spices # nosec B101

def _encrypt(in_bytes: bytes, in_pass: bytes) -> bytes:
signature_spice = secrets.choice(signature_spices)
signature = hash_bytes(in_bytes + signature_spice, signature_hasher)
signature = hash_bytes(in_bytes + signature_spice, signature_params)
merged = merge_data_and_signature(in_bytes, signature)
proc_spice = secrets.choice(proc_spices)
key = derive_key(
Expand All @@ -33,10 +34,10 @@ def _encrypt(in_bytes: bytes, in_pass: bytes) -> bytes:


def get_decrypt(
signature_hasher,
proc_spices: typing.List[bytes],
signature_spices: typing.List[bytes],
scrypt_params,
signature_params,
) -> typing.Callable[[bytes, bytes], bytes | None]:
assert proc_spices # nosec B101
assert signature_spices # nosec B101
Expand All @@ -48,11 +49,11 @@ def _decrypt(in_bytes: bytes, in_pass: bytes) -> bytes | None:
)
data = xor_bytes(in_bytes, key)
decrypted, signature = split_data_and_signature(
data, signature_hasher().digest_size
data, digest_size(signature_params)
)

if any(
hash_bytes(decrypted + _, signature_hasher) == signature
hash_bytes(decrypted + _, signature_params) == signature
for _ in signature_spices
):
return decrypted
Expand Down
10 changes: 10 additions & 0 deletions tests/encryption_algorithms/test_common.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
import pytest

import puzzle_generator.encryption_algorithms.simple.common as eac
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)
32 changes: 1 addition & 31 deletions tests/encryption_algorithms/test_encryption_algorithms.py
Original file line number Diff line number Diff line change
@@ -1,41 +1,11 @@
import hashlib
import itertools
import pytest

from .. import utils


_SOME_HASHES = [
hashlib.md5,
hashlib.sha512,
hashlib.sha3_512,
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"]


@pytest.mark.parametrize("in_bytes", utils.BYTES_LIST)
@pytest.mark.parametrize("in_pass", utils.BYTES_LIST)
@pytest.mark.parametrize(
("encrypt", "decrypt"),
[
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(
hash, _PROC_SPICES, _SIGNATURE_SPICES, scrypt_params
)
for hash, scrypt_params in itertools.product(_SOME_HASHES, _SOME_SCRYPT_PARAMS)
],
)
@pytest.mark.parametrize(("encrypt", "decrypt"), utils.ENCRYPT_DECRYPT_PAIRS)
def test_encryption_decryption(in_bytes, in_pass, encrypt, decrypt):
encrypted = encrypt(in_bytes, in_pass)

Expand Down
Loading

0 comments on commit 995110c

Please sign in to comment.