#
# This file is part of pyasn1-modules software.
#
# Copyright (c) 2019, Vigil Security, LLC
# License: https://www.pysnmp.com/pyasn1/license.html
#
import sys
import unittest

from pyasn1.codec.der.decoder import decode as der_decoder
from pyasn1.codec.der.encoder import encode as der_encoder
from pyasn1.type import univ

from pyasn1_modules import pem, rfc4357, rfc5280, rfc5652


class SignedTestCase(unittest.TestCase):
    signed_pem_text = """\
MIIBKAYJKoZIhvcNAQcCoIIBGTCCARUCAQExDDAKBgYqhQMCAgkFADAbBgkqhkiG
9w0BBwGgDgQMc2FtcGxlIHRleHQKMYHkMIHhAgEBMIGBMG0xHzAdBgNVBAMMFkdv
c3RSMzQxMC0yMDAxIGV4YW1wbGUxEjAQBgNVBAoMCUNyeXB0b1BybzELMAkGA1UE
BhMCUlUxKTAnBgkqhkiG9w0BCQEWGkdvc3RSMzQxMC0yMDAxQGV4YW1wbGUuY29t
AhAr9cYewhG9F8fc1GJmtC4hMAoGBiqFAwICCQUAMAoGBiqFAwICEwUABEDAw0LZ
P4/+JRERiHe/icPbg0IE1iD5aCqZ9v4wO+T0yPjVtNr74caRZzQfvKZ6DRJ7/RAl
xlHbjbL0jHF+7XKp
"""

    def setUp(self):
        self.asn1Spec = rfc5652.ContentInfo()

    def testDerCodec(self):
        substrate = pem.readBase64fromText(self.signed_pem_text)
        asn1Object, rest = der_decoder(substrate, asn1Spec=self.asn1Spec)
        self.assertFalse(rest)
        self.assertTrue(asn1Object.prettyPrint())
        self.assertEqual(substrate, der_encoder(asn1Object))
        self.assertEqual(rfc5652.id_signedData, asn1Object["contentType"])

        sd, rest = der_decoder(asn1Object["content"], asn1Spec=rfc5652.SignedData())
        self.assertFalse(rest)
        self.assertTrue(sd.prettyPrint())
        self.assertEqual(asn1Object["content"], der_encoder(sd))

        encoded_null = der_encoder(univ.Null(""))

        si = sd["signerInfos"][0]
        self.assertEqual(rfc4357.id_GostR3411_94, si["digestAlgorithm"]["algorithm"])
        self.assertEqual(encoded_null, si["digestAlgorithm"]["parameters"])

        self.assertEqual(
            rfc4357.id_GostR3410_2001, si["signatureAlgorithm"]["algorithm"]
        )
        self.assertEqual(encoded_null, si["signatureAlgorithm"]["parameters"])
        self.assertEqual(64, len(si["signature"]))

    def testOpenTypes(self):
        substrate = pem.readBase64fromText(self.signed_pem_text)
        asn1Object, rest = der_decoder(
            substrate, asn1Spec=self.asn1Spec, decodeOpenTypes=True
        )
        self.assertFalse(rest)
        self.assertTrue(asn1Object.prettyPrint())
        self.assertEqual(substrate, der_encoder(asn1Object))
        self.assertEqual(rfc5652.id_signedData, asn1Object["contentType"])

        si = asn1Object["content"]["signerInfos"][0]
        self.assertEqual(rfc4357.id_GostR3411_94, si["digestAlgorithm"]["algorithm"])
        self.assertEqual(univ.Null(""), si["digestAlgorithm"]["parameters"])

        self.assertEqual(
            rfc4357.id_GostR3410_2001, si["signatureAlgorithm"]["algorithm"]
        )
        self.assertEqual(univ.Null(""), si["signatureAlgorithm"]["parameters"])

        self.assertEqual(64, len(si["signature"]))


