Better generalization for PEM certs

IdentityPython · Dec 14, 2020 · daf2142 · daf2142
1 parent 867c8db
commit daf2142
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 15 deletions.
diff --git 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 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 @@ def _encrypt_assertion(self, encrypt_cert, sp_entity_id, response, node_xpath=No
             _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 @@ def _encrypt_assertion(self, encrypt_cert, sp_entity_id, response, node_xpath=No
                 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
@@ -8,6 +8,7 @@
 import itertools
 import logging
 import os
+import re
 import six
 from uuid import uuid4 as gen_random_key
 
@@ -61,11 +62,8 @@
 
 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):