diff --git a/datadog_sync/commands/shared/options.py b/datadog_sync/commands/shared/options.py index 55262e34..d7cf83a1 100644 --- a/datadog_sync/commands/shared/options.py +++ b/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, + ), ] diff --git a/datadog_sync/model/restriction_policies.py b/datadog_sync/model/restriction_policies.py index 1bf032ec..b870c5bf 100644 --- a/datadog_sync/model/restriction_policies.py +++ b/datadog_sync/model/restriction_policies.py @@ -111,7 +111,15 @@ 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"] @@ -119,7 +127,15 @@ 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"] diff --git a/datadog_sync/utils/configuration.py b/datadog_sync/utils/configuration.py index dfd4d0df..660e7cb2 100644 --- a/datadog_sync/utils/configuration.py +++ b/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, ) diff --git a/tests/conftest.py b/tests/conftest.py index c9410bb3..fd61481a 100644 --- a/tests/conftest.py +++ b/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)