class KeyAgreeTestCase(unittest.TestCase):
    keyagree_pem_text = """\
MIIBpAYJKoZIhvcNAQcDoIIBlTCCAZECAQIxggFQoYIBTAIBA6BloWMwHAYGKoUD
AgITMBIGByqFAwICJAAGByqFAwICHgEDQwAEQLNVOfRngZcrpcTZhB8n+4HtCDLm
mtTyAHi4/4Nk6tIdsHg8ff4DwfQG5DvMFrnF9vYZNxwXuKCqx9GhlLOlNiChCgQI
L/D20YZLMoowHgYGKoUDAgJgMBQGByqFAwICDQAwCQYHKoUDAgIfATCBszCBsDCB
gTBtMR8wHQYDVQQDDBZHb3N0UjM0MTAtMjAwMSBleGFtcGxlMRIwEAYDVQQKDAlD
cnlwdG9Qcm8xCzAJBgNVBAYTAlJVMSkwJwYJKoZIhvcNAQkBFhpHb3N0UjM0MTAt
MjAwMUBleGFtcGxlLmNvbQIQK/XGHsIRvRfH3NRiZrQuIQQqMCgEIBajHOfOTukN
8ex0aQRoHsefOu24Ox8dSn75pdnLGdXoBAST/YZ+MDgGCSqGSIb3DQEHATAdBgYq
hQMCAhUwEwQItzXhegc1oh0GByqFAwICHwGADDmxivS/qeJlJbZVyQ==
"""

    def setUp(self):
        self.asn1Spec = rfc5652.ContentInfo()

    def testDerCodec(self):
        substrate = pem.readBase64fromText(self.keyagree_pem_text)
        asn1Object, rest = der_decoder(substrate, asn1Spec=self.asn1Spec)
        self.assertFalse(rest)
        self.assertTrue(asn1Object.prettyPrint())
        self.assertEqual(substrate, der_encoder(asn1Object))
        self.assertEqual(rfc5652.id_envelopedData, asn1Object["contentType"])

        ed, rest = der_decoder(asn1Object["content"], asn1Spec=rfc5652.EnvelopedData())
        self.assertFalse(rest)
        self.assertTrue(ed.prettyPrint())
        self.assertEqual(asn1Object["content"], der_encoder(ed))

        ri = ed["recipientInfos"][0]
        alg1 = ri["kari"]["originator"]["originatorKey"]["algorithm"]
        self.assertEqual(rfc4357.id_GostR3410_2001, alg1["algorithm"])
        param1, rest = der_decoder(
            alg1["parameters"], asn1Spec=rfc4357.GostR3410_2001_PublicKeyParameters()
        )
        self.assertFalse(rest)
        self.assertTrue(param1.prettyPrint())
        self.assertEqual(alg1["parameters"], der_encoder(param1))

        self.assertEqual(
            rfc4357.id_GostR3410_2001_CryptoPro_XchA_ParamSet,
            param1["publicKeyParamSet"],
        )
        self.assertEqual(
            rfc4357.id_GostR3411_94_CryptoProParamSet, param1["digestParamSet"]
        )

        self.assertEqual(8, len(ri["kari"]["ukm"]))

        alg2 = ed["encryptedContentInfo"]["contentEncryptionAlgorithm"]
        self.assertEqual(
            rfc4357.id_Gost28147_89,
            alg2["algorithm"],
        )
        param2, rest = der_decoder(
            alg2["parameters"], asn1Spec=rfc4357.Gost28147_89_Parameters()
        )
        self.assertFalse(rest)
        self.assertTrue(param1.prettyPrint())
        self.assertEqual(alg2["parameters"], der_encoder(param2))

        self.assertEqual(8, len(param2["iv"]))
        self.assertEqual(
            rfc4357.id_Gost28147_89_CryptoPro_A_ParamSet, param2["encryptionParamSet"]
        )

    def testOpenTypes(self):
        openTypeMap = {
            rfc4357.id_GostR3410_2001: rfc4357.GostR3410_2001_PublicKeyParameters(),
            rfc4357.id_Gost28147_89: rfc4357.Gost28147_89_Parameters(),
        }

        substrate = pem.readBase64fromText(self.keyagree_pem_text)
        asn1Object, rest = der_decoder(
            substrate,
            asn1Spec=self.asn1Spec,
            openTypes=openTypeMap,
            decodeOpenTypes=True,
        )
        self.assertFalse(rest)
        self.assertTrue(asn1Object.prettyPrint())
        self.assertEqual(substrate, der_encoder(asn1Object))
        self.assertEqual(rfc5652.id_envelopedData, asn1Object["contentType"])

        ri = asn1Object["content"]["recipientInfos"][0]
        alg1 = ri["kari"]["originator"]["originatorKey"]["algorithm"]
        self.assertEqual(rfc4357.id_GostR3410_2001, alg1["algorithm"])
        param1 = alg1["parameters"]
        self.assertEqual(
            rfc4357.id_GostR3410_2001_CryptoPro_XchA_ParamSet,
            param1["publicKeyParamSet"],
        )
        self.assertEqual(
            rfc4357.id_GostR3411_94_CryptoProParamSet, param1["digestParamSet"]
        )

        self.assertEqual(8, len(ri["kari"]["ukm"]))

        alg2 = asn1Object["content"]["encryptedContentInfo"][
            "contentEncryptionAlgorithm"
        ]
        self.assertEqual(rfc4357.id_Gost28147_89, alg2["algorithm"])
        param2 = alg2["parameters"]
        self.assertEqual(8, len(param2["iv"]))
        self.assertEqual(
            rfc4357.id_Gost28147_89_CryptoPro_A_ParamSet, param2["encryptionParamSet"]
        )


