#
# This file is part of pyasn1-modules software.
#
# Created by Russ Housley.
#
# Copyright (c) 2019, Vigil Security, LLC
# License: https://www.pysnmp.com/pyasn1/license.html
#
# PKCS #5: Password-Based Cryptography Specification, Version 2.1
#
# ASN.1 source from:
# https://www.rfc-editor.org/rfc/rfc8018.txt
#

from pyasn1.type import constraint, namedtype, namedval, univ

from pyasn1_modules import rfc3565, rfc5280

MAX = float("inf")


def _OID(*components):
    output = []
    for x in tuple(components):
        if isinstance(x, univ.ObjectIdentifier):
            output.extend(list(x))
        else:
            output.append(int(x))

    return univ.ObjectIdentifier(output)


# Import from RFC 3565

AES_IV = rfc3565.AES_IV


# Import from RFC 5280

AlgorithmIdentifier = rfc5280.AlgorithmIdentifier


# Basic object identifiers

nistAlgorithms = _OID(2, 16, 840, 1, 101, 3, 4)

aes = _OID(nistAlgorithms, 1)

oiw = _OID(1, 3, 14)

rsadsi = _OID(1, 2, 840, 113549)

pkcs = _OID(rsadsi, 1)

digestAlgorithm = _OID(rsadsi, 2)

encryptionAlgorithm = _OID(rsadsi, 3)

pkcs_5 = _OID(pkcs, 5)


# HMAC object identifiers

id_hmacWithSHA1 = _OID(digestAlgorithm, 7)

id_hmacWithSHA224 = _OID(digestAlgorithm, 8)

id_hmacWithSHA256 = _OID(digestAlgorithm, 9)

id_hmacWithSHA384 = _OID(digestAlgorithm, 10)

id_hmacWithSHA512 = _OID(digestAlgorithm, 11)

id_hmacWithSHA512_224 = _OID(digestAlgorithm, 12)

id_hmacWithSHA512_256 = _OID(digestAlgorithm, 13)


# PBES1 object identifiers

pbeWithMD2AndDES_CBC = _OID(pkcs_5, 1)

pbeWithMD2AndRC2_CBC = _OID(pkcs_5, 4)

pbeWithMD5AndDES_CBC = _OID(pkcs_5, 3)

pbeWithMD5AndRC2_CBC = _OID(pkcs_5, 6)

pbeWithSHA1AndDES_CBC = _OID(pkcs_5, 10)

pbeWithSHA1AndRC2_CBC = _OID(pkcs_5, 11)


# Supporting techniques object identifiers

desCBC = _OID(oiw, 3, 2, 7)

des_EDE3_CBC = _OID(encryptionAlgorithm, 7)

rc2CBC = _OID(encryptionAlgorithm, 2)

rc5_CBC_PAD = _OID(encryptionAlgorithm, 9)

aes128_CBC_PAD = _OID(aes, 2)

aes192_CBC_PAD = _OID(aes, 22)

aes256_CBC_PAD = _OID(aes, 42)


# PBES1


class PBEParameter(univ.Sequence):
    pass


PBEParameter.componentType = namedtype.NamedTypes(
    namedtype.NamedType(
        "salt",
        univ.OctetString().subtype(subtypeSpec=constraint.ValueSizeConstraint(8, 8)),
    ),
    namedtype.NamedType("iterationCount", univ.Integer()),
)


# PBES2

id_PBES2 = _OID(pkcs_5, 13)


class PBES2_params(univ.Sequence):
    pass


PBES2_params.componentType = namedtype.NamedTypes(
    namedtype.NamedType("keyDerivationFunc", AlgorithmIdentifier()),
    namedtype.NamedType("encryptionScheme", AlgorithmIdentifier()),
)


# PBMAC1

id_PBMAC1 = _OID(pkcs_5, 14)


class PBMAC1_params(univ.Sequence):
    pass


PBMAC1_params.componentType = namedtype.NamedTypes(
    namedtype.NamedType("keyDerivationFunc", AlgorithmIdentifier()),
    namedtype.NamedType("messageAuthScheme", AlgorithmIdentifier()),
)


# PBKDF2

id_PBKDF2 = _OID(pkcs_5, 12)


algid_hmacWithSHA1 = AlgorithmIdentifier()
algid_hmacWithSHA1["algorithm"] = id_hmacWithSHA1
algid_hmacWithSHA1["parameters"] = univ.Null("")


