From 6663bef144765990f1e5ae4ef407b836909ae29a Mon Sep 17 00:00:00 2001
From: Sixto Martin <pitbulk@gmail.com>
Date: Mon, 9 Oct 2023 13:13:06 +0200
Subject: [PATCH] Improve get_metadata method from Parser, allowing to set
 timeouts and headers

---
 README.md                                 | 12 +++++++++---
 src/onelogin/saml2/idp_metadata_parser.py | 22 +++++++++++++++++-----
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 09eee6c..d7e418d 100644
--- a/README.md
+++ b/README.md
@@ -589,7 +589,7 @@ The method above requires a little extra work to manually specify attributes abo
 
 There's an easier method -- use a metadata exchange.  Metadata is just an XML file that defines the capabilities of both the IdP and the SP application.  It also contains the X.509 public key certificates which add to the trusted relationship.  The IdP administrator can also configure custom settings for an SP based on the metadata.
 
-Using ````parse_remote```` IdP metadata can be obtained and added to the settings withouth further ado.
+Using ````parse_remote```` IdP metadata can be obtained and added to the settings without further ado.
 
 But take in mind that the OneLogin_Saml2_IdPMetadataParser class does not validate in any way the URL that is introduced in order to be parsed.
 
@@ -598,9 +598,15 @@ Usually the same administrator that handles the Service Provider also sets the U
 But there are other scenarios, like a SAAS app where the administrator of the app delegates this functionality to other users. In this case, extra precaution should be taken in order to validate such URL inputs and avoid attacks like SSRF.
 
 
-``
+```
 idp_data = OneLogin_Saml2_IdPMetadataParser.parse_remote('https://example.com/auth/saml2/idp/metadata')
-``
+```
+
+You can specify a timeout in seconds for metadata retrieval, if not it is not guaranteed that the request will complete
+
+```
+idp_data = OneLogin_Saml2_IdPMetadataParser.parse_remote('https://example.com/auth/saml2/idp/metadata', timeout=5)
+```
 
 If the Metadata contains several entities, the relevant ``EntityDescriptor`` can be specified when retrieving the settings from the ``IdpMetadataParser`` by its ``EntityId`` value:
 ```
diff --git a/src/onelogin/saml2/idp_metadata_parser.py b/src/onelogin/saml2/idp_metadata_parser.py
index 7d03bf5..8981211 100644
--- a/src/onelogin/saml2/idp_metadata_parser.py
+++ b/src/onelogin/saml2/idp_metadata_parser.py
@@ -27,7 +27,7 @@ class OneLogin_Saml2_IdPMetadataParser(object):
     """
 
     @staticmethod
-    def get_metadata(url, validate_cert=True):
+    def get_metadata(url, validate_cert=True, timeout=None, headers=None):
         """
         Gets the metadata XML from the provided URL
 
@@ -37,17 +37,26 @@ def get_metadata(url, validate_cert=True):
         :param validate_cert: If the url uses https schema, that flag enables or not the verification of the associated certificate.
         :type validate_cert: bool
 
+        :param timeout: Timeout in seconds to wait for metadata response
+        :type timeout: int
+
+        :param headers: Extra headers to send in the request
+        :type headers: dict
+
         :returns: metadata XML
         :rtype: string
         """
         valid = False
+
+        request = urllib2.Request(url, headers=headers or {})
+
         if validate_cert:
-            response = urllib2.urlopen(url)
+            response = urllib2.urlopen(request, timeout=timeout)
         else:
             ctx = ssl.create_default_context()
             ctx.check_hostname = False
             ctx.verify_mode = ssl.CERT_NONE
-            response = urllib2.urlopen(url, context=ctx)
+            response = urllib2.urlopen(request, context=ctx, timeout=timeout)
         xml = response.read()
 
         if xml:
@@ -65,7 +74,7 @@ def get_metadata(url, validate_cert=True):
         return xml
 
     @staticmethod
-    def parse_remote(url, validate_cert=True, entity_id=None, **kwargs):
+    def parse_remote(url, validate_cert=True, entity_id=None, timeout=None, **kwargs):
         """
         Gets the metadata XML from the provided URL and parse it, returning a dict with extracted data
 
@@ -79,10 +88,13 @@ def parse_remote(url, validate_cert=True, entity_id=None, **kwargs):
                           that contains multiple EntityDescriptor.
         :type entity_id: string
 
+        :param timeout: Timeout in seconds to wait for metadata response
+        :type timeout: int
+
         :returns: settings dict with extracted data
         :rtype: dict
         """
-        idp_metadata = OneLogin_Saml2_IdPMetadataParser.get_metadata(url, validate_cert)
+        idp_metadata = OneLogin_Saml2_IdPMetadataParser.get_metadata(url, validate_cert, timeout, headers=kwargs.pop('headers', None))
         return OneLogin_Saml2_IdPMetadataParser.parse(idp_metadata, entity_id=entity_id, **kwargs)
 
     @staticmethod