# Copyright (c) 2013-2019 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"""

import os

from hashlib import md5, sha1

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


# pylint: disable=bad-whitespace

_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')

# pylint: enable=bad-whitespace

_pkcs1_cipher = {}
_pkcs1_dek_name = {}

_pkcs8_handler = {}
_pkcs8_cipher_oid = {}

_pbes2_cipher = {}
_pbes2_cipher_oid = {}

_pbes2_kdf = {}
_pbes2_kdf_oid = {}

_pbes2_prf = {}
_pbes2_prf_oid = {}


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, block_size, key, iv):
        self._cipher = BasicCipher(cipher_name, key, iv)
        self._block_size = block_size

    def encrypt(self, data):
        """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):
        """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, passphrase, salt, count, key_size):
    """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, passphrase, salt, count, key_size, idx):
    """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, v):
        """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, passphrase, hash_alg, cipher_name):
    """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, passphrase, hash_alg, cipher_name):
    """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 = 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, cipher_name, key):
    """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, passphrase, default_key_size):
    """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, passphrase):
    """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, pkcs1_dek_name, cipher_name):
    """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, hash_name, pkcs8_cipher_oid,
                          handler, hash_alg, cipher_name):
    """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, pbes2_cipher_oid,
                          handler, cipher_name):
    """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, kdf_oid, handler, *args):
    """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, prf_oid):
    """Register a PBES2 pseudo-random function"""

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


def pkcs1_encrypt(data, pkcs1_cipher_name, passphrase):
    """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, pkcs1_dek_name, iv, passphrase):
    """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, pkcs8_cipher_name, hash_name, version, passphrase):
    """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 = (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, passphrase):
    """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')


# pylint: disable=bad-whitespace

_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)
