<commit_before>import base64
import os

from Crypto.Cipher import AES
from Crypto.Hash import HMAC, SHA256
from Crypto.Protocol.KDF import PBKDF2

class UnknownAlgorithmError(Exception):
    pass

class MessageAuthenticationError(Exception):
    pass

class CipherInstance():
    def __init__(self, key, algo=None):
        self.key = key
        self.algo = algo or ("aes-128-cfb-siv" if key else "none")

    def _generate_key(self, nbits) -> "bytes":
        if nbits % 8:
            raise ValueError("nbits not divisible by 8")
        return os.urandom(nbits // 8)

    def _get_key_bits(self, nbits) -> "bytes":
        if nbits % 8:
            raise ValueError("nbits not divisible by 8")
        nbytes = nbits // 8
        if len(self.key) < nbytes:
            raise ValueError("key too short (%d < %d)" % (len(self.key), nbytes))
        return self.key[:nbits // 8]

    def _deterministic_iv(self, clear: "bytes", nbytes) -> "bytes":
        mac = HMAC.new(self.key, clear, SHA256).digest()
        if len(mac) < nbytes:
            raise ValueError("resulting mac too short (%d < %d)" % (len(mac), nbytes))
        return mac[:nbytes]

    def _encrypt_data(self, clear: "bytes", algo: "str") -> "bytes":
        algo = algo.split("-")
        if algo[0] == "none":
            return clear
        elif algo[0] == "aes":
            if algo[1] in {"128", "192", "256"}:
                nbits = int(algo[1])
                key = self._get_key_bits(nbits)
                if algo[2] == "cfb":
                    if "siv" in algo[3:]:
                        iv = self._deterministic_iv(clear, AES.block_size)
                    else:
                        iv = os.urandom(AES.block_size)
                    cipher = AES.new(key, AES.MODE_CFB, iv)
                    return iv + cipher.encrypt(clear)
        else:
            raise UnknownAlgorithmError()

    def _decrypt_data(self, wrapped: "bytes", algo: "str") -> "bytes":
        algo = algo.split("-")
        if algo[0] == "none":
            return wrapped
        elif algo[0] == "aes":
            if algo[1] in {"128", "192", "256"}:
                nbits = int(algo[1])
                key = self._get_key_bits(nbits)
                if algo[2] == "cfb":
                    iv = wrapped[:AES.block_size]
                    buf = wrapped[AES.block_size:]
                    cipher = AES.new(key, AES.MODE_CFB, iv)
                    clear = cipher.decrypt(buf)
                    if "siv" in algo[3:]:
                        if iv != self._deterministic_iv(clear, AES.block_size):
                            raise MessageAuthenticationError()
                    return clear
        else:
            raise UnknownAlgorithmError()

    def wrap_bytes(self, clear: "bytes") -> "str":
        wrapped = self._encrypt_data(clear, self.algo)
        wrapped = base64.b64encode(wrapped).decode("utf-8")
        wrapped = "%s;%s" % (self.algo, wrapped)
        return wrapped

    def unwrap_bytes(self, wrapped: "str") -> "str":
        algo, wrapped = wrapped.split(";", 1)
        wrapped = base64.b64decode(wrapped.encode("utf-8"))
        clear = self._decrypt_data(wrapped, algo)
        return clear

    def wrap_str(self, clear: "str") -> "str":
        return self.wrap_bytes(clear.encode("utf-8"))

    def unwrap_str(self, wrapped: "str") -> "str":
        return self.unwrap_bytes(wrapped).decode("utf-8")

class SecureStorage():
    def __init__(self):
        self.kdf_salt = b"\x25\xa9\x7b\xc5\x7a\x59\x0d\xa6"
        self.kek_cipher = None
        self.dek_cipher = None

    # KEK

    @property
    def has_kek(self):
        return bool(self.kek_cipher and self.kek_cipher.key)

    def set_raw_kek(self, kek):
        if self.kek_cipher:
            raise Exception("KEK already set")
        else:
            self.kek_cipher = CipherInstance(kek)

    def change_raw_kek(self, new_kek):
        if not self.kek_cipher:
            raise Exception("KEK not yet set")
        elif not self.dek_cipher:
            raise Exception("DEK not yet set")
        else:
            self.kek_cipher = CipherInstance(new_kek)

    def kdf(self, passwd):
        return PBKDF2(passwd.encode("utf-8"), self.kdf_salt)

    def set_password(self, passwd):
        return self.set_raw_kek(self.kdf(passwd))

    def change_password(self, passwd):
        return self.change_raw_kek(self.kdf(passwd))

    # DEK

    @property
    def has_dek(self):
        return bool(self.dek_cipher and self.dek_cipher.key)<commit_msg>encryption: explicitly specify the default PBKDF2 parameters<commit_after>import base64
import os

from Crypto.Cipher import AES
from Crypto.Hash import HMAC, SHA1, SHA256
from Crypto.Protocol.KDF import PBKDF2

class UnknownAlgorithmError(Exception):
    pass

class MessageAuthenticationError(Exception):
    pass

class CipherInstance():
    def __init__(self, key, algo=None):
        self.key = key
        self.algo = algo or ("aes-128-cfb-siv" if key else "none")

    def _generate_key(self, nbits) -> "bytes":
        if nbits % 8:
            raise ValueError("nbits not divisible by 8")
        return os.urandom(nbits // 8)

    def _get_key_bits(self, nbits) -> "bytes":
        if nbits % 8:
            raise ValueError("nbits not divisible by 8")
        nbytes = nbits // 8
        if len(self.key) < nbytes:
            raise ValueError("key too short (%d < %d)" % (len(self.key), nbytes))
        return self.key[:nbits // 8]

    def _deterministic_iv(self, clear: "bytes", nbytes) -> "bytes":
        mac = HMAC.new(self.key, clear, SHA256).digest()
        if len(mac) < nbytes:
            raise ValueError("resulting mac too short (%d < %d)" % (len(mac), nbytes))
        return mac[:nbytes]

    def _encrypt_data(self, clear: "bytes", algo: "str") -> "bytes":
        algo = algo.split("-")
        if algo[0] == "none":
            return clear
        elif algo[0] == "aes":
            if algo[1] in {"128", "192", "256"}:
                nbits = int(algo[1])
                key = self._get_key_bits(nbits)
                if algo[2] == "cfb":
                    if "siv" in algo[3:]:
                        iv = self._deterministic_iv(clear, AES.block_size)
                    else:
                        iv = os.urandom(AES.block_size)
                    cipher = AES.new(key, AES.MODE_CFB, iv)
                    return iv + cipher.encrypt(clear)
        else:
            raise UnknownAlgorithmError()

    def _decrypt_data(self, wrapped: "bytes", algo: "str") -> "bytes":
        algo = algo.split("-")
        if algo[0] == "none":
            return wrapped
        elif algo[0] == "aes":
            if algo[1] in {"128", "192", "256"}:
                nbits = int(algo[1])
                key = self._get_key_bits(nbits)
                if algo[2] == "cfb":
                    iv = wrapped[:AES.block_size]
                    buf = wrapped[AES.block_size:]
                    cipher = AES.new(key, AES.MODE_CFB, iv)
                    clear = cipher.decrypt(buf)
                    if "siv" in algo[3:]:
                        if iv != self._deterministic_iv(clear, AES.block_size):
                            raise MessageAuthenticationError()
                    return clear
        else:
            raise UnknownAlgorithmError()

    def wrap_bytes(self, clear: "bytes") -> "str":
        wrapped = self._encrypt_data(clear, self.algo)
        wrapped = base64.b64encode(wrapped).decode("utf-8")
        wrapped = "%s;%s" % (self.algo, wrapped)
        return wrapped

    def unwrap_bytes(self, wrapped: "str") -> "str":
        algo, wrapped = wrapped.split(";", 1)
        wrapped = base64.b64decode(wrapped.encode("utf-8"))
        clear = self._decrypt_data(wrapped, algo)
        return clear

    def wrap_str(self, clear: "str") -> "str":
        return self.wrap_bytes(clear.encode("utf-8"))

    def unwrap_str(self, wrapped: "str") -> "str":
        return self.unwrap_bytes(wrapped).decode("utf-8")

class SecureStorage():
    def __init__(self):
        self.kdf_salt = b"\x25\xa9\x7b\xc5\x7a\x59\x0d\xa6"
        self.kek_cipher = None
        self.dek_cipher = None

    # KEK

    @property
    def has_kek(self):
        return bool(self.kek_cipher and self.kek_cipher.key)

    def set_raw_kek(self, kek):
        if self.kek_cipher:
            raise Exception("KEK already set")
        else:
            self.kek_cipher = CipherInstance(kek)

    def change_raw_kek(self, new_kek):
        if not self.kek_cipher:
            raise Exception("KEK not yet set")
        elif not self.dek_cipher:
            raise Exception("DEK not yet set")
        else:
            self.kek_cipher = CipherInstance(new_kek)

    def kdf(self, passwd):
        # The defaults of Crypto.Protocol.KDF.PBKDF2() were iter=1000, hash=SHA1.
        # TODO: Migration path to better parameters
        return PBKDF2(passwd.encode("utf-8"),
                      self.kdf_salt,
                      dkLen=16,
                      count=1000,
                      hmac_hash_module=SHA1)

    def set_password(self, passwd):
        return self.set_raw_kek(self.kdf(passwd))

    def change_password(self, passwd):
        return self.change_raw_kek(self.kdf(passwd))

    # DEK

    @property
    def has_dek(self):
        return bool(self.dek_cipher and self.dek_cipher.key)
{'fraction_non_alphanumeric': 0.12406417112299466, 'fraction_numerical': 0.017433155080213904, 'mean_word_length': 1.9424166142227817, 'pattern_counts': {'":': 30, '<': 11, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '184876', 'n_tokens_mistral': 3155, 'n_tokens_neox': 2939, 'n_words': 694}