class PBKDF2_params(univ.Sequence):
    pass


PBKDF2_params.componentType = namedtype.NamedTypes(
    namedtype.NamedType(
        "salt",
        univ.Choice(
            componentType=namedtype.NamedTypes(
                namedtype.NamedType("specified", univ.OctetString()),
                namedtype.NamedType("otherSource", AlgorithmIdentifier()),
            )
        ),
    ),
    namedtype.NamedType(
        "iterationCount",
        univ.Integer().subtype(subtypeSpec=constraint.ValueRangeConstraint(1, MAX)),
    ),
    namedtype.OptionalNamedType(
        "keyLength",
        univ.Integer().subtype(subtypeSpec=constraint.ValueRangeConstraint(1, MAX)),
    ),
    namedtype.DefaultedNamedType("prf", algid_hmacWithSHA1),
)


# RC2 CBC algorithm parameter


class RC2_CBC_Parameter(univ.Sequence):
    pass


RC2_CBC_Parameter.componentType = namedtype.NamedTypes(
    namedtype.OptionalNamedType("rc2ParameterVersion", univ.Integer()),
    namedtype.NamedType(
        "iv",
        univ.OctetString().subtype(subtypeSpec=constraint.ValueSizeConstraint(8, 8)),
    ),
)


# RC5 CBC algorithm parameter


class RC5_CBC_Parameters(univ.Sequence):
    pass


RC5_CBC_Parameters.componentType = namedtype.NamedTypes(
    namedtype.NamedType(
        "version",
        univ.Integer(namedValues=namedval.NamedValues(("v1_0", 16))).subtype(
            subtypeSpec=constraint.SingleValueConstraint(16)
        ),
    ),
    namedtype.NamedType(
        "rounds",
        univ.Integer().subtype(subtypeSpec=constraint.ValueRangeConstraint(8, 127)),
    ),
    namedtype.NamedType(
        "blockSizeInBits",
        univ.Integer().subtype(subtypeSpec=constraint.SingleValueConstraint(64, 128)),
    ),
    namedtype.OptionalNamedType("iv", univ.OctetString()),
)


# Initialization Vector for AES: OCTET STRING (SIZE(16))


class AES_IV(univ.OctetString):
    pass


AES_IV.subtypeSpec = constraint.ValueSizeConstraint(16, 16)


# Initialization Vector for DES: OCTET STRING (SIZE(8))


class DES_IV(univ.OctetString):
    pass


DES_IV.subtypeSpec = constraint.ValueSizeConstraint(8, 8)


# Update the Algorithm Identifier map

_algorithmIdentifierMapUpdate = {
    # PBKDF2-PRFs
    id_hmacWithSHA1: univ.Null(),
    id_hmacWithSHA224: univ.Null(),
    id_hmacWithSHA256: univ.Null(),
    id_hmacWithSHA384: univ.Null(),
    id_hmacWithSHA512: univ.Null(),
    id_hmacWithSHA512_224: univ.Null(),
    id_hmacWithSHA512_256: univ.Null(),
    # PBES1Algorithms
    pbeWithMD2AndDES_CBC: PBEParameter(),
    pbeWithMD2AndRC2_CBC: PBEParameter(),
    pbeWithMD5AndDES_CBC: PBEParameter(),
    pbeWithMD5AndRC2_CBC: PBEParameter(),
    pbeWithSHA1AndDES_CBC: PBEParameter(),
    pbeWithSHA1AndRC2_CBC: PBEParameter(),
    # PBES2Algorithms
    id_PBES2: PBES2_params(),
    # PBES2-KDFs
    id_PBKDF2: PBKDF2_params(),
    # PBMAC1Algorithms
    id_PBMAC1: PBMAC1_params(),
    # SupportingAlgorithms
    desCBC: DES_IV(),
    des_EDE3_CBC: DES_IV(),
    rc2CBC: RC2_CBC_Parameter(),
    rc5_CBC_PAD: RC5_CBC_Parameters(),
    aes128_CBC_PAD: AES_IV(),
    aes192_CBC_PAD: AES_IV(),
    aes256_CBC_PAD: AES_IV(),
}

rfc5280.algorithmIdentifierMap.update(_algorithmIdentifierMapUpdate)