Improve decrypt method, Add an option to decrypt an element in place or copy it before decryption.

pitbulk · pitbulk · commit c193ee3458d4 · 2017-09-15T14:02:42.000+02:00
diff --git a/src/onelogin/saml2/response.py b/src/onelogin/saml2/response.py
@@ -778,7 +778,7 @@ def __decrypt_assertion(self, xml):
                             keyinfo.append(encrypted_key[0])
 
                 encrypted_data = encrypted_data_nodes[0]
-                decrypted = OneLogin_Saml2_Utils.decrypt_element(encrypted_data, key, debug)
+                decrypted = OneLogin_Saml2_Utils.decrypt_element(encrypted_data, key, debug=debug, inplace=True)
                 xml.replace(encrypted_assertion_nodes[0], decrypted)
         return xml
 
diff --git a/src/onelogin/saml2/utils.py b/src/onelogin/saml2/utils.py
@@ -10,14 +10,17 @@
 """
 
 import base64
-from datetime import datetime
+from copy import deepcopy
 import calendar
+from datetime import datetime
+from defusedxml.lxml import fromstring
 from hashlib import sha1, sha256, sha384, sha512
 from isodate import parse_duration as duration_parser
 import re
 from textwrap import wrap
 from functools import wraps
 from uuid import uuid4
+from xml.dom.minidom import Element
 
 import zlib
 import xmlsec
@@ -655,7 +658,7 @@ def get_status(dom):
         return status
 
     @staticmethod
-    def decrypt_element(encrypted_data, key, debug=False):
+    def decrypt_element(encrypted_data, key, debug=False, inplace=False):
         """
         Decrypts an encrypted element.
 
@@ -668,10 +671,20 @@ def decrypt_element(encrypted_data, key, debug=False):
         :param debug: Activate the xmlsec debug
         :type: bool
 
