Add support for authn_mappings (#255)

* add support for authnmappings * add tests * Update scripts/cleanup_org.py Co-authored-by: Kevin Zou <17015060+nkzou@users.noreply.github.com> * Update scripts/cleanup_org.py Co-authored-by: Kevin Zou <17015060+nkzou@users.noreply.github.com> --------- Co-authored-by: Kevin Zou <17015060+nkzou@users.noreply.github.com>
DataDog · Jun 27, 2024 · 5c10d6f · 5c10d6f
1 parent ab8f998
commit 5c10d6f
Show file tree

Hide file tree

Showing 15 changed files with 2,700 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -208,6 +208,7 @@ When running againts multiple destination organizations, a seperate working dire
 
 | Resource                               | Description                                              |
 |----------------------------------------|----------------------------------------------------------|
+| authn_mappings                         | Sync Datadog authn mappings.                            |
 | dashboard_lists                        | Sync Datadog dashboard lists.                            |
 | dashboards                             | Sync Datadog dashboards.                                 |
 | downtime_schedules                     | Sync Datadog downtimes.                                  |
@@ -249,6 +250,7 @@ See [Supported resources](#supported-resources) section below for potential reso
 
 | Resource                               | Dependencies                                                     |
 |----------------------------------------|------------------------------------------------------------------|
+| authn_mappings                         | roles, teams                                                     |
 | dashboard_lists                        | dashboards                                                       |
 | dashboards                             | monitors, roles, powerpacks, service_level_objectives            |
 | downtime_schedules                     | monitors                                                         |

diff --git a/datadog_sync/model/authn_mappings.py b/datadog_sync/model/authn_mappings.py
@@ -0,0 +1,83 @@
+# Unless explicitly stated otherwise all files in this repository are licensed
+# under the 3-clause BSD style license (see LICENSE).
+# This product includes software developed at Datadog (https://www.datadoghq.com/).
+# Copyright 2019 Datadog, Inc.
+from typing import Optional, List, Dict, Tuple
+
+from datadog_sync.utils.base_resource import BaseResource, ResourceConfig
+from datadog_sync.utils.custom_client import CustomClient
+
+
+class AuthNMappings(BaseResource):
+    resource_type = "authn_mappings"
+    resource_config = ResourceConfig(
+        base_path="/api/v2/authn_mappings",
+        excluded_attributes=[
+            "id",
+            "attributes.created_at",
+            "attributes.modified_at",
+            "attributes.saml_assertion_attribute_id",
+            "relationships.saml_assertion_attribute",
+        ],
+        resource_connections={"roles": ["relationships.role.data.id"], "teams": ["relationships.team.data.id"]},
+    )
+    # Additional AuthNMappings specific attributes
+
+    async def get_resources(self, client: CustomClient) -> List[Dict]:
+        role_resp = await client.paginated_request(client.get)(
+            self.resource_config.base_path, params={"resource_type": "role"}
+        )
+        team_resp = await client.paginated_request(client.get)(
+            self.resource_config.base_path, params={"resource_type": "team"}
+        )
+
+        return role_resp + team_resp
+
+    async def import_resource(self, _id: Optional[str] = None, resource: Optional[Dict] = None) -> Tuple[str, Dict]:
+        if _id:
+            source_client = self.config.source_client
+            resource = (await source_client.get(self.resource_config.base_path + f"/{_id}"))["data"]
+
+        return resource["id"], resource
+
+    async def pre_resource_action_hook(self, _id, resource: Dict) -> None:
+        pass
+
+    async def pre_apply_hook(self) -> None:
+        pass
+
+    async def create_resource(self, _id: str, resource: Dict) -> Tuple[str, Dict]:
+        destination_client = self.config.destination_client
+        payload = {"data": resource}
+        resp = await destination_client.post(self.resource_config.base_path, payload)
+        self.remove_null_relationships(resp)
+
+        return _id, resp["data"]
+
+    async def update_resource(self, _id: str, resource: Dict) -> Tuple[str, Dict]:
+        destination_client = self.config.destination_client
+        d_id = self.resource_config.destination_resources[_id]["id"]
+        resource["id"] = d_id
+        payload = {"data": resource}
+        resp = await destination_client.patch(self.resource_config.base_path + f"/{d_id}", payload)
+        self.remove_null_relationships(resp)
+
+        return _id, resp["data"]
+
+    async def delete_resource(self, _id: str) -> None:
+        destination_client = self.config.destination_client
+        await destination_client.delete(
+            self.resource_config.base_path + f"/{self.resource_config.destination_resources[_id]['id']}"
+        )
+
+    def connect_id(self, key: str, r_obj: Dict, resource_to_connect: str) -> Optional[List[str]]:
+        return super(AuthNMappings, self).connect_id(key, r_obj, resource_to_connect)
+
+    @staticmethod
+    def remove_null_relationships(resp: Dict) -> None:
+        if resp["data"]["relationships"].get("role", {}).get("data") is None:
+            resp["data"]["relationships"].pop("role", None)
+        if resp["data"]["relationships"].get("team", {}).get("data") is None:
+            resp["data"]["relationships"].pop("team", None)
+
+        return resp
diff --git a/datadog_sync/models/__init__.py b/datadog_sync/models/__init__.py
@@ -4,6 +4,7 @@
 # Copyright 2019 Datadog, Inc.
 # ruff: noqa
 
+from datadog_sync.model.authn_mappings import AuthNMappings
 from datadog_sync.model.dashboard_lists import DashboardLists
 from datadog_sync.model.dashboards import Dashboards
 from datadog_sync.model.downtime_schedules import DowntimeSchedules

diff --git a/datadog_sync/utils/custom_client.py b/datadog_sync/utils/custom_client.py
@@ -44,9 +44,9 @@ async def wrapper(*args, **kwargs):
                         except ValueError:
                             sleep_duration = retry_count * default_backoff
                         if (sleep_duration + time.time()) > timeout:
-                            log.debug("retry timeout has or will exceed timeout duration")
+                            log.debug(f"{e}. retry timeout has or will exceed timeout duration")
                             raise CustomClientHTTPError(e, message=err_text)
-                        log.debug(f"retrying request after {sleep_duration}s")
+                        log.debug(f"{e}. retrying request after {sleep_duration}s")
                         await asyncio.sleep(sleep_duration)
                         retry_count += 1
                         continue
@@ -55,7 +55,7 @@ async def wrapper(*args, **kwargs):
                         if (sleep_duration + time.time()) > timeout:
                             log.debug("retry timeout has or will exceed timeout duration")
                             raise CustomClientHTTPError(e, message=err_text)
-                        log.debug(f"retrying request after {sleep_duration}s")
+                        log.debug(f"{e}. retrying request after {sleep_duration}s")
                         await asyncio.sleep(retry_count * default_backoff)
                         retry_count += 1
                         continue

diff --git a/scripts/cleanup_org.py b/scripts/cleanup_org.py
@@ -24,6 +24,7 @@ def __init__(self):
         self.validate_org()
 
         # Delete all supported resources
+        self.cleanup_authn_mappings()
         self.cleanup_service_level_objectives()
         self.cleanup_slo_corrections()
         self.cleanup_synthetics_tests()
@@ -59,6 +60,15 @@ def validate_org(self):
             print("Error getting organization. Validate api+app keys %s: %s", url, e)
             exit(1)
 
+    def cleanup_authn_mappings(
+        self,
+    ):
+        path = "/api/v2/authn_mappings"
+        res_role = self.get_resources(path, params={"resource_type": "role"})
+        res_team = self.get_resources(path, params={"resource_type": "team"})
+        for resource in res_role["data"] + res_team["data"]:
+            self.delete_resource(resource["id"], path)
+
     def cleanup_dashboards(
         self,
     ):

diff --git a/...es/cassettes/test_authn_mappings/TestAuthNMappingsResources.test_no_resource_diffs.frozen b/...es/cassettes/test_authn_mappings/TestAuthNMappingsResources.test_no_resource_diffs.frozen
@@ -0,0 +1 @@
+2024-06-14T13:40:03.401775-04:00
diff --git a/...ces/cassettes/test_authn_mappings/TestAuthNMappingsResources.test_resource_cleanup.frozen b/...ces/cassettes/test_authn_mappings/TestAuthNMappingsResources.test_resource_cleanup.frozen
@@ -0,0 +1 @@
+2024-06-14T13:40:03.411769-04:00
diff --git a/...urces/cassettes/test_authn_mappings/TestAuthNMappingsResources.test_resource_cleanup.yaml b/...urces/cassettes/test_authn_mappings/TestAuthNMappingsResources.test_resource_cleanup.yaml
@@ -0,0 +1,82 @@
+interactions:
+- request:
+    body: null
+    headers:
+      Content-Type:
+      - application/json
+    method: DELETE
+    uri: https://api.datadoghq.eu/api/v2/authn_mappings/2158979a-2a75-11ef-b6cb-da7ad0900005
+  response:
+    body:
+      string: ''
+    headers:
+      Content-Type:
+      - text/html; charset=utf-8
+    status:
+      code: 204
+      message: No Content
+- request:
+    body: null
+    headers:
+      Content-Type:
+      - application/json
+    method: DELETE
+    uri: https://api.datadoghq.eu/api/v2/authn_mappings/2151d2ca-2a75-11ef-9510-da7ad0900005
+  response:
+    body:
+      string: ''
+    headers:
+      Content-Type:
+      - text/html; charset=utf-8
+    status:
+      code: 204
+      message: No Content
+- request:
+    body: null
+    headers:
+      Content-Type:
+      - application/json
+    method: DELETE
+    uri: https://api.datadoghq.eu/api/v2/authn_mappings/21593696-2a75-11ef-bbf5-da7ad0900005
+  response:
+    body:
+      string: ''
+    headers:
+      Content-Type:
+      - text/html; charset=utf-8
+    status:
+      code: 204
+      message: No Content
+- request:
+    body: null
+    headers:
+      Content-Type:
+      - application/json
+    method: DELETE
+    uri: https://api.datadoghq.eu/api/v2/authn_mappings/21534c0e-2a75-11ef-8c13-da7ad0900005
+  response:
+    body:
+      string: ''
+    headers:
+      Content-Type:
+      - text/html; charset=utf-8
+    status:
+      code: 204
+      message: No Content
+- request:
+    body: null
+    headers:
+      Content-Type:
+      - application/json
+    method: DELETE
+    uri: https://api.datadoghq.eu/api/v2/authn_mappings/214e8a66-2a75-11ef-9113-da7ad0900005
+  response:
+    body:
+      string: ''
+    headers:
+      Content-Type:
+      - text/html; charset=utf-8
+    status:
+      code: 204
+      message: No Content
+version: 1
diff --git a/...rces/cassettes/test_authn_mappings/TestAuthNMappingsResources.test_resource_import.frozen b/...rces/cassettes/test_authn_mappings/TestAuthNMappingsResources.test_resource_import.frozen
@@ -0,0 +1 @@
+2024-06-14T13:40:01.957206-04:00
diff --git a/...ources/cassettes/test_authn_mappings/TestAuthNMappingsResources.test_resource_import.yaml b/...ources/cassettes/test_authn_mappings/TestAuthNMappingsResources.test_resource_import.yaml
@@ -0,0 +1,94 @@
+interactions:
+- request:
+    body: null
+    headers:
+      Content-Type:
+      - application/json
+    method: GET
+    uri: https://api.datadoghq.com/api/v2/authn_mappings?page%5Bnumber%5D=0&page%5Bsize%5D=100&resource_type=role
+  response:
+    body:
+      string: '{"data": [{"type": "authn_mappings", "id": "29556498-29bf-11ef-abeb-da7ad0900002",
+        "attributes": {"attribute_key": "Member-of", "attribute_value": "TestOrgFoo",
+        "created_at": "2024-06-13T19:57:28.072750+00:00", "modified_at": "2024-06-13T19:58:14.970554+00:00"},
+        "relationships": {"role": {"data": {"type": "roles", "id": "fa017a37-bfcb-11eb-a4d7-da7ad0900002"}}}},
+        {"type": "authn_mappings", "id": "29562f9a-29bf-11ef-a243-da7ad0900002", "attributes":
+        {"attribute_key": "Member-of", "attribute_value": "TestOrg", "created_at":
+        "2024-06-13T19:57:28.077725+00:00", "modified_at": "2024-06-13T19:58:14.967881+00:00"},
+        "relationships": {"role": {"data": {"type": "roles", "id": "fa017a37-bfcb-11eb-a4d7-da7ad0900002"}}}},
+        {"type": "authn_mappings", "id": "454a4024-29bf-11ef-8d8e-da7ad0900002", "attributes":
+        {"attribute_key": "Member-of", "attribute_value": "TestOrgFooBar", "created_at":
+        "2024-06-13T19:58:14.975381+00:00", "modified_at": "2024-06-13T19:58:14.975381+00:00"},
+        "relationships": {"role": {"data": {"type": "roles", "id": "fa017a37-bfcb-11eb-a4d7-da7ad0900002"}}}}],
+        "included": [{"type": "roles", "id": "fa017a37-bfcb-11eb-a4d7-da7ad0900002",
+        "attributes": {"name": "Datadog Read Only Role", "created_at": "2021-05-28T15:47:16.084615+00:00",
+        "modified_at": "2021-05-28T15:47:16.084615+00:00"}, "relationships": {"permissions":
+        {"data": [{"type": "permissions", "id": "5e605652-dd12-11e8-9e53-375565b8970e"},
+        {"type": "permissions", "id": "6f66600e-dd12-11e8-9e55-7f30fbb45e73"}, {"type":
+        "permissions", "id": "d90f6830-d3d8-11e9-a77a-b3404e5e9ee2"}, {"type": "permissions",
+        "id": "4441648c-d8b1-11e9-a77a-1b899a04b304"}, {"type": "permissions", "id":
+        "1af86ce4-7823-11ea-93dc-d7cad1b1c6cb"}, {"type": "permissions", "id": "b382b982-8535-11ea-93de-2bf1bdf20798"},
+        {"type": "permissions", "id": "7314eb20-aa58-11ea-95e2-6fb6e4a451d5"}, {"type":
+        "permissions", "id": "80de1ec0-aa58-11ea-95e2-aff381626d5d"}, {"type": "permissions",
+        "id": "5025ee24-f923-11ea-adbc-576ea241df8d"}, {"type": "permissions", "id":
+        "417ba636-2dce-11eb-84c0-6bce5b0d9de0"}, {"type": "permissions", "id": "43fa188e-2dce-11eb-84c0-835ad1fd6287"},
+        {"type": "permissions", "id": "4916eebe-2dce-11eb-84c0-271cb2c672e8"}, {"type":
+        "permissions", "id": "edfd5e75-801f-11eb-96d8-da7ad0900002"}, {"type": "permissions",
+        "id": "98b984f4-b16d-11eb-a2c6-da7ad0900002"}, {"type": "permissions", "id":
+        "12efc20e-d36c-11eb-a9b8-da7ad0900002"}, {"type": "permissions", "id": "97971c1c-e895-11eb-b13c-da7ad0900002"},
+        {"type": "permissions", "id": "7605ef24-f376-11eb-b90b-da7ad0900002"}, {"type":
+        "permissions", "id": "7605ef25-f376-11eb-b90b-da7ad0900002"}, {"type": "permissions",
+        "id": "f4473c60-4792-11ec-a27b-da7ad0900002"}, {"type": "permissions", "id":
+        "8e4d6b6e-5750-11ec-a9f4-da7ad0900002"}, {"type": "permissions", "id": "945b3bb4-5884-11ec-aa6d-da7ad0900002"},
+        {"type": "permissions", "id": "f6e917a8-8502-11ec-bf20-da7ad0900002"}, {"type":
+        "permissions", "id": "f6e917a6-8502-11ec-bf20-da7ad0900002"}, {"type": "permissions",
+        "id": "b6bf9ac6-9a59-11ec-8480-da7ad0900002"}, {"type": "permissions", "id":
+        "f8e941cf-e746-11ec-b22d-da7ad0900002"}, {"type": "permissions", "id": "ee68fba8-173a-11ed-b00b-da7ad0900002"},
+        {"type": "permissions", "id": "6be119a6-1cd8-11ed-b185-da7ad0900002"}, {"type":
+        "permissions", "id": "4ee674f6-55d9-11ed-b10d-da7ad0900002"}, {"type": "permissions",
+        "id": "4ee5731c-55d9-11ed-b10b-da7ad0900002"}, {"type": "permissions", "id":
+        "8247acc4-7a4c-11ed-958f-da7ad0900002"}, {"type": "permissions", "id": "6c5ad874-7aff-11ed-a5cd-da7ad0900002"},
+        {"type": "permissions", "id": "c13a2368-7d61-11ed-b5b7-da7ad0900002"}, {"type":
+        "permissions", "id": "4e61a95e-de98-11ed-aa23-da7ad0900002"}, {"type": "permissions",
+        "id": "a773e3d8-fff2-11ed-965c-da7ad0900002"}, {"type": "permissions", "id":
+        "1377d9e4-0ec7-11ee-aebc-da7ad0900002"}, {"type": "permissions", "id": "de0e73c2-3d23-11ee-aa7d-da7ad0900002"},
+        {"type": "permissions", "id": "a8b4d6e8-4ea4-11ee-b482-da7ad0900002"}, {"type":
+        "permissions", "id": "50c270de-69ee-11ee-9151-da7ad0900002"}, {"type": "permissions",
+        "id": "bdda0cea-c1f0-11ee-b427-da7ad0900002"}, {"type": "permissions", "id":
+        "f5f475d4-0197-11ef-be1f-da7ad0900002"}, {"type": "permissions", "id": "8c3a9cde-0973-11ef-a2be-da7ad0900002"}]}}}],
+        "meta": {"page": {"total_count": 3, "total_filtered_count": 3}}}'
+    headers:
+      Content-Type:
+      - application/json
+    status:
+      code: 200
+      message: OK
+- request:
+    body: null
+    headers:
+      Content-Type:
+      - application/json
+    method: GET
+    uri: https://api.datadoghq.com/api/v2/authn_mappings?page%5Bnumber%5D=0&page%5Bsize%5D=100&resource_type=team
+  response:
+    body:
+      string: '{"data": [{"type": "authn_mappings", "id": "5f23356e-29bf-11ef-b407-da7ad0900002",
+        "attributes": {"attribute_key": "Member-of", "attribute_value": "TestOrgFoo",
+        "created_at": "2024-06-13T19:58:58.340586+00:00", "modified_at": "2024-06-13T19:58:58.340586+00:00"},
+        "relationships": {"team": {"data": {"type": "team", "id": "65f966e4-4189-4b52-bff1-ba5f7e572b73"}}}},
+        {"type": "authn_mappings", "id": "5f2351fc-29bf-11ef-929a-da7ad0900002", "attributes":
+        {"attribute_key": "Member-of", "attribute_value": "TestOrg", "created_at":
+        "2024-06-13T19:58:58.341089+00:00", "modified_at": "2024-06-13T19:58:58.341089+00:00"},
+        "relationships": {"team": {"data": {"type": "team", "id": "65f966e4-4189-4b52-bff1-ba5f7e572b73"}}}}],
+        "included": [{"type": "team", "id": "65f966e4-4189-4b52-bff1-ba5f7e572b73",
+        "attributes": {"name": "Example team", "handle": "example-team", "summary":
+        "This is the description of a team", "avatar": null, "banner": 0, "user_count":
+        1, "link_count": 0, "is_open_membership": true, "handles": ["example-team"]}}],
+        "meta": {"page": {"total_count": 2, "total_filtered_count": 2}}}'
+    headers:
+      Content-Type:
+      - application/json
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/...ources/cassettes/test_authn_mappings/TestAuthNMappingsResources.test_resource_sync.frozen b/...ources/cassettes/test_authn_mappings/TestAuthNMappingsResources.test_resource_sync.frozen
@@ -0,0 +1 @@
+2024-06-14T13:40:02.236270-04:00