Use `extension_elements` to extract `KeyInfo`

Instead of explicitly declaring `KeyInfo` as child of `SubjectConfirmationData`, use `extension_elements` to extract `KeyInfo` element(s).

4 files changed, 10 insertions, 12 deletions

diff --git a/src/saml2/response.py b/src/saml2/response.py
index 118f7fe0..4c884c30 100644
--- a/src/saml2/response.py
+++ b/src/saml2/response.py
@@ -722,12 +722,11 @@ class AuthnResponse(StatusResponse):
         return True
 
     def _holder_of_key_confirmed(self, data):
-        if not data or not data.key_info:
+        if not data or not data.extension_elements:
             return False
 
         has_keyinfo = False
-        key_info = data.key_info or ()
-        for element in extension_elements_to_elements(key_info,
+        for element in extension_elements_to_elements(data.extension_elements,
                                                       [samlp, saml, xenc, ds]):
             if isinstance(element, ds.KeyInfo):
                 has_keyinfo = True
diff --git a/src/saml2/saml.py b/src/saml2/saml.py
index 0d6728e5..bdb1ec60 100644
--- a/src/saml2/saml.py
+++ b/src/saml2/saml.py
@@ -482,12 +482,8 @@ class SubjectConfirmationDataType_(SamlBase):
     c_any = {"namespace": "##any", "processContents": "lax", "minOccurs": "0",
              "maxOccurs": "unbounded"}
     c_any_attribute = {"namespace": "##other", "processContents": "lax"}
-    c_children['{http://www.w3.org/2000/09/xmldsig#}KeyInfo'] = ('key_info',
-                                                                 [ds.KeyInfo])
-    c_cardinality['key_info'] = {"min": 0, "max": 1}
 
     def __init__(self,
-                 key_info=None,
                  not_before=None,
                  not_on_or_after=None,
                  recipient=None,
@@ -500,7 +496,6 @@ class SubjectConfirmationDataType_(SamlBase):
                           text=text,
                           extension_elements=extension_elements,
                           extension_attributes=extension_attributes)
-        self.key_info = key_info
         self.not_before = not_before
         self.not_on_or_after = not_on_or_after
         self.recipient = recipient
diff --git a/tests/test_02_saml.py b/tests/test_02_saml.py
index ec5f15d6..b8083273 100644
--- a/tests/test_02_saml.py
+++ b/tests/test_02_saml.py
@@ -886,7 +886,8 @@ class TestSubjectConfirmation:
         assert sc.subject_confirmation_data.recipient == "recipient"
         assert sc.subject_confirmation_data.in_response_to == "responseID"
         assert sc.subject_confirmation_data.address == "127.0.0.1"
-        assert sc.subject_confirmation_data.key_info is None
+        key_info = sc.subject_confirmation_data.extensions_as_elements(ds.KeyInfo.c_tag, ds)
+        assert len(key_info) == 0
 
     def testHolderOfKeyUsingTestData(self):
         """Test subject_confirmation_from_string() using test data for 'holder-of-key' SubjectConfirmation"""
@@ -898,7 +899,7 @@ class TestSubjectConfirmation:
         assert sc.subject_confirmation_data.not_on_or_after == "2007-09-14T01:05:02Z"
         assert sc.subject_confirmation_data.recipient == "recipient"
         assert sc.subject_confirmation_data.in_response_to == "responseID"
-        key_info = sc.subject_confirmation_data.key_info
+        key_info = sc.subject_confirmation_data.extensions_as_elements(ds.KeyInfo.c_tag, ds)
         assert len(key_info) == 1
         assert len(key_info[0].x509_data) == 1
 
diff --git a/tests/test_93_hok.py b/tests/test_93_hok.py
index dc6aac6e..962d2be5 100644
--- a/tests/test_93_hok.py
+++ b/tests/test_93_hok.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
+from saml2 import xmldsig as ds
 from saml2.response import authn_response, VerificationError
 from saml2.config import config_factory
 
@@ -18,8 +19,10 @@ class TestHolderOfKeyResponse:
 
         assert resp.get_subject() is not None
         assert len(resp.assertion.subject.subject_confirmation) == 2
-        actual_hok_certs = [sc.subject_confirmation_data.key_info[0].x509_data[0].x509_certificate.text.strip() 
-                            for sc in resp.assertion.subject.subject_confirmation]
+        key_infos = [sc.subject_confirmation_data.extensions_as_elements(ds.KeyInfo.c_tag, ds)[0]
+                    for sc in resp.assertion.subject.subject_confirmation]
+        actual_hok_certs = [key_info_element.x509_data[0].x509_certificate.text.strip() 
+                            for key_info_element in key_infos]
         assert actual_hok_certs == self._expected_hok_certs()
 
     def _expected_hok_certs(self):