from Crypto.Hash import (
    SHA256 as _sha256,
    SHA512 as _sha512,
)
from Crypto.Protocol.KDF import (
    scrypt as _scrypt,
    HKDF as _HKDF,
    PBKDF2 as _PBKDF2,
)
from Crypto.Cipher import (
    AES as _AES
)


def SHA256(x):
    return _sha256.new(x).digest()


def scrypt(*, password: str, salt: str, n: int, r: int, p: int, dklen: int) -> bytes:
    assert(n < 2**(128 * r / 8))
    res = _scrypt(password=password, salt=salt, key_len=dklen, N=n, r=r, p=p)
    return res if isinstance(res, bytes) else res[0]  # PyCryptodome can return Tuple[bytes]


def PBKDF2(*, password: str, salt: bytes, dklen: int, c: int, prf: str) -> bytes:
    assert('sha' in prf)
    _hash = _sha256 if 'sha256' in prf else _sha512
    password_bytes = password.encode("utf-8")
    res = _PBKDF2(password=password_bytes, salt=salt, dkLen=dklen, count=c, hmac_hash_module=_hash)  # type: ignore
    return res if isinstance(res, bytes) else res[0]  # PyCryptodome can return Tuple[bytes]


def HKDF(*, salt: bytes, IKM: bytes, L: int) -> bytes:
    res = _HKDF(master=IKM, key_len=L, salt=salt, hashmod=_sha256)
    return res if isinstance(res, bytes) else res[0]  # PyCryptodome can return Tuple[bytes]


def AES_128_CTR(*, key: bytes, iv: bytes):
    return _AES.new(key=key, mode=_AES.MODE_CTR, initial_value=iv, nonce=b'')