# Source code for asyncssh.pbe

    # Copyright (c) 2013-2021 by Ron Frederick <ronf@timeheart.net> and others.
    #
    # This program and the accompanying materials are made available under
    # the terms of the Eclipse Public License v2.0 which accompanies this
    # distribution and is available at:
    #
    #     http://www.eclipse.org/legal/epl-2.0/
    #
    # This program may also be made available under the following secondary
    # licenses when the conditions for such availability set forth in the
    # Eclipse Public License v2.0 are satisfied:
    #
    #    GNU General Public License, Version 2.0, or any later versions of
    #    that license
    #
    # SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-or-later
    #
    # Contributors:
    #     Ron Frederick - initial implementation, API, and documentation

    """Asymmetric key password based encryption functions"""

    from hashlib import md5, sha1
    import os
    from typing import Callable, Dict, Sequence, Tuple, Union

    from .asn1 import ASN1DecodeError, ObjectIdentifier, der_encode, der_decode
    from .crypto import BasicCipher, get_cipher_params, pbkdf2_hmac
    from .misc import BytesOrStr, HashType

    _Cipher = Union[BasicCipher, '_RFC1423Pad']

    _PKCS8CipherHandler = Callable[[object, BytesOrStr, Callable, str], _Cipher]
    _PKCS8Cipher = Tuple[_PKCS8CipherHandler, Callable, str]

    _PBES2CipherHandler = Callable[[Sequence, str, bytes], _Cipher]
    _PBES2Cipher = Tuple[_PBES2CipherHandler, str]

    _PBES2KDFHandler = Callable[[Sequence, BytesOrStr, int], bytes]
    _PBES2KDF = Tuple[_PBES2KDFHandler, Tuple[object, ...]]

    _ES1_MD5_DES    = ObjectIdentifier('1.2.840.113549.1.5.3')
    _ES1_SHA1_DES   = ObjectIdentifier('1.2.840.113549.1.5.10')

    _ES2            = ObjectIdentifier('1.2.840.113549.1.5.13')

    _P12_RC4_128    = ObjectIdentifier('1.2.840.113549.1.12.1.1')
    _P12_RC4_40     = ObjectIdentifier('1.2.840.113549.1.12.1.2')
    _P12_DES3       = ObjectIdentifier('1.2.840.113549.1.12.1.3')
    _P12_DES2       = ObjectIdentifier('1.2.840.113549.1.12.1.4')

    _ES2_CAST128    = ObjectIdentifier('1.2.840.113533.7.66.10')
    _ES2_DES3       = ObjectIdentifier('1.2.840.113549.3.7')
    _ES2_BF         = ObjectIdentifier('1.3.6.1.4.1.3029.1.2')
    _ES2_DES        = ObjectIdentifier('1.3.14.3.2.7')
    _ES2_AES128     = ObjectIdentifier('2.16.840.1.101.3.4.1.2')
    _ES2_AES192     = ObjectIdentifier('2.16.840.1.101.3.4.1.22')
    _ES2_AES256     = ObjectIdentifier('2.16.840.1.101.3.4.1.42')

    _ES2_PBKDF2     = ObjectIdentifier('1.2.840.113549.1.5.12')

    _ES2_SHA1       = ObjectIdentifier('1.2.840.113549.2.7')
    _ES2_SHA224     = ObjectIdentifier('1.2.840.113549.2.8')
    _ES2_SHA256     = ObjectIdentifier('1.2.840.113549.2.9')
    _ES2_SHA384     = ObjectIdentifier('1.2.840.113549.2.10')
    _ES2_SHA512     = ObjectIdentifier('1.2.840.113549.2.11')

    _pkcs1_cipher: Dict[bytes, str] = {}
    _pkcs1_dek_name: Dict[str, bytes] = {}

    _pkcs8_handler: Dict[ObjectIdentifier, _PKCS8Cipher] = {}
    _pkcs8_cipher_oid: Dict[Tuple[str, str], ObjectIdentifier] = {}

    _pbes2_cipher: Dict[ObjectIdentifier, _PBES2Cipher] = {}
    _pbes2_cipher_oid: Dict[str, ObjectIdentifier] = {}

    _pbes2_kdf: Dict[ObjectIdentifier, _PBES2KDF] = {}
    _pbes2_kdf_oid: Dict[str, ObjectIdentifier] = {}

    _pbes2_prf: Dict[ObjectIdentifier, str] = {}
    _pbes2_prf_oid: Dict[str, ObjectIdentifier] = {}

    class KeyEncryptionError(ValueError):
        """Key encryption error

           This exception is raised by key decryption functions when the data
           provided is not a valid encrypted private key.

        """

    class _RFC1423Pad:
        """RFC 1423 padding functions

           This class implements RFC 1423 padding for encryption and
           decryption of data by block ciphers. On encryption, the data is
           padded by between 1 and the cipher's block size number of bytes,
           with the padding value being equal to the length of the padding.

        """

        def __init__(self, cipher_name: str, block_size: int,
                     key: bytes, iv: bytes):
            self._cipher = BasicCipher(cipher_name, key, iv)
            self._block_size = block_size

        def encrypt(self, data: bytes) -> bytes:
            """Pad data before encrypting it"""

            pad = self._block_size - (len(data) % self._block_size)
            data += pad * bytes((pad,))
            return self._cipher.encrypt(data)

        def decrypt(self, data: bytes) -> bytes:
            """Remove padding from data after decrypting it"""

            data = self._cipher.decrypt(data)

            if data:
                pad = data[-1]
                if (1 <= pad <= self._block_size and
                        data[-pad:] == pad * bytes((pad,))):
                    return data[:-pad]

            raise KeyEncryptionError('Unable to decrypt key')

    def _pbkdf1(hash_alg: HashType, passphrase: BytesOrStr, salt: bytes,
                count: int, key_size: int) -> bytes:
        """PKCS#5 v1.5 key derivation function for password-based encryption

           This function implements the PKCS#5 v1.5 algorithm for deriving
           an encryption key from a passphrase and salt.

           The standard PBKDF1 function cannot generate more key bytes than
           the hash digest size, but 3DES uses a modified form of it which
           calls PBKDF1 recursively on the result to generate more key data.
           Support for this is implemented here.

        """

        if isinstance(passphrase, str):
            passphrase = passphrase.encode('utf-8')

        key = passphrase + salt
        for _ in range(count):
            key = hash_alg(key).digest()

        if len(key) <= key_size:
            return key + _pbkdf1(hash_alg, key + passphrase, salt, count,
                                 key_size - len(key))
        else:
            return key[:key_size]

    def _pbkdf_p12(hash_alg: HashType, passphrase: BytesOrStr, salt: bytes,
                   count: int, key_size: int, idx: int) -> bytes:
        """PKCS#12 key derivation function for password-based encryption

           This function implements the PKCS#12 algorithm for deriving an
           encryption key from a passphrase and salt.

        """

        def _make_block(data: bytes, v: int) -> bytes:
            """Make a block a multiple of v bytes long by repeating data"""

            l = len(data)
            size = ((l + v - 1) // v) * v
            return (((size + l - 1) // l) * data)[:size]

        v = hash_alg().block_size
        D = v * bytes((idx,))

        if isinstance(passphrase, str):
            passphrase = passphrase.encode('utf-16be')

        I = bytearray(_make_block(salt, v) + _make_block(passphrase + b'\0\0', v))

        key = b''
        while len(key) < key_size:
            A = D + I
            for i in range(count):
                A = hash_alg(A).digest()

            B = int.from_bytes(_make_block(A, v), 'big')
            for i in range(0, len(I), v):
                x = (int.from_bytes(I[i:i+v], 'big') + B + 1) % (1 << v*8)
                I[i:i+v] = x.to_bytes(v, 'big')

            key += A

        return key[:key_size]

    def _pbes1(params: object, passphrase: BytesOrStr, hash_alg: HashType,
               cipher_name: str) -> _Cipher:
        """PKCS#5 v1.5 cipher selection function for password-based encryption

           This function implements the PKCS#5 v1.5 algorithm for password-based
           encryption. It returns a cipher object which can be used to encrypt
           or decrypt data based on the specified encryption parameters,
           passphrase, and salt.

        """

        if (not isinstance(params, tuple) or len(params) != 2 or
                not isinstance(params[0], bytes) or
                not isinstance(params[1], int)):
            raise KeyEncryptionError('Invalid PBES1 encryption parameters')

        salt, count = params

        key_size, iv_size, block_size = get_cipher_params(cipher_name)
        key = _pbkdf1(hash_alg, passphrase, salt, count, key_size + iv_size)
        key, iv = key[:key_size], key[key_size:]

        return _RFC1423Pad(cipher_name, block_size, key, iv)

    def _pbe_p12(params: object, passphrase: BytesOrStr, hash_alg: HashType,
                 cipher_name: str) -> _Cipher:
        """PKCS#12 cipher selection function for password-based encryption

           This function implements the PKCS#12 algorithm for password-based
           encryption. It returns a cipher object which can be used to encrypt
           or decrypt data based on the specified encryption parameters,
           passphrase, and salt.

        """

        if (not isinstance(params, tuple) or len(params) != 2 or
                not isinstance(params[0], bytes) or not params[0] or
                not isinstance(params[1], int) or params[1] == 0):
            raise KeyEncryptionError('Invalid PBES1 PKCS#12 encryption parameters')

        salt, count = params

        key_size, iv_size, block_size = get_cipher_params(cipher_name)
        key = _pbkdf_p12(hash_alg, passphrase, salt, count, key_size, 1)

        if block_size == 1:
            cipher: _Cipher = BasicCipher(cipher_name, key, b'')
        else:
            iv = _pbkdf_p12(hash_alg, passphrase, salt, count, iv_size, 2)
            cipher = _RFC1423Pad(cipher_name, block_size, key, iv)

        return cipher

    def _pbes2_iv(enc_params: Sequence, cipher_name: str, key: bytes) -> _Cipher:
        """PKCS#5 v2.0 handler for PBES2 ciphers with an IV as a parameter

           This function returns the appropriate cipher object to use for
           PBES2 encryption for ciphers that have only an IV as an encryption
           parameter.

        """

        _, iv_size, block_size = get_cipher_params(cipher_name)

        if len(enc_params) != 1 or not isinstance(enc_params[0], bytes):
            raise KeyEncryptionError('Invalid PBES2 encryption parameters')

        if len(enc_params[0]) != iv_size:
            raise KeyEncryptionError('Invalid length IV for PBES2 encryption')

        return _RFC1423Pad(cipher_name, block_size, key, enc_params[0])

    def _pbes2_pbkdf2(kdf_params: Sequence, passphrase: BytesOrStr,
                      default_key_size: int) -> bytes:
        """PKCS#5 v2.0 handler for PBKDF2 key derivation

           This function parses the PBKDF2 arguments from a PKCS#8 encrypted key
           and returns the encryption key to use for encryption.

        """

        if (len(kdf_params) != 1 or not isinstance(kdf_params[0], tuple) or
                len(kdf_params[0]) < 2):
            raise KeyEncryptionError('Invalid PBES2 key derivation parameters')

        kdf_params = list(kdf_params[0])

        if (not isinstance(kdf_params[0], bytes) or
                not isinstance(kdf_params[1], int)):
            raise KeyEncryptionError('Invalid PBES2 key derivation parameters')

        salt = kdf_params.pop(0)
        count = kdf_params.pop(0)

        if kdf_params and isinstance(kdf_params[0], int):
            key_size = kdf_params.pop(0)    # pragma: no cover, used only by RC2
        else:
            key_size = default_key_size

        if kdf_params:
            if (isinstance(kdf_params[0], tuple) and len(kdf_params[0]) == 2 and
                    isinstance(kdf_params[0][0], ObjectIdentifier)):
                prf_alg = kdf_params[0][0]
                if prf_alg in _pbes2_prf:
                    hash_name = _pbes2_prf[prf_alg]
                else:
                    raise KeyEncryptionError('Unknown PBES2 pseudo-random '
                                             'function')
            else:
                raise KeyEncryptionError('Invalid PBES2 pseudo-random function '
                                         'parameters')
        else:
            hash_name = 'sha1'

        if isinstance(passphrase, str):
            passphrase = passphrase.encode('utf-8')

        return pbkdf2_hmac(hash_name, passphrase, salt, count, key_size)

    def _pbes2(params: object, passphrase: BytesOrStr) -> _Cipher:
        """PKCS#5 v2.0 cipher selection function for password-based encryption

           This function implements the PKCS#5 v2.0 algorithm for password-based
           encryption. It returns a cipher object which can be used to encrypt
           or decrypt data based on the specified encryption parameters and
           passphrase.

        """

        if (not isinstance(params, tuple) or len(params) != 2 or
                not isinstance(params[0], tuple) or len(params[0]) < 1 or
                not isinstance(params[1], tuple) or len(params[1]) < 1):
            raise KeyEncryptionError('Invalid PBES2 encryption parameters')

        kdf_params = list(params[0])

        kdf_alg = kdf_params.pop(0)
        if kdf_alg not in _pbes2_kdf:
            raise KeyEncryptionError('Unknown PBES2 key derivation function')

        enc_params = list(params[1])

        enc_alg = enc_params.pop(0)
        if enc_alg not in _pbes2_cipher:
            raise KeyEncryptionError('Unknown PBES2 encryption algorithm')

        kdf_handler, kdf_args = _pbes2_kdf[kdf_alg]
        enc_handler, cipher_name = _pbes2_cipher[enc_alg]
        default_key_size, _, _ = get_cipher_params(cipher_name)

        key = kdf_handler(kdf_params, passphrase, default_key_size, *kdf_args)
        return enc_handler(enc_params, cipher_name, key)

    def register_pkcs1_cipher(pkcs1_cipher_name: str, pkcs1_dek_name: bytes,
                              cipher_name: str) -> None:
        """Register a cipher used for PKCS#1 private key encryption"""

        _pkcs1_cipher[pkcs1_dek_name] = cipher_name
        _pkcs1_dek_name[pkcs1_cipher_name] = pkcs1_dek_name

    def register_pkcs8_cipher(pkcs8_cipher_name: str, hash_name: str,
                              pkcs8_cipher_oid: ObjectIdentifier,
                              handler: _PKCS8CipherHandler, hash_alg: HashType,
                              cipher_name: str) -> None:
        """Register a cipher used for PKCS#8 private key encryption"""

        _pkcs8_handler[pkcs8_cipher_oid] = (handler, hash_alg, cipher_name)
        _pkcs8_cipher_oid[pkcs8_cipher_name, hash_name] = pkcs8_cipher_oid

    def register_pbes2_cipher(pbes2_cipher_name: str,
                              pbes2_cipher_oid: ObjectIdentifier,
                              handler: _PBES2CipherHandler,
                              cipher_name: str) -> None:
        """Register a PBES2 encryption algorithm"""

        _pbes2_cipher[pbes2_cipher_oid] = (handler, cipher_name)
        _pbes2_cipher_oid[pbes2_cipher_name] = pbes2_cipher_oid

    def register_pbes2_kdf(kdf_name: str, kdf_oid: ObjectIdentifier,
                           handler: _PBES2KDFHandler, *args: object) -> None:
        """Register a PBES2 key derivation function"""

        _pbes2_kdf[kdf_oid] = (handler, args)
        _pbes2_kdf_oid[kdf_name] = kdf_oid

    def register_pbes2_prf(hash_name: str, prf_oid: ObjectIdentifier) -> None:
        """Register a PBES2 pseudo-random function"""

        _pbes2_prf[prf_oid] = hash_name
        _pbes2_prf_oid[hash_name] = prf_oid

    def pkcs1_encrypt(data: bytes, pkcs1_cipher_name: str,
                      passphrase: BytesOrStr) -> Tuple[bytes, bytes, bytes]:
        """Encrypt PKCS#1 key data

           This function encrypts PKCS#1 key data using the specified cipher
           and passphrase. Available ciphers include:

               aes128-cbc, aes192-cbc, aes256-cbc, des-cbc, des3-cbc

        """

        if pkcs1_cipher_name in _pkcs1_dek_name:
            pkcs1_dek_name = _pkcs1_dek_name[pkcs1_cipher_name]
            cipher_name = _pkcs1_cipher[pkcs1_dek_name]
            key_size, iv_size, block_size = get_cipher_params(cipher_name)

            iv = os.urandom(iv_size)
            key = _pbkdf1(md5, passphrase, iv[:8], 1, key_size)

            cipher = _RFC1423Pad(cipher_name, block_size, key, iv)
            return pkcs1_dek_name, iv, cipher.encrypt(data)
        else:
            raise KeyEncryptionError('Unknown PKCS#1 encryption algorithm')

    def pkcs1_decrypt(data: bytes, pkcs1_dek_name: bytes, iv: bytes,
                      passphrase: BytesOrStr) -> bytes:
        """Decrypt PKCS#1 key data

           This function decrypts PKCS#1 key data using the specified algorithm,
           initialization vector, and passphrase. The algorithm name and IV
           should be taken from the PEM DEK-Info header.

        """

        if pkcs1_dek_name in _pkcs1_cipher:
            cipher_name = _pkcs1_cipher[pkcs1_dek_name]
            key_size, _, block_size = get_cipher_params(cipher_name)
            key = _pbkdf1(md5, passphrase, iv[:8], 1, key_size)

            cipher = _RFC1423Pad(cipher_name, block_size, key, iv)
            return cipher.decrypt(data)
        else:
            raise KeyEncryptionError('Unknown PKCS#1 encryption algorithm')

    def pkcs8_encrypt(data: bytes, pkcs8_cipher_name: str, hash_name: str,
                      version: int, passphrase: BytesOrStr) -> bytes:
        """Encrypt PKCS#8 key data

           This function encrypts PKCS#8 key data using the specified cipher,
           hash, encryption version, and passphrase.

           Available ciphers include:

               aes128-cbc, aes192-cbc, aes256-cbc, blowfish-cbc, cast128-cbc,
               des-cbc, des2-cbc, des3-cbc, rc4-40, and rc4-128

           Available hashes include:

               md5, sha1, sha256, sha384, sha512

           Available versions include 1 for PBES1 and 2 for PBES2.

           Only some combinations of cipher, hash, and version are supported.

        """

        if version == 1 and (pkcs8_cipher_name, hash_name) in _pkcs8_cipher_oid:
            pkcs8_cipher_oid = _pkcs8_cipher_oid[pkcs8_cipher_name, hash_name]
            handler, hash_alg, cipher_name = _pkcs8_handler[pkcs8_cipher_oid]

            alg = pkcs8_cipher_oid
            params: object = (os.urandom(8), 2048)
            cipher = handler(params, passphrase, hash_alg, cipher_name)
        elif version == 2 and pkcs8_cipher_name in _pbes2_cipher_oid:
            pbes2_cipher_oid = _pbes2_cipher_oid[pkcs8_cipher_name]
            _, cipher_name = _pbes2_cipher[pbes2_cipher_oid]
            _, iv_size, _ = get_cipher_params(cipher_name)

            kdf_params = [os.urandom(8), 2048]
            iv = os.urandom(iv_size)
            enc_params = (pbes2_cipher_oid, iv)

            if hash_name != 'sha1':
                if hash_name in _pbes2_prf_oid:
                    kdf_params.append((_pbes2_prf_oid[hash_name], None))
                else:
                    raise KeyEncryptionError('Unknown PBES2 hash function')

            alg = _ES2
            params = ((_ES2_PBKDF2, tuple(kdf_params)), enc_params)
            cipher = _pbes2(params, passphrase)
        else:
            raise KeyEncryptionError('Unknown PKCS#8 encryption algorithm')

        return der_encode(((alg, params), cipher.encrypt(data)))

    def pkcs8_decrypt(key_data: object, passphrase: BytesOrStr) -> object:
        """Decrypt PKCS#8 key data

           This function decrypts key data in PKCS#8 EncryptedPrivateKeyInfo
           format using the specified passphrase.

        """

        if not isinstance(key_data, tuple) or len(key_data) != 2:
            raise KeyEncryptionError('Invalid PKCS#8 encrypted key format')

        alg_params, data = key_data

        if (not isinstance(alg_params, tuple) or len(alg_params) != 2 or
                not isinstance(data, bytes)):
            raise KeyEncryptionError('Invalid PKCS#8 encrypted key format')

        alg, params = alg_params

        if alg == _ES2:
            cipher = _pbes2(params, passphrase)
        elif alg in _pkcs8_handler:
            handler, hash_alg, cipher_name = _pkcs8_handler[alg]
            cipher = handler(params, passphrase, hash_alg, cipher_name)
        else:
            raise KeyEncryptionError('Unknown PKCS#8 encryption algorithm')

        try:
            return der_decode(cipher.decrypt(data))
        except (ASN1DecodeError, UnicodeDecodeError):
            raise KeyEncryptionError('Invalid PKCS#8 encrypted key data') from None

    _pkcs1_cipher_list = (
        ('aes128-cbc', b'AES-128-CBC',  'aes128-cbc'),
        ('aes192-cbc', b'AES-192-CBC',  'aes192-cbc'),
        ('aes256-cbc', b'AES-256-CBC',  'aes256-cbc'),
        ('des-cbc',    b'DES-CBC',      'des-cbc'),
        ('des3-cbc',   b'DES-EDE3-CBC', 'des3-cbc')
    )

    _pkcs8_cipher_list = (
        ('des-cbc', 'md5',  _ES1_MD5_DES,  _pbes1,   md5,  'des-cbc'),
        ('des-cbc', 'sha1', _ES1_SHA1_DES, _pbes1,   sha1, 'des-cbc'),

        ('des2-cbc','sha1', _P12_DES2,     _pbe_p12, sha1, 'des2-cbc'),
        ('des3-cbc','sha1', _P12_DES3,     _pbe_p12, sha1, 'des3-cbc'),
        ('rc4-40',  'sha1', _P12_RC4_40,   _pbe_p12, sha1, 'arcfour40'),
        ('rc4-128', 'sha1', _P12_RC4_128,  _pbe_p12, sha1, 'arcfour')
    )

    _pbes2_cipher_list = (
        ('aes128-cbc',   _ES2_AES128,  _pbes2_iv,  'aes128-cbc'),
        ('aes192-cbc',   _ES2_AES192,  _pbes2_iv,  'aes192-cbc'),
        ('aes256-cbc',   _ES2_AES256,  _pbes2_iv,  'aes256-cbc'),
        ('blowfish-cbc', _ES2_BF,      _pbes2_iv,  'blowfish-cbc'),
        ('cast128-cbc',  _ES2_CAST128, _pbes2_iv,  'cast128-cbc'),
        ('des-cbc',      _ES2_DES,     _pbes2_iv,  'des-cbc'),
        ('des3-cbc',     _ES2_DES3,    _pbes2_iv,  'des3-cbc')
    )

    _pbes2_kdf_list = (
        ('pbkdf2', _ES2_PBKDF2, _pbes2_pbkdf2),
    )

    _pbes2_prf_list = (
        ('sha1',   _ES2_SHA1),
        ('sha224', _ES2_SHA224),
        ('sha256', _ES2_SHA256),
        ('sha384', _ES2_SHA384),
        ('sha512', _ES2_SHA512)
    )

    for _pkcs1_cipher_args in _pkcs1_cipher_list:
        register_pkcs1_cipher(*_pkcs1_cipher_args)

    for _pkcs8_cipher_args in _pkcs8_cipher_list:
        register_pkcs8_cipher(*_pkcs8_cipher_args)

    for _pbes2_cipher_args in _pbes2_cipher_list:
        register_pbes2_cipher(*_pbes2_cipher_args)

    for _pbes2_kdf_args in _pbes2_kdf_list:
        register_pbes2_kdf(*_pbes2_kdf_args)

    for _pbes2_prf_args in _pbes2_prf_list:
        register_pbes2_prf(*_pbes2_prf_args)