class KeyTransportTestCase(unittest.TestCase):
    keytrans_pem_text = """\
MIIBpwYJKoZIhvcNAQcDoIIBmDCCAZQCAQAxggFTMIIBTwIBADCBgTBtMR8wHQYD
VQQDDBZHb3N0UjM0MTAtMjAwMSBleGFtcGxlMRIwEAYDVQQKDAlDcnlwdG9Qcm8x
CzAJBgNVBAYTAlJVMSkwJwYJKoZIhvcNAQkBFhpHb3N0UjM0MTAtMjAwMUBleGFt
cGxlLmNvbQIQK/XGHsIRvRfH3NRiZrQuITAcBgYqhQMCAhMwEgYHKoUDAgIkAAYH
KoUDAgIeAQSBpzCBpDAoBCBqL6ghBpVon5/kR6qey2EVK35BYLxdjfv1PSgbGJr5
dQQENm2Yt6B4BgcqhQMCAh8BoGMwHAYGKoUDAgITMBIGByqFAwICJAAGByqFAwIC
HgEDQwAEQE0rLzOQ5tyj3VUqzd/g7/sx93N+Tv+/eImKK8PNMZQESw5gSJYf28dd
Em/askCKd7W96vLsNMsjn5uL3Z4SwPYECJeV4ywrrSsMMDgGCSqGSIb3DQEHATAd
BgYqhQMCAhUwEwQIvBCLHwv/NCkGByqFAwICHwGADKqOch3uT7Mu4w+hNw==
"""

    def setUp(self):
        self.asn1Spec = rfc5652.ContentInfo()

    def testDerCodec(self):
        substrate = pem.readBase64fromText(self.keytrans_pem_text)
        asn1Object, rest = der_decoder(substrate, asn1Spec=self.asn1Spec)
        self.assertFalse(rest)
        self.assertTrue(asn1Object.prettyPrint())
        self.assertEqual(substrate, der_encoder(asn1Object))
        self.assertEqual(rfc5652.id_envelopedData, asn1Object["contentType"])

        ed, rest = der_decoder(asn1Object["content"], asn1Spec=rfc5652.EnvelopedData())
        self.assertFalse(rest)
        self.assertTrue(ed.prettyPrint())
        self.assertEqual(asn1Object["content"], der_encoder(ed))

        ri = ed["recipientInfos"][0]
        alg1 = ri["ktri"]["keyEncryptionAlgorithm"]
        self.assertEqual(rfc4357.id_GostR3410_2001, alg1["algorithm"])
        param1, rest = der_decoder(
            alg1["parameters"], asn1Spec=rfc4357.GostR3410_2001_PublicKeyParameters()
        )
        self.assertFalse(rest)
        self.assertTrue(param1.prettyPrint())
        self.assertEqual(alg1["parameters"], der_encoder(param1))
        self.assertEqual(
            rfc4357.id_GostR3410_2001_CryptoPro_XchA_ParamSet,
            param1["publicKeyParamSet"],
        )
        self.assertEqual(
            rfc4357.id_GostR3411_94_CryptoProParamSet, param1["digestParamSet"]
        )

        alg2 = ed["encryptedContentInfo"]["contentEncryptionAlgorithm"]
        self.assertEqual(rfc4357.id_Gost28147_89, alg2["algorithm"])
        param2, rest = der_decoder(
            alg2["parameters"], asn1Spec=rfc4357.Gost28147_89_Parameters()
        )
        self.assertFalse(rest)
        self.assertTrue(param2.prettyPrint())
        self.assertEqual(alg2["parameters"], der_encoder(param2))
        self.assertEqual(8, len(param2["iv"]))
        self.assertEqual(
            rfc4357.id_Gost28147_89_CryptoPro_A_ParamSet, param2["encryptionParamSet"]
        )

    def testOpenTypes(self):
        openTypeMap = {
            rfc4357.id_GostR3410_2001: rfc4357.GostR3410_2001_PublicKeyParameters(),
            rfc4357.id_Gost28147_89: rfc4357.Gost28147_89_Parameters(),
        }

        substrate = pem.readBase64fromText(self.keytrans_pem_text)
        asn1Object, rest = der_decoder(
            substrate,
            asn1Spec=self.asn1Spec,
            openTypes=openTypeMap,
            decodeOpenTypes=True,
        )
        self.assertFalse(rest)
        self.assertTrue(asn1Object.prettyPrint())
        self.assertEqual(substrate, der_encoder(asn1Object))

        ri = asn1Object["content"]["recipientInfos"][0]
        alg1 = ri["ktri"]["keyEncryptionAlgorithm"]
        self.assertEqual(rfc4357.id_GostR3410_2001, alg1["algorithm"])
        param1 = alg1["parameters"]
        self.assertEqual(
            rfc4357.id_GostR3410_2001_CryptoPro_XchA_ParamSet,
            param1["publicKeyParamSet"],
        )
        self.assertEqual(
            rfc4357.id_GostR3411_94_CryptoProParamSet, param1["digestParamSet"]
        )

        alg2 = asn1Object["content"]["encryptedContentInfo"][
            "contentEncryptionAlgorithm"
        ]
        self.assertEqual(rfc4357.id_Gost28147_89, alg2["algorithm"])
        param2 = alg2["parameters"]
        self.assertEqual(8, len(param2["iv"]))
        self.assertEqual(
            rfc4357.id_Gost28147_89_CryptoPro_A_ParamSet, param2["encryptionParamSet"]
        )


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

if __name__ == "__main__":
    result = unittest.TextTestRunner(verbosity=2).run(suite)
    sys.exit(not result.wasSuccessful())