Make request/response/metadata classes overridable by subclassing

Author: akx

src/onelogin/saml2/auth.py

@@ -34,6 +34,11 @@ class OneLogin_Saml2_Auth(object):
     SAML Response, a Logout Request or a Logout Response).
     """
 
+    authn_request_class = OneLogin_Saml2_Authn_Request
+    logout_request_class = OneLogin_Saml2_Logout_Request
+    logout_response_class = OneLogin_Saml2_Logout_Response
+    response_class = OneLogin_Saml2_Response
+
     def __init__(self, request_data, old_settings=None, custom_base_path=None):
         """
         Initializes the SP SAML instance.
@@ -102,7 +107,7 @@ def process_response(self, request_id=None):
 
         if 'post_data' in self.__request_data and 'SAMLResponse' in self.__request_data['post_data']:
             # AuthnResponse -- HTTP_POST Binding
-            response = OneLogin_Saml2_Response(self.__settings, self.__request_data['post_data']['SAMLResponse'])
+            response = self.response_class(self.__settings, self.__request_data['post_data']['SAMLResponse'])
             self.__last_response = response.get_xml_document()
 
             if response.is_valid(self.__request_data, request_id):
@@ -147,7 +152,7 @@ def process_slo(self, keep_local_session=False, request_id=None, delete_session_
 
         get_data = 'get_data' in self.__request_data and self.__request_data['get_data']
         if get_data and 'SAMLResponse' in get_data:
-            logout_response = OneLogin_Saml2_Logout_Response(self.__settings, get_data['SAMLResponse'])
+            logout_response = self.logout_response_class(self.__settings, get_data['SAMLResponse'])
             self.__last_response = logout_response.get_xml()
             if not self.validate_response_signature(get_data):
                 self.__errors.append('invalid_logout_response_signature')
@@ -163,7 +168,7 @@ def process_slo(self, keep_local_session=False, request_id=None, delete_session_
                     OneLogin_Saml2_Utils.delete_local_session(delete_session_cb)
 
         elif get_data and 'SAMLRequest' in get_data:
-            logout_request = OneLogin_Saml2_Logout_Request(self.__settings, get_data['SAMLRequest'])
+            logout_request = self.logout_request_class(self.__settings, get_data['SAMLRequest'])
             self.__last_request = logout_request.get_xml()
             if not self.validate_request_signature(get_data):
                 self.__errors.append("invalid_logout_request_signature")
@@ -177,7 +182,7 @@ def process_slo(self, keep_local_session=False, request_id=None, delete_session_
 
                 in_response_to = logout_request.id
                 self.__last_message_id = logout_request.id
-                response_builder = OneLogin_Saml2_Logout_Response(self.__settings)
+                response_builder = self.logout_response_class(self.__settings)
                 response_builder.build(in_response_to)
                 self.__last_response = response_builder.get_xml()
                 logout_response = response_builder.get_response()
@@ -371,7 +376,7 @@ def login(self, return_to=None, force_authn=False, is_passive=False, set_nameid_
         :returns: Redirection URL
         :rtype: string
         """
-        authn_request = OneLogin_Saml2_Authn_Request(self.__settings, force_authn, is_passive, set_nameid_policy, name_id_value_req)
+        authn_request = self.authn_request_class(self.__settings, force_authn, is_passive, set_nameid_policy, name_id_value_req)
         self.__last_request = authn_request.get_xml()
         self.__last_request_id = authn_request.get_id()
 
@@ -425,7 +430,7 @@ def logout(self, return_to=None, name_id=None, session_index=None, nq=None, name
         if name_id_format is None and self.__nameid_format is not None:
             name_id_format = self.__nameid_format
 
-        logout_request = OneLogin_Saml2_Logout_Request(
+        logout_request = self.logout_request_class(
             self.__settings,
             name_id=name_id,
             session_index=session_index,

src/onelogin/saml2/settings.py

@@ -66,6 +66,8 @@ class OneLogin_Saml2_Settings(object):
 
     """
 
+    metadata_class = OneLogin_Saml2_Metadata
+
     def __init__(self, settings=None, custom_base_path=None, sp_validation_only=False):
         """
         Initializes the settings:
@@ -616,7 +618,7 @@ def get_sp_metadata(self):
         :returns: SP metadata (xml)
         :rtype: string
         """
-        metadata = OneLogin_Saml2_Metadata.builder(
+        metadata = self.metadata_class.builder(
             self.__sp, self.__security['authnRequestsSigned'],
             self.__security['wantAssertionsSigned'],
             self.__security['metadataValidUntil'],
@@ -627,10 +629,10 @@ def get_sp_metadata(self):
         add_encryption = self.__security['wantNameIdEncrypted'] or self.__security['wantAssertionsEncrypted']
 
         cert_new = self.get_sp_cert_new()
-        metadata = OneLogin_Saml2_Metadata.add_x509_key_descriptors(metadata, cert_new, add_encryption)
+        metadata = self.metadata_class.add_x509_key_descriptors(metadata, cert_new, add_encryption)
 
         cert = self.get_sp_cert()
-        metadata = OneLogin_Saml2_Metadata.add_x509_key_descriptors(metadata, cert, add_encryption)
+        metadata = self.metadata_class.add_x509_key_descriptors(metadata, cert, add_encryption)
 
         # Sign metadata
         if 'signMetadata' in self.__security and self.__security['signMetadata'] is not False:
@@ -684,7 +686,7 @@ def get_sp_metadata(self):
             signature_algorithm = self.__security['signatureAlgorithm']
             digest_algorithm = self.__security['digestAlgorithm']
 
-            metadata = OneLogin_Saml2_Metadata.sign_metadata(metadata, key_metadata, cert_metadata, signature_algorithm, digest_algorithm)
+            metadata = self.metadata_class.sign_metadata(metadata, key_metadata, cert_metadata, signature_algorithm, digest_algorithm)
 
         return metadata