import unittest

from lxml import etree

import xmlsec
from tests import base

consts = xmlsec.constants


class TestTemplates(base.TestMemoryLeaks):
    def test_create(self):
        root = self.load_xml('doc.xml')
        sign = xmlsec.template.create(
            root, c14n_method=consts.TransformExclC14N, sign_method=consts.TransformRsaSha1, id='Id', ns='test'
        )
        self.assertEqual('Id', sign.get('Id'))
        self.assertEqual('test', sign.prefix)

    def test_create_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.create('', c14n_method=consts.TransformExclC14N, sign_method=consts.TransformRsaSha1)

    def test_encrypt_data_create(self):
        root = self.load_xml('doc.xml')
        enc = xmlsec.template.encrypted_data_create(
            root, method=consts.TransformDes3Cbc, id='Id', type='Type', mime_type='MimeType', encoding='Encoding', ns='test'
        )
        for a in ('Id', 'Type', 'MimeType', 'Encoding'):
            self.assertEqual(a, enc.get(a))
        self.assertEqual('test', enc.prefix)

    def test_ensure_key_info(self):
        root = self.load_xml('doc.xml')
        sign = xmlsec.template.create(root, c14n_method=consts.TransformExclC14N, sign_method=consts.TransformRsaSha1)
        ki = xmlsec.template.ensure_key_info(sign, id='Id')
        self.assertEqual('Id', ki.get('Id'))

    def test_ensure_key_info_fail(self):
        with self.assertRaisesRegex(xmlsec.Error, 'cannot ensure key info.'):
            xmlsec.template.ensure_key_info(etree.fromstring(b'<Data/>'), id='Id')

    def test_ensure_key_info_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.ensure_key_info('', id=0)

    def test_add_encrypted_key(self):
        root = self.load_xml('doc.xml')
        sign = xmlsec.template.create(root, c14n_method=consts.TransformExclC14N, sign_method=consts.TransformRsaSha1)
        ki = xmlsec.template.ensure_key_info(sign)
        ek = xmlsec.template.add_encrypted_key(ki, consts.TransformRsaOaep)
        self.assertEqual(ek, xmlsec.tree.find_node(self.load_xml('sign_template.xml'), consts.NodeEncryptedKey, consts.EncNs))
        ek2 = xmlsec.template.add_encrypted_key(ki, consts.TransformRsaOaep, id='Id', type='Type', recipient='Recipient')
        for a in ('Id', 'Type', 'Recipient'):
            self.assertEqual(a, ek2.get(a))

    def test_add_key_name(self):
        root = self.load_xml('doc.xml')
        sign = xmlsec.template.create(root, c14n_method=consts.TransformExclC14N, sign_method=consts.TransformRsaSha1)
        ki = xmlsec.template.ensure_key_info(sign)
        kn = xmlsec.template.add_key_name(ki)
        self.assertEqual(kn, xmlsec.tree.find_node(self.load_xml('sign_template.xml'), consts.NodeKeyName, consts.DSigNs))
        kn2 = xmlsec.template.add_key_name(ki, name='name')
        self.assertEqual('name', kn2.text)

    def test_add_key_name_none(self):
        root = self.load_xml('doc.xml')
        sign = xmlsec.template.create(root, c14n_method=consts.TransformExclC14N, sign_method=consts.TransformRsaSha1)
        ki = xmlsec.template.ensure_key_info(sign)
        kn2 = xmlsec.template.add_key_name(ki, name=None)
        self.assertEqual(kn2.text, None)
        print(etree.tostring(kn2))

    def test_add_key_name_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.add_key_name('')

    def test_add_reference(self):
        root = self.load_xml('doc.xml')
        sign = xmlsec.template.create(root, c14n_method=consts.TransformExclC14N, sign_method=consts.TransformRsaSha1)
        ref = xmlsec.template.add_reference(sign, consts.TransformSha1, id='Id', uri='URI', type='Type')
        for a in ('Id', 'URI', 'Type'):
            self.assertEqual(a, ref.get(a))

    def test_add_reference_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.add_reference('', consts.TransformSha1)
        with self.assertRaises(TypeError):
            xmlsec.template.add_reference(etree.Element('root'), '')

    def test_add_reference_fail(self):
        with self.assertRaisesRegex(xmlsec.Error, 'cannot add reference.'):
            xmlsec.template.add_reference(etree.Element('root'), consts.TransformSha1)

    def test_add_transform_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.add_transform('', consts.TransformSha1)
        with self.assertRaises(TypeError):
            xmlsec.template.add_transform(etree.Element('root'), '')

    def test_add_key_value(self):
        root = self.load_xml('doc.xml')
        sign = xmlsec.template.create(root, c14n_method=consts.TransformExclC14N, sign_method=consts.TransformRsaSha1)
        ki = xmlsec.template.ensure_key_info(sign)
        kv = xmlsec.template.add_key_value(ki)
        self.assertEqual(kv, xmlsec.tree.find_node(self.load_xml('sign_template.xml'), consts.NodeKeyValue, consts.DSigNs))

    def test_add_key_value_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.add_key_value('')

    def test_add_x509_data(self):
        root = self.load_xml('doc.xml')
        sign = xmlsec.template.create(root, c14n_method=consts.TransformExclC14N, sign_method=consts.TransformRsaSha1)
        ki = xmlsec.template.ensure_key_info(sign)
        x509 = xmlsec.template.add_x509_data(ki)
        xmlsec.template.x509_data_add_certificate(x509)
        xmlsec.template.x509_data_add_crl(x509)
        issuer = xmlsec.template.x509_data_add_issuer_serial(x509)
        xmlsec.template.x509_data_add_ski(x509)
        xmlsec.template.x509_data_add_subject_name(x509)
        xmlsec.template.x509_issuer_serial_add_issuer_name(issuer)
        xmlsec.template.x509_issuer_serial_add_serial_number(issuer)
        self.assertEqual(x509, xmlsec.tree.find_node(self.load_xml('sign_template.xml'), consts.NodeX509Data, consts.DSigNs))

    def test_add_x509_data_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.add_x509_data('')

    def test_x509_issuer_serial_add_issuer(self):
        root = self.load_xml('doc.xml')
        sign = xmlsec.template.create(root, c14n_method=consts.TransformExclC14N, sign_method=consts.TransformRsaSha1)
        ki = xmlsec.template.ensure_key_info(sign)
        x509 = xmlsec.template.add_x509_data(ki)
        issuer = xmlsec.template.x509_data_add_issuer_serial(x509)
        name = xmlsec.template.x509_issuer_serial_add_issuer_name(issuer, name='Name')
        serial = xmlsec.template.x509_issuer_serial_add_serial_number(issuer, serial='Serial')
        self.assertEqual('Name', name.text)
        self.assertEqual('Serial', serial.text)

    def test_x509_issuer_serial_add_issuer_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.x509_data_add_issuer_serial('')

    def test_x509_issuer_serial_add_issuer_name_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.x509_issuer_serial_add_issuer_name('')

    def test_x509_issuer_serial_add_serial_number_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.x509_issuer_serial_add_serial_number('')

    def test_x509_data_add_subject_name_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.x509_data_add_subject_name('')

    def test_x509_data_add_ski_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.x509_data_add_ski('')

    def test_x509_data_add_certificate_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.x509_data_add_certificate('')

    def test_x509_data_add_crl_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.x509_data_add_crl('')

    def test_add_encrypted_key_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.add_encrypted_key('', 0)

    def test_encrypted_data_create_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.encrypted_data_create('', 0)

    def test_encrypted_data_ensure_cipher_value(self):
        root = self.load_xml('doc.xml')
        enc = xmlsec.template.encrypted_data_create(root, method=consts.TransformDes3Cbc)
        cv = xmlsec.template.encrypted_data_ensure_cipher_value(enc)
        self.assertEqual(cv, xmlsec.tree.find_node(self.load_xml('sign_template.xml'), consts.NodeCipherValue, consts.EncNs))

    def test_encrypted_data_ensure_cipher_value_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.encrypted_data_ensure_cipher_value('')

    def test_encrypted_data_ensure_key_info(self):
        root = self.load_xml('doc.xml')
        enc = xmlsec.template.encrypted_data_create(root, method=consts.TransformDes3Cbc)
        ki = xmlsec.template.encrypted_data_ensure_key_info(enc)
        self.assertEqual(ki, xmlsec.tree.find_node(self.load_xml('enc_template.xml'), consts.NodeKeyInfo, consts.DSigNs))
        ki2 = xmlsec.template.encrypted_data_ensure_key_info(enc, id='Id', ns='test')
        self.assertEqual('Id', ki2.get('Id'))
        self.assertEqual('test', ki2.prefix)

    def test_encrypted_data_ensure_key_info_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.encrypted_data_ensure_key_info('')

    @unittest.skipIf(not hasattr(consts, 'TransformXslt'), reason='XSLT transformations not enabled')
    def test_transform_add_c14n_inclusive_namespaces(self):
        root = self.load_xml('doc.xml')
        sign = xmlsec.template.create(root, c14n_method=consts.TransformExclC14N, sign_method=consts.TransformRsaSha1)
        ref = xmlsec.template.add_reference(sign, consts.TransformSha1)
        trans1 = xmlsec.template.add_transform(ref, consts.TransformEnveloped)
        xmlsec.template.transform_add_c14n_inclusive_namespaces(trans1, 'default')
        trans2 = xmlsec.template.add_transform(ref, consts.TransformXslt)
        xmlsec.template.transform_add_c14n_inclusive_namespaces(trans2, ['ns1', 'ns2'])
        self.assertEqual(ref, xmlsec.tree.find_node(self.load_xml('sign_template.xml'), consts.NodeReference, consts.DSigNs))

    def test_transform_add_c14n_inclusive_namespaces_bad_args(self):
        with self.assertRaises(TypeError):
            xmlsec.template.transform_add_c14n_inclusive_namespaces('', [])
