Paging in ms graph queries (#2492)

* paging in ms graph queries * update changelog * merge dict fix * rename sample * add log * add log
microsoft · Aug 23, 2022 · 1273b5c · 1273b5c
1 parent 21e2e89
commit 1273b5c
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 13 deletions.
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
@@ -89,6 +89,21 @@
             "false"
           ]
         },
+        {
+          "name": "E2E Extended AAD",
+          "type": "python",
+          "request": "launch",
+          "module": "pytest",
+          "justMyCode": true,
+          "cwd": "${workspaceFolder}/e2e_tests/",
+          "preLaunchTask": "Copy_env_file_for_e2e_debug",
+          "args": [
+            "-m",
+            "extended_aad",
+            "--verify",
+            "false"
+          ]
+        },
         {
           "name": "E2E Shared Services",
           "type": "python",

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -26,6 +26,7 @@ BUG FIXES:
 * Enable route table on the Airlock Processor subnet ([#2414](https://github.com/microsoft/AzureTRE/pull/2414))
 * Support for _Standard_ app service plan SKUs ([#2415](https://github.com/microsoft/AzureTRE/pull/2415))
 * Fix Azure ML Workspace deletion ([#2452](https://github.com/microsoft/AzureTRE/pull/2452))
+* Get all pages in MS Graph queries ([#2492](https://github.com/microsoft/AzureTRE/pull/2492))
 
 ## 0.4.1 (August 03, 2022)
 

diff --git a/api_app/_version.py b/api_app/_version.py
@@ -1 +1 @@
-__version__ = "0.4.18"
+__version__ = "0.4.19"
diff --git a/api_app/api/routes/workspaces.py b/api_app/api/routes/workspaces.py
@@ -79,6 +79,7 @@ async def retrieve_workspace_by_workspace_id(user=Depends(get_current_tre_user_o
     if access_service.get_workspace_role(user, workspace, user_role_assignments) != WorkspaceRole.NoRole:
         return WorkspaceInResponse(workspace=workspace)
     else:
+        logging.debug("User doesn't have roles in workspace.")
         raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=strings.ACCESS_USER_IS_NOT_OWNER_OR_RESEARCHER)
 
 

diff --git a/api_app/services/aad_authentication.py b/api_app/services/aad_authentication.py
@@ -287,27 +287,45 @@ def _get_app_auth_info(self, client_id: str) -> dict:
 
         return authInfo
 
-    def _get_role_assignment_graph_data_for_user(self, user_id: str) -> dict:
+    def _ms_graph_query(self, url: str, http_method: str, json=None) -> dict:
         msgraph_token = self._get_msgraph_token()
+        auth_headers = self._get_auth_header(msgraph_token)
+        graph_data = {}
+        while True:
+            if not url:
+                break
+            logging.debug(f"Making request to: {url}")
+            if json:
+                response = requests.request(method=http_method, url=url, json=json, headers=auth_headers)
+            else:
+                response = requests.request(method=http_method, url=url, headers=auth_headers)
+            url = ""
+            if response.status_code == 200:
+                json_response = response.json()
+                graph_data = merge_dict(graph_data, json_response)
+                if '@odata.nextLink' in json_response:
+                    url = json_response['@odata.nextLink']
+            else:
+                logging.error(f"MS Graph query to: {url} failed with status code {response.status_code}")
+                logging.error(f"Full response: {response}")
+        return graph_data
+
+    def _get_role_assignment_graph_data_for_user(self, user_id: str) -> dict:
         user_endpoint = f"https://graph.microsoft.com/v1.0/users/{user_id}/appRoleAssignments"
-        graph_data = requests.get(user_endpoint, headers=self._get_auth_header(msgraph_token)).json()
+        graph_data = self._ms_graph_query(user_endpoint, "GET")
         return graph_data
 
     def _get_role_assignment_graph_data_for_service_principal(self, principal_id: str) -> dict:
-        msgraph_token = self._get_msgraph_token()
-        user_endpoint = f"https://graph.microsoft.com/v1.0/servicePrincipals/{principal_id}/appRoleAssignments"
-        graph_data = requests.get(user_endpoint, headers=self._get_auth_header(msgraph_token)).json()
+        svc_principal_endpoint = f"https://graph.microsoft.com/v1.0/servicePrincipals/{principal_id}/appRoleAssignments"
+        graph_data = self._ms_graph_query(svc_principal_endpoint, "GET")
         return graph_data
 
     def _get_identity_type(self, id: str) -> str:
-        msgraph_token = self._get_msgraph_token()
         objects_endpoint = "https://graph.microsoft.com/v1.0/directoryObjects/getByIds"
         request_body = {"ids": [id], "types": ["user", "servicePrincipal"]}
-        graph_data = requests.post(
-            objects_endpoint,
-            headers=self._get_auth_header(msgraph_token),
-            json=request_body
-        ).json()
+        graph_data = self._ms_graph_query(objects_endpoint, "POST", json=request_body)
+
+        logging.debug(graph_data)
 
         if "value" not in graph_data or len(graph_data["value"]) != 1:
             logging.debug(graph_data)
@@ -344,13 +362,14 @@ def get_identity_role_assignments(self, user_id: str) -> List[RoleAssignment]:
         elif identity_type == "#microsoft.graph.servicePrincipal":
             graph_data = self._get_role_assignment_graph_data_for_service_principal(user_id)
         else:
-            logging.debug(graph_data)
             raise AuthConfigValidationError(f"{strings.ACCESS_UNHANDLED_ACCOUNT_TYPE} {identity_type}")
 
         if 'value' not in graph_data:
             logging.debug(graph_data)
             raise AuthConfigValidationError(f"{strings.ACCESS_UNABLE_TO_GET_ROLE_ASSIGNMENTS_FOR_USER} {user_id}")
 
+        logging.debug(graph_data)
+
         return [RoleAssignment(role_assignment['resourceId'], role_assignment['appRoleId']) for role_assignment in graph_data['value']]
 
     def get_workspace_role(self, user: User, workspace: Workspace, user_role_assignments: List[RoleAssignment]) -> WorkspaceRole:
@@ -370,3 +389,15 @@ def get_workspace_role(self, user: User, workspace: Workspace, user_role_assignm
         if RoleAssignment(resource_id=workspace_sp_id, role_id=workspace.properties['app_role_id_workspace_airlock_manager']) in user_role_assignments:
             return WorkspaceRole.AirlockManager
         return WorkspaceRole.NoRole
+
+
+def merge_dict(d1, d2):
+    dd = defaultdict(list)
+
+    for d in (d1, d2):
+        for key, value in d.items():
+            if isinstance(value, list):
+                dd[key].extend(value)
+            else:
+                dd[key].append(value)
+    return dict(dd)
diff --git a/e2e_tests/.env.tmpl → e2e_tests/.env.sample b/e2e_tests/.env.tmpl → e2e_tests/.env.sample