Merge pull request #212 from akx/overrides

Overridability enhancements

Author: pitbulk

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.
@@ -103,7 +108,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):
@@ -149,7 +154,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')
@@ -165,7 +170,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")
@@ -179,7 +184,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()
@@ -395,7 +400,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()
 
@@ -449,7 +454,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/authn_request.py

@@ -47,8 +47,7 @@ def __init__(self, settings, force_authn=False, is_passive=False, set_nameid_pol
         idp_data = self.__settings.get_idp_data()
         security = self.__settings.get_security_data()
 
-        uid = OneLogin_Saml2_Utils.generate_unique_id()
-        self.__id = uid
+        self.__id = self._generate_request_id()
         issue_instant = OneLogin_Saml2_Utils.parse_time_to_SAML(OneLogin_Saml2_Utils.now())
 
         destination = idp_data['singleSignOnService']['url']
@@ -113,7 +112,7 @@ def __init__(self, settings, force_authn=False, is_passive=False, set_nameid_pol
 
         request = OneLogin_Saml2_Templates.AUTHN_REQUEST % \
             {
-                'id': uid,
+                'id': self.__id,
                 'provider_name': provider_name_str,
                 'force_authn_str': force_authn_str,
                 'is_passive_str': is_passive_str,
@@ -129,6 +128,14 @@ def __init__(self, settings, force_authn=False, is_passive=False, set_nameid_pol
 
         self.__authn_request = request
 
+    def _generate_request_id(self):
+        """
+        Generate an unique request ID.
+
+        You can override this in a subclass.
+        """
+        return OneLogin_Saml2_Utils.generate_unique_id()
+
     def get_request(self, deflate=True):
         """
         Returns unsigned AuthnRequest.

src/onelogin/saml2/idp_metadata_parser.py

@@ -25,8 +25,8 @@ class OneLogin_Saml2_IdPMetadataParser(object):
     A class that contain methods related to obtaining and parsing metadata from IdP
     """
 
-    @staticmethod
-    def get_metadata(url, validate_cert=True):
+    @classmethod
+    def get_metadata(cls, url, validate_cert=True):
         """
         Gets the metadata XML from the provided URL
         :param url: Url where the XML of the Identity Provider Metadata is published.
@@ -63,8 +63,8 @@ def get_metadata(url, validate_cert=True):
 
         return xml
 
-    @staticmethod
-    def parse_remote(url, validate_cert=True, entity_id=None, **kwargs):
+    @classmethod
+    def parse_remote(cls, url, validate_cert=True, entity_id=None, **kwargs):
         """
         Gets the metadata XML from the provided URL and parse it, returning a dict with extracted data
         :param url: Url where the XML of the Identity Provider Metadata is published.
@@ -80,11 +80,12 @@ def parse_remote(url, validate_cert=True, entity_id=None, **kwargs):
         :returns: settings dict with extracted data
         :rtype: dict
         """
-        idp_metadata = OneLogin_Saml2_IdPMetadataParser.get_metadata(url, validate_cert)
-        return OneLogin_Saml2_IdPMetadataParser.parse(idp_metadata, entity_id=entity_id, **kwargs)
+        idp_metadata = cls.get_metadata(url, validate_cert)
+        return cls.parse(idp_metadata, entity_id=entity_id, **kwargs)
 
-    @staticmethod
+    @classmethod
     def parse(
+            cls,
             idp_metadata,
             required_sso_binding=OneLogin_Saml2_Constants.BINDING_HTTP_REDIRECT,
             required_slo_binding=OneLogin_Saml2_Constants.BINDING_HTTP_REDIRECT,

src/onelogin/saml2/logout_request.py

@@ -203,8 +203,8 @@ def get_nameid_data(request, key=None):
 
         return name_id_data
 
-    @staticmethod
-    def get_nameid(request, key=None):
+    @classmethod
+    def get_nameid(cls, request, key=None):
         """
         Gets the NameID of the Logout Request Message
         :param request: Logout Request Message
@@ -214,11 +214,11 @@ def get_nameid(request, key=None):
         :return: Name ID Value
         :rtype: string
         """
-        name_id = OneLogin_Saml2_Logout_Request.get_nameid_data(request, key)
+        name_id = cls.get_nameid_data(request, key)
         return name_id['Value']
 
-    @staticmethod
-    def get_nameid_format(request, key=None):
+    @classmethod
+    def get_nameid_format(cls, request, key=None):
         """
         Gets the NameID Format of the Logout Request Message
         :param request: Logout Request Message
@@ -229,7 +229,7 @@ def get_nameid_format(request, key=None):
         :rtype: string
         """
         name_id_format = None
-        name_id_data = OneLogin_Saml2_Logout_Request.get_nameid_data(request, key)
+        name_id_data = cls.get_nameid_data(request, key)
         if name_id_data and 'Format' in name_id_data.keys():
             name_id_format = name_id_data['Format']
         return name_id_format
@@ -325,7 +325,7 @@ def is_valid(self, request_data, raise_exceptions=False):
                         )
 
                 # Check issuer
-                issuer = OneLogin_Saml2_Logout_Request.get_issuer(root)
+                issuer = self.get_issuer(root)
                 if issuer is not None and issuer != idp_entity_id:
                     raise OneLogin_Saml2_ValidationError(
                         'Invalid issuer in the Logout Request (expected %(idpEntityId)s, got %(issuer)s)' %

src/onelogin/saml2/metadata.py

@@ -34,8 +34,8 @@ class OneLogin_Saml2_Metadata(object):
     TIME_VALID = 172800   # 2 days
     TIME_CACHED = 604800  # 1 week
 
-    @staticmethod
-    def builder(sp, authnsign=False, wsign=False, valid_until=None, cache_duration=None, contacts=None, organization=None):
+    @classmethod
+    def builder(cls, sp, authnsign=False, wsign=False, valid_until=None, cache_duration=None, contacts=None, organization=None):
         """
         Builds the metadata of the SP
 
@@ -61,7 +61,7 @@ def builder(sp, authnsign=False, wsign=False, valid_until=None, cache_duration=N
         :type organization: dict
         """
         if valid_until is None:
-            valid_until = int(time()) + OneLogin_Saml2_Metadata.TIME_VALID
+            valid_until = int(time()) + cls.TIME_VALID
         if not isinstance(valid_until, basestring):
             if isinstance(valid_until, datetime):
                 valid_until_time = valid_until.timetuple()
@@ -72,7 +72,7 @@ def builder(sp, authnsign=False, wsign=False, valid_until=None, cache_duration=N
             valid_until_str = valid_until
 
         if cache_duration is None:
-            cache_duration = OneLogin_Saml2_Metadata.TIME_CACHED
+            cache_duration = cls.TIME_CACHED
         if not isinstance(cache_duration, compat.str_type):
             cache_duration_str = 'PT%sS' % cache_duration  # Period of Time x Seconds
         else:
@@ -228,8 +228,8 @@ def __add_x509_key_descriptors(root, cert, signing):
         x509_certificate.text = OneLogin_Saml2_Utils.format_cert(cert, False)
         key_descriptor.set('use', ('encryption', 'signing')[signing])
 
-    @staticmethod
-    def add_x509_key_descriptors(metadata, cert=None, add_encryption=True):
+    @classmethod
+    def add_x509_key_descriptors(cls, metadata, cert=None, add_encryption=True):
         """
         Adds the x509 descriptors (sign/encryption) to the metadata
         The same cert will be used for sign/encrypt
@@ -260,6 +260,6 @@ def add_x509_key_descriptors(metadata, cert=None, add_encryption=True):
             raise Exception('Malformed metadata.')
 
         if add_encryption:
-            OneLogin_Saml2_Metadata.__add_x509_key_descriptors(sp_sso_descriptor, cert, False)
-        OneLogin_Saml2_Metadata.__add_x509_key_descriptors(sp_sso_descriptor, cert, True)
+            cls.__add_x509_key_descriptors(sp_sso_descriptor, cert, False)
+        cls.__add_x509_key_descriptors(sp_sso_descriptor, cert, True)
         return OneLogin_Saml2_XML.to_string(root)

src/onelogin/saml2/settings.py

@@ -81,6 +81,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:
@@ -661,7 +663,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'],
@@ -672,10 +674,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:
@@ -729,7 +731,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