Tests Fixed

a regexp generalization for wrapped and unwrapped PEM certs (BEGIN/END)

2 files changed, 21 insertions, 15 deletions

diff --git a/src/saml2/entity.py b/src/saml2/entity.py
index 3311e6a2..3b1b5829 100644
--- a/src/saml2/entity.py
+++ b/src/saml2/entity.py
@@ -1,6 +1,7 @@
 import base64
 import copy
 import logging
+import re
 import requests
 import six
 
@@ -650,19 +651,20 @@ class Entity(HTTPBase):
             _certs = self.metadata.certs(sp_entity_id, "any", "encryption")
         exception = None
         for _cert in _certs:
+            begin_cert = "-----BEGIN CERTIFICATE-----\n"
+            end_cert = "\n-----END CERTIFICATE-----\n"
+            unwrapped_cert = re.sub(f'{begin_cert}|{end_cert}', '', _cert)
+            wrapped_cert = f'{begin_cert}{unwrapped_cert}{end_cert}'
             try:
-                begin_cert = "-----BEGIN CERTIFICATE-----\n"
-                end_cert = "\n-----END CERTIFICATE-----\n"
-                if begin_cert not in _cert:
-                    _cert = "%s%s" % (begin_cert, _cert)
-                if end_cert not in _cert:
-                    _cert = "%s%s" % (_cert, end_cert)
-                tmp = make_temp(_cert.encode('ascii'),
+                tmp = make_temp(wrapped_cert.encode('ascii'),
                                 decode=False,
                                 delete_tmpfiles=self.config.delete_tmpfiles)
                 
                 # it would be possibile to handle many other args here ...
-                pre_enc_part = pre_encryption_part(encrypt_cert=encrypt_cert)
+                pre_enc_part_dict = dict()
+                if encrypt_cert:
+                    pre_enc_part_dict['encrypt_cert'] = unwrapped_cert
+                pre_enc_part = pre_encryption_part(**pre_enc_part_dict)
                 
                 response = self.sec.encrypt_assertion(response, tmp.name,
                                                       pre_enc_part,
diff --git a/src/saml2/sigver.py b/src/saml2/sigver.py
index e65cb2c8..5563627b 100644
--- a/src/saml2/sigver.py
+++ b/src/saml2/sigver.py
@@ -62,7 +62,7 @@ logger = logging.getLogger(__name__)
 SIG = '{{{ns}#}}{attribute}'.format(ns=ds.NAMESPACE, attribute='Signature')
 
 # deprecated 
-RSA_1_5 = 'http://www.w3.org/2001/04/xmlenc#rsa-1_5'
+# RSA_1_5 = 'http://www.w3.org/2001/04/xmlenc#rsa-1_5'
 
 TRIPLE_DES_CBC = 'http://www.w3.org/2001/04/xmlenc#tripledes-cbc'
 RSA_OAEP = "http://www.w3.org/2001/04/xmlenc#rsa-oaep-mgf1p"
@@ -754,7 +754,7 @@ class CryptoBackendXmlSec1(CryptoBackend):
         :param key_type: The type of session key to use.
         :return: The encrypted text
         """
-
+        
         if isinstance(statement, SamlBase):
             statement = pre_encrypt_assertion(statement)
 
@@ -1293,7 +1293,7 @@ class SecurityContext(object):
 
         self.metadata = metadata
         self.only_use_keys_in_metadata = only_use_keys_in_metadata
-
+        
         if not template:
             this_dir, this_filename = os.path.split(__file__)
             self.template = os.path.join(this_dir, 'xml_template', 'template.xml')
@@ -1866,13 +1866,17 @@ def pre_encryption_part(msg_enc=TRIPLE_DES_CBC, key_enc=RSA_OAEP,
     ed_id = encrypted_data_id or "ED_{id}".format(id=gen_random_key())
     msg_encryption_method = EncryptionMethod(algorithm=msg_enc)
     key_encryption_method = EncryptionMethod(algorithm=key_enc)
+    
+    enc_key_dict= dict(key_name=ds.KeyName(text=key_name))
+    
+    enc_key_dict['x509_data'] = ds.X509Data(
+        x509_certificate=ds.X509Certificate(text=encrypt_cert))
+    key_info = ds.KeyInfo(**enc_key_dict)
+    
     encrypted_key = EncryptedKey(
         id=ek_id,
         encryption_method=key_encryption_method,
-        key_info=ds.KeyInfo(key_name=ds.KeyName(text=key_name),
-                            x509_data=ds.X509Data(
-                            x509_certificate=ds.X509Certificate(text=encrypt_cert)
-                        )),
+        key_info=key_info,
         cipher_data=CipherData(cipher_value=CipherValue(text='')),
     )
     key_info = ds.KeyInfo(encrypted_key=encrypted_key)