codekingpro/portable-devtools

1# This file is dual licensed under the terms of the Apache License, Version2# 2.0, and the BSD License. See the LICENSE file in the root of this repository3# for complete details.4 5from __future__ import annotations6 7import abc8 9from cryptography import utils10from cryptography.hazmat.bindings._rust import Encoding as Encoding11from cryptography.hazmat.bindings._rust import (12    ParameterFormat as ParameterFormat,13)14from cryptography.hazmat.bindings._rust import PrivateFormat as PrivateFormat15from cryptography.hazmat.bindings._rust import PublicFormat as PublicFormat16from cryptography.hazmat.primitives.hashes import HashAlgorithm17 18# This exists to break an import cycle. These classes are normally accessible19# from the serialization module.20 21 22class PBES(utils.Enum):23    PBESv1SHA1And3KeyTripleDESCBC = "PBESv1 using SHA1 and 3-Key TripleDES"24    PBESv2SHA256AndAES256CBC = "PBESv2 using SHA256 PBKDF2 and AES256 CBC"25 26 27class KeySerializationEncryption(metaclass=abc.ABCMeta):28    pass29 30 31class BestAvailableEncryption(KeySerializationEncryption):32    def __init__(self, password: bytes):33        if not isinstance(password, bytes) or len(password) == 0:34            raise ValueError("Password must be 1 or more bytes.")35 36        self.password = password37 38 39class NoEncryption(KeySerializationEncryption):40    pass41 42 43class KeySerializationEncryptionBuilder:44    def __init__(45        self,46        format: PrivateFormat,47        *,48        _kdf_rounds: int | None = None,49        _hmac_hash: HashAlgorithm | None = None,50        _key_cert_algorithm: PBES | None = None,51    ) -> None:52        self._format = format53 54        self._kdf_rounds = _kdf_rounds55        self._hmac_hash = _hmac_hash56        self._key_cert_algorithm = _key_cert_algorithm57 58    def kdf_rounds(self, rounds: int) -> KeySerializationEncryptionBuilder:59        if self._kdf_rounds is not None:60            raise ValueError("kdf_rounds already set")61 62        if not isinstance(rounds, int):63            raise TypeError("kdf_rounds must be an integer")64 65        if rounds < 1:66            raise ValueError("kdf_rounds must be a positive integer")67 68        return KeySerializationEncryptionBuilder(69            self._format,70            _kdf_rounds=rounds,71            _hmac_hash=self._hmac_hash,72            _key_cert_algorithm=self._key_cert_algorithm,73        )74 75    def hmac_hash(76        self, algorithm: HashAlgorithm77    ) -> KeySerializationEncryptionBuilder:78        if self._format is not PrivateFormat.PKCS12:79            raise TypeError(80                "hmac_hash only supported with PrivateFormat.PKCS12"81            )82 83        if self._hmac_hash is not None:84            raise ValueError("hmac_hash already set")85        return KeySerializationEncryptionBuilder(86            self._format,87            _kdf_rounds=self._kdf_rounds,88            _hmac_hash=algorithm,89            _key_cert_algorithm=self._key_cert_algorithm,90        )91 92    def key_cert_algorithm(93        self, algorithm: PBES94    ) -> KeySerializationEncryptionBuilder:95        if self._format is not PrivateFormat.PKCS12:96            raise TypeError(97                "key_cert_algorithm only supported with PrivateFormat.PKCS12"98            )99        if self._key_cert_algorithm is not None:100            raise ValueError("key_cert_algorithm already set")101        return KeySerializationEncryptionBuilder(102            self._format,103            _kdf_rounds=self._kdf_rounds,104            _hmac_hash=self._hmac_hash,105            _key_cert_algorithm=algorithm,106        )107 108    def build(self, password: bytes) -> KeySerializationEncryption:109        if not isinstance(password, bytes) or len(password) == 0:110            raise ValueError("Password must be 1 or more bytes.")111 112        return _KeySerializationEncryption(113            self._format,114            password,115            kdf_rounds=self._kdf_rounds,116            hmac_hash=self._hmac_hash,117            key_cert_algorithm=self._key_cert_algorithm,118        )119 120 121class _KeySerializationEncryption(KeySerializationEncryption):122    def __init__(123        self,124        format: PrivateFormat,125        password: bytes,126        *,127        kdf_rounds: int | None,128        hmac_hash: HashAlgorithm | None,129        key_cert_algorithm: PBES | None,130    ):131        self._format = format132        self.password = password133 134        self._kdf_rounds = kdf_rounds135        self._hmac_hash = hmac_hash136        self._key_cert_algorithm = key_cert_algorithm137