Better generalization for PEM certsshibsp_enc

2 files changed, 17 insertions, 15 deletions

diff --git a/src/saml2/entity.py b/src/saml2/entity.py
index 3b1b5829..12d882f2 100644
--- a/src/saml2/entity.py
+++ b/src/saml2/entity.py
@@ -1,7 +1,6 @@
 import base64
 import copy
 import logging
-import re
 import requests
 import six
 
@@ -66,6 +65,7 @@ from saml2.sigver import security_context
 from saml2.sigver import SigverError
 from saml2.sigver import SignatureError
 from saml2.sigver import make_temp
+from saml2.sigver import get_pem_wrapped_unwrapped
 from saml2.sigver import pre_encryption_part
 from saml2.sigver import pre_signature_part
 from saml2.sigver import pre_encrypt_assertion
@@ -651,10 +651,7 @@ 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}'
+            wrapped_cert, unwrapped_cert = get_pem_wrapped_unwrapped(_cert)
             try:
                 tmp = make_temp(wrapped_cert.encode('ascii'),
                                 decode=False,
@@ -665,6 +662,7 @@ class Entity(HTTPBase):
                 if encrypt_cert:
                     pre_enc_part_dict['encrypt_cert'] = unwrapped_cert
                 pre_enc_part = pre_encryption_part(**pre_enc_part_dict)
+                # end pre_enc_part
                 
                 response = self.sec.encrypt_assertion(response, tmp.name,
                                                       pre_enc_part,
diff --git a/src/saml2/sigver.py b/src/saml2/sigver.py
index 5563627b..9744bd2c 100644
--- a/src/saml2/sigver.py
+++ b/src/saml2/sigver.py
@@ -8,6 +8,7 @@ import hashlib
 import itertools
 import logging
 import os
+import re
 import six
 from uuid import uuid4 as gen_random_key
 
@@ -61,11 +62,8 @@ 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'
-
-TRIPLE_DES_CBC = 'http://www.w3.org/2001/04/xmlenc#tripledes-cbc'
-RSA_OAEP = "http://www.w3.org/2001/04/xmlenc#rsa-oaep-mgf1p"
+TRIPLEDES_CBC = 'http://www.w3.org/2001/04/xmlenc#tripledes-cbc'
+RSA_OAEP_MGF1P = "http://www.w3.org/2001/04/xmlenc#rsa-oaep-mgf1p"
 
 class SigverError(SAMLError):
     pass
@@ -104,6 +102,14 @@ class CertificateError(SigverError):
     pass
 
 
+def get_pem_wrapped_unwrapped(cert):
+    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}'
+    return wrapped_cert, unwrapped_cert
+
+
 def read_file(*args, **kwargs):
     with open(*args, **kwargs) as handler:
         return handler.read()
@@ -1088,10 +1094,8 @@ def encrypt_cert_from_item(item):
         pass
 
     if _encrypt_cert is not None:
-        if _encrypt_cert.find('-----BEGIN CERTIFICATE-----\n') == -1:
-            _encrypt_cert = '-----BEGIN CERTIFICATE-----\n' + _encrypt_cert
-        if _encrypt_cert.find('\n-----END CERTIFICATE-----') == -1:
-            _encrypt_cert = _encrypt_cert + '\n-----END CERTIFICATE-----'
+        wrapped_cert, unwrapped_cert = get_pem_wrapped_unwrapped(_encrypt_cert)
+        _encrypt_cert = wrapped_cert
     return _encrypt_cert
 
 
@@ -1851,7 +1855,7 @@ def pre_signature_part(
 # </EncryptedData>
 
 
-def pre_encryption_part(msg_enc=TRIPLE_DES_CBC, key_enc=RSA_OAEP, 
+def pre_encryption_part(msg_enc=TRIPLEDES_CBC, key_enc=RSA_OAEP_MGF1P, 
         key_name='my-rsa-key',
         encrypted_key_id=None, encrypted_data_id=None,
         encrypt_cert=None):