+        :param inplace: update passed data with decrypted result
+        :type: bool
+
         :returns: The decrypted element.
         :rtype: lxml.etree.Element
         """
-        encrypted_data = OneLogin_Saml2_XML.to_etree(encrypted_data)
+
+        if isinstance(encrypted_data, Element):
+            encrypted_data = fromstring(str(encrypted_data.toxml()))
+        if not inplace and isinstance(encrypted_data, OneLogin_Saml2_XML._element_class):
+            encrypted_data = deepcopy(encrypted_data)
+        elif isinstance(encrypted_data, OneLogin_Saml2_XML._text_class):
+            encrypted_data = OneLogin_Saml2_XML._parse_etree(encrypted_data)
+
         xmlsec.enable_debug_trace(debug)
         manager = xmlsec.KeysManager()
 
diff --git a/tests/data/misc/sp4.key b/tests/data/misc/sp4.key
@@ -0,0 +1,28 @@
+-----BEGIN PRIVATE KEY-----
+MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQD4ZrcXcjCBOQS7
+stUabuXPYnXKvcoJUrMVPRX1zfrXvpfghCrykbL1TKoqGfmEA9oNRoMBOmZCgLlK
+eb0TfuEO/u1jf4rRFcK7U/dYEiX74bQgUnJUWTfFlhwPjxGhn9zDrc2tSpworJBV
+amyBZIo5Beap5OJLote/Wqp1DZjNyEZ2m8m+lv8udmejmlo5RMoIzuG3VdH6ADC9
+LKF+QsXC/HRZBhLE/y+75/XrNODvX8eM8+9Xp21QlVF1EIZDfNQ2iHsA8GEpJDC5
+aomTW/xExBysejnwP2ROrfm3PIfP64EbB4G01f8eErlXeUD0oQ0gECgIXsJpfBkD
+IWMHwx3/AgMBAAECggEAdbLNvFlJ7GDlAj75RJ4ZXAuOPrNw4LwDyON53U9tNP7F
+HgfiBa/NuPdLhclq9geRMUsg1dsjCw3NPiGy2mL7JszaFJQhZXLHI1Xk1CE9SD0o
+yUvniln/2CqJP0IOG6QQydM3qo24snkZpq9XnHPUHrLSGdwu8aHGUpAWRoJbzdzR
+tBWBn6SlkuaE52vcGh7eMdKSICRCg2/gg6LIi89pkiI9tfozAL2LPcDTRGp3DA3w
+U6OO8k+d1La4s9G0i22OGSwPxGerTHnBIzpeM/ivRwBypFy3EV9bbjQlheI53xAo
+ZMmGeSnQ89MWgY64pnWrX862Mf1EZYTjumDe2dl1kQKBgQD9pBG2BbcQ8qieTf84
+92LeOYTPRdd0N+gdyDKKorRO772zgxBwpSwO285nzy/FKSnpJIDtuee6OFClnDor
+Ui6lG1WPQeoSEdH1V10XkfSaoFOz7Hyv9H2dCLvW/VO9KYq07VAmQcvNZnqIW+tI
+edSHcQ3I8tnw4CiFa0BPvdhk9wKBgQD6tiuN2NvuNFFLvwpBGp3hjGyn6siyXDyP
+8IXQmP66NxKqcX/NafVO3bVh6VrPGd7PL1PloQZ5EBG2PPtRdf/g4aeZKZleCUXm
+9OgMEOUqdbTP9TGrmgNPtNBx3jnhnX/GTy/7GK77YlXEVplezWaerwRM7NCFCtp2
+W6K1M961OQKBgQDDSznr2hirrvuP8GRMW4a/rrAI3DDZplZN4CCySDbm9IcvGgJl
+iXgT9MDHg2q3t0sy3U18PYEkDEpkSZcsVfneXN6TEGCHCzuLWXovNM2O5VWtmrAi
+1vCFIf1nuuRoKP1I89SbsFuYyogcSBIwWsX+h1ji2cJfSmlI2VzKSVW93wKBgQDA
+sqwfRoMkP0oM8jUrfQ3Egm4xUiAYFxTlfXUcs7t13UaXgs08USifCYGUVAvcCoJa
+tIHDiVS0UEmMzKpOHmghrM9oxbR/tpjnv21reMDrNbVX8ZnPz3ykEtHz816BrtC6
+17qFQJ+d0CMj2XvghfdOGC8yAQL0fzcSqbQRmmCe4QKBgFWY9fqHEKdG/UlxZfBB
+C/QRNTJsrbZf9Ok/o1h6BHnK64xUc4elShEwV9IdC4QNW0UCr7WXoGLUkhfUphId
+q//KUDNc7VrWj5URsZcGi7WMkqNm9kPkpeuh3iSvh3+q7tK0/yfuj9ZQOjKzQnit
+VZBooJAJGdSqYgitpyxB71/n
+-----END PRIVATE KEY-----
diff --git a/tests/src/OneLogin/saml2_tests/utils_test.py b/tests/src/OneLogin/saml2_tests/utils_test.py
@@ -8,6 +8,7 @@
 from lxml import etree
 from os.path import dirname, join, exists
 import unittest
+from defusedxml.lxml import fromstring
 from xml.dom.minidom import parseString
 
 from onelogin.saml2 import compat
@@ -613,7 +614,7 @@ def testDecryptElement(self):
         encrypted_nameid_nodes = dom_nameid_enc.find('.//saml:EncryptedID', namespaces=OneLogin_Saml2_Constants.NSMAP)
         encrypted_data = encrypted_nameid_nodes[0]
         decrypted_nameid = OneLogin_Saml2_Utils.decrypt_element(encrypted_data, key)
-        self.assertEqual('{%s}NameID' % OneLogin_Saml2_Constants.NS_SAML, decrypted_nameid.tag)
+        self.assertEqual('saml:NameID', decrypted_nameid.tag)
         self.assertEqual('2de11defd199f8d5bb63f9b7deb265ba5c675c10', decrypted_nameid.text)
 
         xml_assertion_enc = b64decode(self.file_contents(join(self.data_path, 'responses', 'valid_encrypted_assertion_encrypted_nameid.xml.base64')))
@@ -636,24 +637,71 @@ def testDecryptElement(self):
         key2 = f.read()
         f.close()
 
-        self.assertRaises(Exception, OneLogin_Saml2_Utils.decrypt_element, encrypted_data, key2)
+        # sp.key and sp2.key are equivalent we should be able to decrypt the nameID again
+        decrypted_nameid = OneLogin_Saml2_Utils.decrypt_element(encrypted_data, key2)
+        self.assertIn('{%s}NameID' % (OneLogin_Saml2_Constants.NS_SAML), decrypted_nameid.tag)
+        self.assertEqual('457bdb600de717891c77647b0806ce59c089d5b8', decrypted_nameid.text)
 
-        key_3_file_name = join(self.data_path, 'misc', 'sp2.key')
+        key_3_file_name = join(self.data_path, 'misc', 'sp3.key')
         f = open(key_3_file_name, 'r')
         key3 = f.read()
         f.close()
-        self.assertRaises(Exception, OneLogin_Saml2_Utils.decrypt_element, encrypted_data, key3)
+
+        # sp.key and sp3.key are equivalent we should be able to decrypt the nameID again
+        decrypted_nameid = OneLogin_Saml2_Utils.decrypt_element(encrypted_data, key3)
+        self.assertIn('{%s}NameID' % (OneLogin_Saml2_Constants.NS_SAML), decrypted_nameid.tag)
+        self.assertEqual('457bdb600de717891c77647b0806ce59c089d5b8', decrypted_nameid.text)
+
+        key_4_file_name = join(self.data_path, 'misc', 'sp4.key')
+        f = open(key_4_file_name, 'r')
+        key4 = f.read()
+        f.close()
+
+        with self.assertRaisesRegex(Exception, "(1, 'failed to decrypt')"):
+            OneLogin_Saml2_Utils.decrypt_element(encrypted_data, key4)
+
         xml_nameid_enc_2 = b64decode(self.file_contents(join(self.data_path, 'responses', 'invalids', 'encrypted_nameID_without_EncMethod.xml.base64')))
-        dom_nameid_enc_2 = etree.fromstring(xml_nameid_enc_2)
-        encrypted_nameid_nodes_2 = dom_nameid_enc_2.find('.//saml:EncryptedID', namespaces=OneLogin_Saml2_Constants.NSMAP)
-        encrypted_data_2 = encrypted_nameid_nodes_2[0]
-        self.assertRaises(Exception, OneLogin_Saml2_Utils.decrypt_element, encrypted_data_2, key)
+        dom_nameid_enc_2 = parseString(xml_nameid_enc_2)
+        encrypted_nameid_nodes_2 = dom_nameid_enc_2.getElementsByTagName('saml:EncryptedID')
+        encrypted_data_2 = encrypted_nameid_nodes_2[0].firstChild
+
+        with self.assertRaisesRegex(Exception, "(1, 'failed to decrypt')"):
+            OneLogin_Saml2_Utils.decrypt_element(encrypted_data_2, key)
 
         xml_nameid_enc_3 = b64decode(self.file_contents(join(self.data_path, 'responses', 'invalids', 'encrypted_nameID_without_keyinfo.xml.base64')))
-        dom_nameid_enc_3 = etree.fromstring(xml_nameid_enc_3)
-        encrypted_nameid_nodes_3 = dom_nameid_enc_3.find('.//saml:EncryptedID', namespaces=OneLogin_Saml2_Constants.NSMAP)
-        encrypted_data_3 = encrypted_nameid_nodes_3[0]
-        self.assertRaises(Exception, OneLogin_Saml2_Utils.decrypt_element, encrypted_data_3, key)
+        dom_nameid_enc_3 = parseString(xml_nameid_enc_3)
+        encrypted_nameid_nodes_3 = dom_nameid_enc_3.getElementsByTagName('saml:EncryptedID')
+        encrypted_data_3 = encrypted_nameid_nodes_3[0].firstChild
+
+        with self.assertRaisesRegex(Exception, "(1, 'failed to decrypt')"):
+            OneLogin_Saml2_Utils.decrypt_element(encrypted_data_3, key)
+
+    def testDecryptElementInplace(self):
+        """
+        Tests the decrypt_element method of the OneLogin_Saml2_Utils with inplace=True
+        """
+        settings = OneLogin_Saml2_Settings(self.loadSettingsJSON())
+
+        key = settings.get_sp_key()
+
+        xml_nameid_enc = b64decode(self.file_contents(join(self.data_path, 'responses', 'response_encrypted_nameid.xml.base64')))
+        dom = fromstring(xml_nameid_enc)
+        encrypted_node = dom.xpath('//saml:EncryptedID/xenc:EncryptedData', namespaces=OneLogin_Saml2_Constants.NSMAP)[0]
+
+        # can be decrypted twice when copy the node first
+        for _ in range(2):
+            decrypted_nameid = OneLogin_Saml2_Utils.decrypt_element(encrypted_node, key, inplace=False)
+            self.assertIn('NameID', decrypted_nameid.tag)
+            self.assertEqual('2de11defd199f8d5bb63f9b7deb265ba5c675c10', decrypted_nameid.text)
+
+        # can only be decrypted once in place
+        decrypted_nameid = OneLogin_Saml2_Utils.decrypt_element(encrypted_node, key, inplace=True)
+        self.assertIn('NameID', decrypted_nameid.tag)
+        self.assertEqual('2de11defd199f8d5bb63f9b7deb265ba5c675c10', decrypted_nameid.text)
+
+        # can't be decrypted twice since it has been decrypted inplace
+        with self.assertRaisesRegex(Exception, "(1, 'failed to decrypt')"):
+            OneLogin_Saml2_Utils.decrypt_element(encrypted_node, key, inplace=True)
 
     def testAddSign(self):
         """
