From 2e78ef8434bbf17fe330df9d001c0a0a888f0810 Mon Sep 17 00:00:00 2001
From: Noah Paige <69586985+noah-paige@users.noreply.github.com>
Date: Fri, 1 Dec 2023 16:36:12 -0500
Subject: [PATCH] Add Pagination to Return Full List Cognito Groups (#891)

### Feature or Bugfix
<!-- please choose -->
- Bugfix

### Detail
- When inviting a new team to an Organization or Environment we call the
API `listCognitoGroups` which calls boto3's `list_groups()` on the
Cognito User Pool in the backend
- By default - if the the Cognito User Pool has over 25 Groups in it,
the `list_groups()` will only return the first 25 in the response as a
default limit
- This PR adds a paginator to `list_groups()` to that backend API
returns the full list of groups in the Cognito User Pool that can be
invited to the Organization or Environment in the data.all UI


### Relates
- N/A

### Security
Please answer the questions below briefly where applicable, or write
`N/A`. Based on
[OWASP 10](https://owasp.org/Top10/en/).

- Does this PR introduce or modify any input fields or queries - this
includes
fetching data from storage outside the application (e.g. a database, an
S3 bucket)?
  - Is the input sanitized?
- What precautions are you taking before deserializing the data you
consume?
  - Is injection prevented by parametrizing queries?
  - Have you ensured no `eval` or similar functions are used?
- Does this PR introduce any functionality or component that requires
authorization?
- How have you ensured it respects the existing AuthN/AuthZ mechanisms?
  - Are you logging failed auth attempts?
- Are you using or adding any cryptographic features?
  - Do you use a standard proven implementations?
  - Are the used keys controlled by the customer? Where are they stored?
- Are you introducing any new policies/roles/users?
  - Have you used the least-privilege principle? How?


By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 license.
---
 backend/dataall/base/aws/cognito.py           | 20 +++++++++++++++----
 .../core/cognito_groups/api/resolvers.py      |  4 ++--
 tests/core/cognito_groups/test_group.py       |  2 +-
 3 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/backend/dataall/base/aws/cognito.py b/backend/dataall/base/aws/cognito.py
index 5e63172ab..9bddef3f3 100644
--- a/backend/dataall/base/aws/cognito.py
+++ b/backend/dataall/base/aws/cognito.py
@@ -18,7 +18,14 @@ def get_user_emailids_from_group(self, groupName):
             parameter_path = f'/dataall/{envname}/cognito/userpool'
             ssm = boto3.client('ssm', region_name=os.getenv('AWS_REGION', 'eu-west-1'))
             user_pool_id = ssm.get_parameter(Name=parameter_path)['Parameter']['Value']
-            cognito_user_list = self.client.list_users_in_group(UserPoolId=user_pool_id, GroupName=groupName)["Users"]
+            paginator = self.client.get_paginator('list_users_in_group')
+            pages = paginator.paginate(
+                UserPoolId=user_pool_id,
+                GroupName=groupName
+            )
+            cognito_user_list = []
+            for page in pages:
+                cognito_user_list += page['Users']
             group_email_ids = []
             attributes = []
             # Make a flat list
@@ -40,15 +47,20 @@ def get_user_emailids_from_group(self, groupName):
 
     @staticmethod
     def list_cognito_groups(envname: str, region: str):
+        user_pool_id = None
+        groups = []
         try:
             parameter_path = f'/dataall/{envname}/cognito/userpool'
             ssm = boto3.client('ssm', region_name=region)
             user_pool_id = ssm.get_parameter(Name=parameter_path)['Parameter']['Value']
             cognito = boto3.client('cognito-idp', region_name=region)
-            groups = cognito.list_groups(UserPoolId=user_pool_id)['Groups']
+            paginator = cognito.get_paginator('list_groups')
+            pages = paginator.paginate(UserPoolId=user_pool_id)
+            for page in pages:
+                groups += [gr['GroupName'] for gr in page['Groups']]
         except Exception as e:
             log.error(
                 f'Failed to list groups of user pool {user_pool_id} due to {e}'
             )
-        else:
-            return groups
+            raise e
+        return groups
diff --git a/backend/dataall/core/cognito_groups/api/resolvers.py b/backend/dataall/core/cognito_groups/api/resolvers.py
index 7040cf569..9487447a2 100644
--- a/backend/dataall/core/cognito_groups/api/resolvers.py
+++ b/backend/dataall/core/cognito_groups/api/resolvers.py
@@ -68,6 +68,6 @@ def list_cognito_groups(context, source, filter: dict = None):
     invited_group_uris = [item.groupUri for item in invited_groups]
     res = []
     for group in groups:
-        if group['GroupName'] not in invited_group_uris:
-            res.append({"groupName": group['GroupName']})
+        if group not in invited_group_uris:
+            res.append({"groupName": group})
     return res
diff --git a/tests/core/cognito_groups/test_group.py b/tests/core/cognito_groups/test_group.py
index 9eb8ec26e..e2474f01f 100644
--- a/tests/core/cognito_groups/test_group.py
+++ b/tests/core/cognito_groups/test_group.py
@@ -3,7 +3,7 @@
 def test_list_cognito_groups_env(client, env_fixture, group, module_mocker):
     module_mocker.patch(
         'dataall.base.aws.cognito.Cognito.list_cognito_groups',
-        return_value=[{"GroupName": 'cognitos'}, {"GroupName": 'testadmins'}],
+        return_value=['cognitos', 'testadmins'],
     )
     response = client.query(
         """