Allow lockout when creating restriction policies

datadog_sync/commands/shared/options.py

@@ -367,6 +367,15 @@ def click_config_file_provider(ctx: Context, opts: CustomOptionClass, value: Non
         "If a role has a permission that doesn't exist in the destination, it will be removed and retried.",
         cls=CustomOptionClass,
     ),
+    option(
+        "--allow-self-lockout",
+        required=False,
+        type=bool,
+        default=False,
+        show_default=True,
+        help="Allow self-lockout when syncing restriction policies.",
+        cls=CustomOptionClass,
+    ),
 ]
 
 

datadog_sync/model/restriction_policies.py

@@ -111,15 +111,31 @@ async def create_resource(self, _id: str, resource: Dict) -> Tuple[str, Dict]:
         destination_client = self.config.destination_client
         resource_id = resource["id"]
         payload = {"data": resource}
-        resp = await destination_client.post(self.resource_config.base_path + f"/{resource_id}", payload)
+
+        # Add query parameter if allow_self_lockout is enabled
+        params = {}
+        if self.config.allow_self_lockout:
+            params["allow_self_lockout"] = "true"
+
+        resp = await destination_client.post(
+            self.resource_config.base_path + f"/{resource_id}", payload, params=params if params else None
+        )
 
         return _id, resp["data"]
 
     async def update_resource(self, _id: str, resource: Dict) -> Tuple[str, Dict]:
         destination_client = self.config.destination_client
         resource_id = resource["id"]
         payload = {"data": resource}
-        resp = await destination_client.post(self.resource_config.base_path + f"/{resource_id}", payload)
+
+        # Add query parameter if allow_self_lockout is enabled
+        params = {}
+        if self.config.allow_self_lockout:
+            params["allow_self_lockout"] = "true"
+
+        resp = await destination_client.post(
+            self.resource_config.base_path + f"/{resource_id}", payload, params=params if params else None
+        )
 
         return _id, resp["data"]
 

datadog_sync/utils/configuration.py

@@ -59,6 +59,7 @@ class Configuration(object):
     verify_ddr_status: bool
     backup_before_reset: bool
     show_progress_bar: bool
+    allow_self_lockout: bool
     allow_partial_permissions_roles: List[str] = field(default_factory=list)
     resources: Dict[str, BaseResource] = field(default_factory=dict)
     resources_arg: List[str] = field(default_factory=list)
@@ -173,6 +174,7 @@ def build_config(cmd: Command, **kwargs: Optional[Any]) -> Configuration:
     verify_ddr_status = kwargs.get("verify_ddr_status")
     backup_before_reset = not kwargs.get("do_not_backup")
     show_progress_bar = kwargs.get("show_progress_bar")
+    allow_self_lockout = kwargs.get("allow_self_lockout", False)
 
     # Parse allow_partial_permissions_roles
     allow_partial_permissions_roles = []
@@ -257,6 +259,7 @@ def build_config(cmd: Command, **kwargs: Optional[Any]) -> Configuration:
         verify_ddr_status=verify_ddr_status,
         backup_before_reset=backup_before_reset,
         show_progress_bar=show_progress_bar,
+        allow_self_lockout=allow_self_lockout,
         allow_partial_permissions_roles=allow_partial_permissions_roles,
     )
 

tests/conftest.py

@@ -162,6 +162,7 @@ def config():
         send_metrics=True,
         backup_before_reset=True,
         show_progress_bar=True,
+        allow_self_lockout=False,
     )
 
     resources = init_resources(cfg)