From dae4f4947e8d092c87246befc8c2694afc096b7e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adri=C3=A1n=20Tom=C3=A1s?=
 <22175056+adriantomas@users.noreply.github.com>
Date: Fri, 24 Feb 2023 23:14:08 +0100
Subject: [PATCH] Cognito: add groups to idToken (#5977)

---
 moto/cognitoidp/models.py                |  3 ++
 tests/test_cognitoidp/test_cognitoidp.py | 53 ++++++++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/moto/cognitoidp/models.py b/moto/cognitoidp/models.py
index 3c6e97c62..0c1cda80e 100644
--- a/moto/cognitoidp/models.py
+++ b/moto/cognitoidp/models.py
@@ -567,6 +567,9 @@ class CognitoIdpUserPool(BaseModel):
 
     def create_id_token(self, client_id: str, username: str) -> Tuple[str, int]:
         extra_data = self.get_user_extra_data_by_client_id(client_id, username)
+        user = self._get_user(username)
+        if len(user.groups) > 0:
+            extra_data["cognito:groups"] = [group.group_name for group in user.groups]
         id_token, expires_in = self.create_jwt(
             client_id, username, "id", extra_data=extra_data
         )
diff --git a/tests/test_cognitoidp/test_cognitoidp.py b/tests/test_cognitoidp/test_cognitoidp.py
index c8425128c..fb91777bc 100644
--- a/tests/test_cognitoidp/test_cognitoidp.py
+++ b/tests/test_cognitoidp/test_cognitoidp.py
@@ -1472,6 +1472,59 @@ def test_group_in_access_token():
     claims["cognito:groups"].should.equal([group_name])
 
 
+@mock_cognitoidp
+def test_group_in_id_token():
+    conn = boto3.client("cognito-idp", "us-west-2")
+
+    username = str(uuid.uuid4())
+    temporary_password = "P2$Sword"
+    user_pool_id = conn.create_user_pool(PoolName=str(uuid.uuid4()))["UserPool"]["Id"]
+    user_attribute_name = str(uuid.uuid4())
+    user_attribute_value = str(uuid.uuid4())
+    group_name = str(uuid.uuid4())
+    client_id = conn.create_user_pool_client(
+        UserPoolId=user_pool_id,
+        ClientName=str(uuid.uuid4()),
+        ReadAttributes=[user_attribute_name],
+    )["UserPoolClient"]["ClientId"]
+
+    conn.create_group(GroupName=group_name, UserPoolId=user_pool_id)
+
+    conn.admin_create_user(
+        UserPoolId=user_pool_id,
+        Username=username,
+        TemporaryPassword=temporary_password,
+        UserAttributes=[{"Name": user_attribute_name, "Value": user_attribute_value}],
+    )
+
+    conn.admin_add_user_to_group(
+        UserPoolId=user_pool_id, Username=username, GroupName=group_name
+    )
+
+    result = conn.admin_initiate_auth(
+        UserPoolId=user_pool_id,
+        ClientId=client_id,
+        AuthFlow="ADMIN_NO_SRP_AUTH",
+        AuthParameters={"USERNAME": username, "PASSWORD": temporary_password},
+    )
+
+    # A newly created user is forced to set a new password
+    result["ChallengeName"].should.equal("NEW_PASSWORD_REQUIRED")
+    result["Session"].should_not.equal(None)
+
+    # This sets a new password and logs the user in (creates tokens)
+    new_password = "P2$Sword"
+    result = conn.respond_to_auth_challenge(
+        Session=result["Session"],
+        ClientId=client_id,
+        ChallengeName="NEW_PASSWORD_REQUIRED",
+        ChallengeResponses={"USERNAME": username, "NEW_PASSWORD": new_password},
+    )
+
+    claims = jwt.get_unverified_claims(result["AuthenticationResult"]["IdToken"])
+    claims["cognito:groups"].should.equal([group_name])
+
+
 @mock_cognitoidp
 def test_create_group_with_duplicate_name_raises_error():
     conn = boto3.client("cognito-idp", "us-west-2")