From dc1598cce5d8f8d3467e350c38528cf7a96cc057 Mon Sep 17 00:00:00 2001
From: Will <will@digitaloutback.co.uk>
Date: Thu, 8 Aug 2024 18:53:13 +0100
Subject: [PATCH] Allow claims-like attribute keys (#253)

* added feature and tests
* added author
* cleanup
* cleanup and add .idea (intellij) to .gitignore
* remove nested .idea/
* cleanup debugging
* Remove unused file
* Reformat with ruff
---
 .gitignore                           |  1 +
 AUTHORS.md                           |  1 +
 django_saml2_auth/saml.py            | 10 ++++----
 django_saml2_auth/tests/test_saml.py | 38 ++++++++++++++++++++++++++++
 4 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/.gitignore b/.gitignore
index 3adb2f0..9de9b0e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -63,3 +63,4 @@ target/
 
 # IDEs
 .vscode/
+.idea/
diff --git a/AUTHORS.md b/AUTHORS.md
index b315885..7e4a93e 100644
--- a/AUTHORS.md
+++ b/AUTHORS.md
@@ -59,3 +59,4 @@ an issue.
 - [Santiago Gandolfo](https://github.com/santigandolfo)
 - [Greg Wong](https://github.com/gregorywong)
 - [Michael V. Battista](https://github.com/mvbattista)
+- [William Abbott](https://github.com/wrabit)
diff --git a/django_saml2_auth/saml.py b/django_saml2_auth/saml.py
index 00c5382..8085bee 100644
--- a/django_saml2_auth/saml.py
+++ b/django_saml2_auth/saml.py
@@ -402,15 +402,15 @@ def extract_user_identity(user_identity: Dict[str, Any]) -> Dict[str, Optional[A
     )
 
     user = {}
-    user["email"] = dictor(user_identity, f"{email_field}/0", pathsep="/")  # Path includes "."
-    user["username"] = dictor(user_identity, f"{username_field}/0", pathsep="/")
-    user["first_name"] = dictor(user_identity, f"{firstname_field}/0", pathsep="/")
-    user["last_name"] = dictor(user_identity, f"{lastname_field}/0", pathsep="/")
+    user["email"] = dictor(user_identity, f"{email_field}|0", pathsep="|")  # Path includes "."
+    user["username"] = dictor(user_identity, f"{username_field}|0", pathsep="|")
+    user["first_name"] = dictor(user_identity, f"{firstname_field}|0", pathsep="|")
+    user["last_name"] = dictor(user_identity, f"{lastname_field}|0", pathsep="|")
 
     token_required = dictor(saml2_auth_settings, "TOKEN_REQUIRED", default=True)
     if token_required:
         token_field = dictor(saml2_auth_settings, "ATTRIBUTES_MAP.token", default="token")
-        user["token"] = dictor(user_identity, f"{token_field}.0")
+        user["token"] = dictor(user_identity, f"{token_field}|0", pathsep="|")
 
     if user["email"]:
         user["email"] = user["email"].lower()
diff --git a/django_saml2_auth/tests/test_saml.py b/django_saml2_auth/tests/test_saml.py
index 1785da4..93c2586 100644
--- a/django_saml2_auth/tests/test_saml.py
+++ b/django_saml2_auth/tests/test_saml.py
@@ -123,6 +123,21 @@ def get_user_identity() -> Mapping[str, List[str]]:
     }
 
 
+def get_user_identify_with_slashed_keys() -> Mapping[str, List[str]]:
+    """Fixture for returning user identity produced by pysaml2 with slashed, claim-like keys.
+
+    Returns:
+        dict: keys are SAML attributes and values are lists of attribute values
+    """
+    return {
+        "http://schemas.org/user/username": ["test@example.com"],
+        "http://schemas.org/user/claim2.0/email": ["test@example.com"],
+        "http://schemas.org/user/claim2.0/first_name": ["John"],
+        "http://schemas.org/user/claim2.0/last_name": ["Doe"],
+        "http://schemas.org/auth/server/token": ["TOKEN"],
+    }
+
+
 def mock_parse_authn_request_response(
     self: Saml2Client, response: AuthnResponse, binding: str
 ) -> "MockAuthnResponse":  # type: ignore # noqa: F821
@@ -447,6 +462,29 @@ def test_extract_user_identity_success():
     assert result["user_identity"] == get_user_identity()
 
 
+def test_extract_user_identity_with_slashed_attribute_keys_success(settings: SettingsWrapper):
+    """Test extract_user_identity function to verify if it correctly extracts user identity
+    information from a (pysaml2) parsed SAML response with slashed attribute keys."""
+    settings.SAML2_AUTH = {
+        "ATTRIBUTES_MAP": {
+            "email": "http://schemas.org/user/claim2.0/email",
+            "username": "http://schemas.org/user/username",
+            "first_name": "http://schemas.org/user/claim2.0/first_name",
+            "last_name": "http://schemas.org/user/claim2.0/last_name",
+            "token": "http://schemas.org/auth/server/token",
+        }
+    }
+
+    result = extract_user_identity(get_user_identify_with_slashed_keys())  # type: ignore
+
+    assert len(result) == 6
+    assert result["username"] == result["email"] == "test@example.com"
+    assert result["first_name"] == "John"
+    assert result["last_name"] == "Doe"
+    assert result["token"] == "TOKEN"
+    assert result["user_identity"] == get_user_identify_with_slashed_keys()
+
+
 def test_extract_user_identity_token_not_required(settings: SettingsWrapper):
     """Test extract_user_identity function to verify if it correctly extracts user identity
     information from a (pysaml2) parsed SAML response when token is not required."""