#
# This file is part of pyasn1-modules software.
#
# Created by Russ Housley
# Copyright (c) 2019, Vigil Security, LLC
# License: http://snmplabs.com/pyasn1/license.html
#
import sys

from pyasn1.codec.der.decoder import decode as der_decode
from pyasn1.codec.der.encoder import encode as der_encode

from pyasn1.type import univ

from pyasn1_modules import pem
from pyasn1_modules import rfc2985
from pyasn1_modules import rfc5280
from pyasn1_modules import rfc5652
from pyasn1_modules import rfc7292


try:
    import unittest2 as unittest

except ImportError:
    import unittest


class PKCS9AttrsTestCase(unittest.TestCase):
    pem_text = """\

HktMK+isIjxOTk4yJTOOAgIH0A==
"""

    def setUp(self):
        self.asn1Spec = rfc2985.AttributeSet()

    def testDerCodec(self):
        substrate = pem.readBase64fromText(self.pem_text)
        asn1Object, rest = der_decode(substrate, asn1Spec=self.asn1Spec)
        assert not rest
        assert asn1Object.prettyPrint()
        assert der_encode(asn1Object) == substrate

        openTypesMap = {
            rfc2985.pkcs_9_at_smimeCapabilities: rfc2985.SMIMECapabilities(),
        }
        openTypesMap.update(rfc5280.certificateAttributesMap)
        openTypesMap.update(rfc5652.cmsAttributesMap)

        for attr in asn1Object:
            assert attr['type'] in openTypesMap.keys()
            av, rest = der_decode(attr['values'][0],
                asn1Spec=openTypesMap[attr['type']])
            assert not rest
            assert av.prettyPrint()
            assert der_encode(av) == attr['values'][0]

            if attr['type'] == rfc2985.pkcs_9_at_userPKCS12:
                assert av['version'] == univ.Integer(3)
                assert av['authSafe']['contentType'] == rfc5652.id_data
                outdata, rest = der_decode(av['authSafe']['content'],
                    asn1Spec=univ.OctetString())
                assert not rest
                authsafe, rest = der_decode(outdata,
                    asn1Spec=rfc7292.AuthenticatedSafe())
                assert not rest

                for ci in authsafe:
                    assert ci['contentType'] == rfc5652.id_data
                    indata, rest = der_decode(ci['content'],
                        asn1Spec=univ.OctetString())
                    assert not rest
                    sc, rest = der_decode(indata,
                        asn1Spec=rfc7292.SafeContents())
                    assert not rest

                    for sb in sc:
                        if sb['bagId'] in rfc7292.pkcs12BagTypeMap:
                            bv, rest = der_decode(sb['bagValue'],
                                asn1Spec=rfc7292.pkcs12BagTypeMap[sb['bagId']])
                            assert not rest

                            for bagattr in sb['bagAttributes']:
                                if bagattr['attrType'] in openTypesMap:
                                    inav, rest = der_decode(bagattr['attrValues'][0],
                                        asn1Spec=openTypesMap[bagattr['attrType']])
                                    assert not rest

                                    if bagattr['attrType'] == rfc2985.pkcs_9_at_friendlyName:
                                        assert inav == "3f71af65-1687-444a-9f46-c8be194c3e8e"

                                    if bagattr['attrType'] == rfc2985.pkcs_9_at_localKeyId:
                                        assert inav == univ.OctetString(hexValue='01000000')

            if attr['type'] == rfc2985.pkcs_9_at_pkcs7PDU:
                ci, rest = der_decode(attr['values'][0],
                    asn1Spec=rfc5652.ContentInfo())
                assert not rest
                assert ci['contentType'] == rfc5652.id_signedData

                sd, rest = der_decode(ci['content'],
                    asn1Spec=rfc5652.SignedData())
                assert not rest
                assert sd['version'] == 1
        
                for si in sd['signerInfos']:
                    assert si['version'] == 1

                    for siattr in si['signedAttrs']:
                        if siattr['attrType'] in openTypesMap:
                            siav, rest = der_decode(siattr['attrValues'][0],
                                asn1Spec=openTypesMap[siattr['attrType']])
                            assert not rest

                            if siattr['attrType'] == rfc2985.pkcs_9_at_contentType:
                                assert siav == rfc5652.id_data

                            if siattr['attrType'] == rfc2985.pkcs_9_at_messageDigest:
                                assert siav.prettyPrint()[2:10] == 'b6e422a4'

                            if siattr['attrType'] == rfc2985.pkcs_9_at_signingTime:
                                assert siav['utcTime'] == '190529182319Z'

                for choices in sd['certificates']:
                    for rdn in choices[0]['tbsCertificate']['subject']['rdnSequence']:
                        if rdn[0]['type'] in openTypesMap:
                            nv, rest = der_decode(rdn[0]['value'],
                                 asn1Spec=openTypesMap[rdn[0]['type']])
                            assert not rest

                            if rdn[0]['type'] == rfc2985.pkcs_9_at_emailAddress:
                                assert nv == 'alice@example.com'

    def testOpenTypes(self):
        openTypesMap = {
            rfc2985.pkcs_9_at_smimeCapabilities: rfc2985.SMIMECapabilities(),
        }
        openTypesMap.update(rfc5280.certificateAttributesMap)
        openTypesMap.update(rfc5652.cmsAttributesMap)

        substrate = pem.readBase64fromText(self.pem_text)
        asn1Object, rest = der_decode(substrate,
            asn1Spec=self.asn1Spec,
            openTypes=openTypesMap,
            decodeOpenTypes=True)
        assert not rest
        assert asn1Object.prettyPrint()
        assert der_encode(asn1Object) == substrate

        for attr in asn1Object:
            assert attr['type'] in openTypesMap.keys()

            if attr['type'] == rfc2985.pkcs_9_at_userPKCS12:
                assert attr['values'][0]['version'] == univ.Integer(3)
                assert attr['values'][0]['authSafe']['contentType'] == rfc5652.id_data
                authsafe, rest = der_decode(attr['values'][0]['authSafe']['content'],
                    asn1Spec=rfc7292.AuthenticatedSafe())
                assert not rest

                for ci in authsafe:
                    assert ci['contentType'] == rfc5652.id_data
                    indata, rest = der_decode(ci['content'],
                       asn1Spec=univ.OctetString())
                    assert not rest

                    sc, rest = der_decode(indata,
                        asn1Spec=rfc7292.SafeContents(),
                        decodeOpenTypes=True)
                    assert not rest

                    for sb in sc:
                        if sb['bagId'] in rfc7292.pkcs12BagTypeMap:
                            for bagattr in sb['bagAttributes']:
                                if bagattr['attrType'] in openTypesMap:

                                    if bagattr['attrType'] == rfc2985.pkcs_9_at_friendlyName:
                                        assert bagattr['attrValues'][0] == "3f71af65-1687-444a-9f46-c8be194c3e8e"

                                    if bagattr['attrType'] == rfc2985.pkcs_9_at_localKeyId:
                                        assert bagattr['attrValues'][0] == univ.OctetString(hexValue='01000000')

            if attr['type'] == rfc2985.pkcs_9_at_pkcs7PDU:
                assert attr['values'][0]['contentType'] == rfc5652.id_signedData
                assert attr['values'][0]['content']['version'] == 1
        
                for si in attr['values'][0]['content']['signerInfos']:
                    assert si['version'] == 1

                    for siattr in si['signedAttrs']:
                        if siattr['attrType'] in openTypesMap:
 
                            if siattr['attrType'] == rfc2985.pkcs_9_at_contentType:
                                assert siattr['attrValues'][0] == rfc5652.id_data

                            if siattr['attrType'] == rfc2985.pkcs_9_at_messageDigest:
                                assert siattr['attrValues'][0].prettyPrint()[2:10] == 'b6e422a4'

                            if siattr['attrType'] == rfc2985.pkcs_9_at_signingTime:
                                assert siattr['attrValues'][0]['utcTime'] == '190529182319Z'

                for choices in attr['values'][0]['content']['certificates']:
                    for rdn in choices[0]['tbsCertificate']['subject']['rdnSequence']:
                        if rdn[0]['type'] in openTypesMap:
                            if rdn[0]['type'] == rfc2985.pkcs_9_at_emailAddress:
                                assert rdn[0]['value'] == 'alice@example.com'


suite = unittest.TestLoader().loadTestsFromModule(sys.modules[__name__])

if __name__ == '__main__':
    unittest.TextTestRunner(verbosity=2).run(suite)
