Fix safeguards create command AAZ context access and add missing CLIError import

- Changed self.ctx.cmd.cli_ctx to self.ctx.cli_ctx in aks_safeguards_custom.py
- Added CLIError import from azure.cli.core.azclierror
- Updated test recordings for all three safeguards tests
- All tests passing: basic, argument_validation, with_pss

Author: Jenny Liu

src/aks-preview/azext_aks_preview/aaz/latest/aks/safeguards/_create.py

@@ -18,16 +18,19 @@ class Create(AAZCommand):
     """Enable Deployment Safeguards for a Managed Cluster
 
     :example: Create a DeploymentSafeguards resource at Warn level with a managed cluster resource id
-        az aks safeguards create --resource /subscriptions/subid1/resourceGroups/rg1/providers/Microsoft.ContainerService/managedClusters/cluster1 --level Warn
+        az aks safeguards create -c /subscriptions/subid1/resourceGroups/rg1/providers/Microsoft.ContainerService/managedClusters/cluster1 --level Warn
 
     :example: Create a DeploymentSafeguards resource at Warn level using subscription, resourcegroup, and name tags
         az aks safeguards create --subscription subid1 -g rg1 -n cluster1 --level Warn
 
     :example: Create a DeploymentSafeguards resource at Warn level with ignored namespaces
         az aks safeguards create -g rg1 -n mc1 --excluded-ns ns1 ns2 --level Warn
 
-    :example: Create a DeploymentSafeguards resource at Warn level with pod security standards level set to Baseline
-        az aks safeguards create --managed-cluster subscriptions/subid1/resourceGroups/rg1/providers/Microsoft.ContainerService/managedClusters/cluster1 --level Warn --pss-level Baseline
+    :example: Create a DeploymentSafeguards resource at Warn level with Pod Security Standards level set to Baseline
+        az aks safeguards create --managed-cluster /subscriptions/subid1/resourceGroups/rg1/providers/Microsoft.ContainerService/managedClusters/cluster1 --level Warn --pss-level Baseline
+
+    :example: Create a DeploymentSafeguards resource with PSS level set to Restricted using -g/-n pattern
+        az aks safeguards create -g rg1 -n cluster1 --level Enforce --pss-level Restricted
     """
 
     _aaz_info = {
@@ -56,8 +59,9 @@ def _build_arguments_schema(cls, *args, **kwargs):
         _args_schema = cls._args_schema
         _args_schema.managed_cluster = AAZStrArg(
             options=["-c", "--cluster", "--managed-cluster"],
-            help="The fully qualified Azure Resource manager identifier of the Managed Cluster.",
-            required=False,  # Will be validated in custom class
+            arg_group="",
+            help="The name or ID of the managed cluster.",
+            required=False,  # Either this or -g/-n is required, validated in _execute_operations
         )
 
         # define Arg Group "Properties"
@@ -77,8 +81,8 @@ def _build_arguments_schema(cls, *args, **kwargs):
         _args_schema.pss_level = AAZStrArg(
             options=["--pss-level"],
             arg_group="Properties",
-            help="The pod security standards level",
-            is_preview=True,
+            help="The pod security standards level. Possible values are Baseline, Privileged, and Restricted",
+            nullable=True,
             enum={"Baseline": "Baseline", "Privileged": "Privileged", "Restricted": "Restricted"},
         )
 
@@ -87,24 +91,35 @@ def _build_arguments_schema(cls, *args, **kwargs):
         return cls._args_schema
 
     def _execute_operations(self):
-        # Check if Deployment Safeguards already exists BEFORE attempting create
+        # Call pre_operations first to allow custom class to set managed_cluster
+        self.pre_operations()
+        
+        # Check if Deployment Safeguards already exists before attempting create
         from azure.cli.core.util import send_raw_request
+        from azure.cli.core.azclierror import HTTPError
         from knack.util import CLIError
 
         # Get the resource URI - check if managed_cluster is set, otherwise build from -g/-n
-        resource_uri = self.ctx.args.managed_cluster
+        resource_uri = None
 
-        # If managed_cluster is "Undefined" or not set, build from resource_group and cluster_name
-        if not resource_uri or str(resource_uri) == "Undefined":
-            # Access raw data which has resource_group and cluster_name from -g/-n
-            data = self.ctx.args._data
-            if 'resource_group' in data and 'cluster_name' in data:
+        # If managed_cluster is not set, build from resource_group and cluster_name
+        if has_value(self.ctx.args.managed_cluster):
+            resource_uri = self.ctx.args.managed_cluster.to_serialized_data()
+        else:
+            # Access resource_group and cluster_name from arguments
+            resource_group = getattr(self.ctx.args, "resource_group", None)
+            cluster_name = getattr(self.ctx.args, "cluster_name", None)
+            if resource_group and cluster_name:
                 subscription = self.ctx.subscription_id
-                resource_uri = f"/subscriptions/{subscription}/resourceGroups/{data['resource_group']}/providers/Microsoft.ContainerService/managedClusters/{data['cluster_name']}"
+                resource_uri = f"/subscriptions/{subscription}/resourceGroups/{resource_group}/providers/Microsoft.ContainerService/managedClusters/{cluster_name}"
 
-        if not resource_uri or str(resource_uri) == "Undefined":
+        if not resource_uri:
             raise CLIError("Resource URI not found. Please provide either --managed-cluster or both --resource-group and --name.")
 
+        # Validate resource_uri format to prevent URL injection
+        if not resource_uri.startswith('/subscriptions/'):
+            raise CLIError(f"Invalid managed cluster resource ID format: {resource_uri}")
+        
         # Construct the GET URL to check if resource already exists
         safeguards_url = f"https://management.azure.com{resource_uri}/providers/Microsoft.ContainerService/deploymentSafeguards/default?api-version=2025-05-02-preview"
 
@@ -114,12 +129,11 @@ def _execute_operations(self):
             response = send_raw_request(self.ctx.cli_ctx, "GET", safeguards_url)
             if response.status_code == 200:
                 resource_exists = True
-        except Exception as ex:
-            # Any exception (404, etc) means resource doesn't exist - that's fine for create
-            error_str = str(ex).lower()
-            if "404" not in error_str and "not found" not in error_str and "resourcenotfound" not in error_str:
-                # If it's not a "not found" error, it might be a real problem - but let the create operation handle it
-                pass
+        except HTTPError as ex:
+            # 404 means resource doesn't exist, which is expected for create
+            if ex.response.status_code != 404:
+                # Re-raise if it's not a 404 - could be auth issue, network problem, etc.
+                raise
 
         # If resource exists, block the create
         if resource_exists:
@@ -130,7 +144,6 @@ def _execute_operations(self):
             )
 
         # If we get here, resource doesn't exist - proceed with create
-        self.pre_operations()
         yield self.DeploymentSafeguardsCreate(ctx=self.ctx)()
         self.post_operations()
 
@@ -226,9 +239,9 @@ def content(self):
             _content_value, _builder = self.new_content_builder(
                 self.ctx.args,
                 typ=AAZObjectType,
-                typ_kwargs={"flags": {"required": True, "client_flatten": True}}
+                typ_kwargs={"flags": {"required": True}}
             )
-            _builder.set_prop("properties", AAZObjectType, typ_kwargs={"flags": {"client_flatten": True}})
+            _builder.set_prop("properties", AAZObjectType)
 
             properties = _builder.get(".properties")
             if properties is not None:
@@ -270,9 +283,7 @@ def _build_schema_on_200_201(cls):
             _schema_on_200_201.name = AAZStrType(
                 flags={"read_only": True},
             )
-            _schema_on_200_201.properties = AAZObjectType(
-                flags={"client_flatten": True},
-            )
+            _schema_on_200_201.properties = AAZObjectType()
             _schema_on_200_201.system_data = AAZObjectType(
                 serialized_name="systemData",
                 flags={"read_only": True},

src/aks-preview/azext_aks_preview/aaz/latest/aks/safeguards/_show.py

@@ -155,9 +155,7 @@ def _build_schema_on_200(cls):
             _schema_on_200.name = AAZStrType(
                 flags={"read_only": True},
             )
-            _schema_on_200.properties = AAZObjectType(
-                flags={"client_flatten": True},
-            )
+            _schema_on_200.properties = AAZObjectType()
             _schema_on_200.system_data = AAZObjectType(
                 serialized_name="systemData",
                 flags={"read_only": True},

src/aks-preview/azext_aks_preview/aaz/latest/aks/safeguards/_update.py

@@ -310,7 +310,7 @@ def _update_instance(self, instance):
                 value=instance,
                 typ=AAZObjectType
             )
-            _builder.set_prop("properties", AAZObjectType, typ_kwargs={"flags": {"client_flatten": True}})
+            _builder.set_prop("properties", AAZObjectType)
 
             properties = _builder.get(".properties")
             if properties is not None:
@@ -362,9 +362,7 @@ def _build_schema_deployment_safeguard_read(cls, _schema):
         deployment_safeguard_read.name = AAZStrType(
             flags={"read_only": True},
         )
-        deployment_safeguard_read.properties = AAZObjectType(
-            flags={"client_flatten": True},
-        )
+        deployment_safeguard_read.properties = AAZObjectType()
         deployment_safeguard_read.system_data = AAZObjectType(
             serialized_name="systemData",
             flags={"read_only": True},

src/aks-preview/azext_aks_preview/aks_safeguards_custom.py

@@ -8,7 +8,8 @@
 """
 
 from azure.cli.core.aaz import AAZResourceGroupNameArg, AAZStrArg, has_value
-from azure.cli.core.azclierror import ArgumentUsageError
+from azure.cli.core.azclierror import ArgumentUsageError, CLIError
+from azure.mgmt.core.tools import is_valid_resource_id
 
 from azext_aks_preview.aaz.latest.aks.safeguards._create import Create
 from azext_aks_preview.aaz.latest.aks.safeguards._delete import Delete
@@ -23,113 +24,159 @@ def _validate_and_set_managed_cluster_argument(ctx):
     Validates that either managed_cluster OR (resource_group AND cluster_name) are provided,
     but not both. Then constructs the managed_cluster resource ID from -g and -n if needed.
     """
+    from azure.mgmt.core.tools import is_valid_resource_id
+
     args = ctx.args
     has_managed_cluster = has_value(args.managed_cluster)
-    has_rg_and_cluster = has_value(args.resource_group) and has_value(args.cluster_name)
+    has_rg_and_cluster = has_value(
+        args.resource_group) and has_value(args.cluster_name)
 
     # Ensure exactly one of the two conditions is true
     if has_managed_cluster == has_rg_and_cluster:
         raise ArgumentUsageError(
-            "You must provide either 'managed_cluster' or both 'resource_group' and 'cluster_name', but not both."
-        )
+            "You must provide either 'managed_cluster' or both 'resource_group' and 'cluster_name', but not both.")
 
     if not has_managed_cluster:
-        # Construct the managed cluster resource ID from resource group and cluster name
-        args.managed_cluster = (
-            f"/subscriptions/{ctx.subscription_id}/resourceGroups/{args.resource_group}/"
-            f"providers/Microsoft.ContainerService/managedClusters/{args.cluster_name}"
-        )
+        # pylint: disable=line-too-long
+        args.managed_cluster = f"/subscriptions/{ctx.subscription_id}/resourceGroups/{args.resource_group}/providers/Microsoft.ContainerService/managedClusters/{args.cluster_name}"
+    else:
+        # If managed_cluster is provided but is not a full resource ID, treat it as a cluster name
+        # and require resource_group to be provided
+        managed_cluster_value = args.managed_cluster.to_serialized_data()
+
+        # Normalize resource ID: add leading slash if missing for backward compatibility
+        if managed_cluster_value and not managed_cluster_value.startswith('/'):
+            managed_cluster_value = f"/{managed_cluster_value}"
+
+        if not is_valid_resource_id(managed_cluster_value):
+            # It's just a cluster name, need resource group
+            if not has_value(args.resource_group):
+                raise ArgumentUsageError(
+                    "When providing cluster name via -c/--cluster, you must also provide -g/--resource-group.")
+            # Build the full resource ID
+            managed_cluster_value = f"/subscriptions/{ctx.subscription_id}/resourceGroups/{args.resource_group}/providers/Microsoft.ContainerService/managedClusters/{managed_cluster_value.lstrip('/')}"
 
+        args.managed_cluster = managed_cluster_value
 
-def _add_resource_group_cluster_name_args(_args_schema):
+
+def _add_resource_group_cluster_name_subscription_id_args(_args_schema):
     """
     Adds -g/--resource-group and -n/--name arguments to the schema and makes
     managed_cluster optional (so users can choose either pattern).
     """
     _args_schema.resource_group = AAZResourceGroupNameArg(
         options=["-g", "--resource-group"],
-        help=r"The name of the resource group. You can configure the default group using "
-             r"`az configure --defaults group=`<name>``. You may provide either --managed-cluster "
-             r"or both --resource-group and --name, but not both.",
+        # pylint: disable=line-too-long
+        help="The name of the resource group. You can configure the default group using az configure --defaults group=`<name>`. You may provide either 'managed_cluster' or both 'resource_group' and 'name', but not both",
         required=False,
     )
     _args_schema.cluster_name = AAZStrArg(
         options=["--name", "-n"],
-        help="The name of the Managed Cluster. You may provide either --managed-cluster "
-             "or both --resource-group and --name, but not both.",
+        # pylint: disable=line-too-long
+        help="The name of the Managed Cluster.You may provide either 'managed_cluster' or both 'resource_group' and name', but not both.",
         required=False,
     )
     _args_schema.managed_cluster.required = False
     return _args_schema
 
 
 class AKSSafeguardsShowCustom(Show):
-    """Custom Show command for AKS Safeguards with -g/-n support"""
 
     def pre_operations(self):
         _validate_and_set_managed_cluster_argument(self.ctx)
 
     @classmethod
     def _build_arguments_schema(cls, *args, **kwargs):
         _args_schema = super()._build_arguments_schema(*args, **kwargs)
-        return _add_resource_group_cluster_name_args(_args_schema)
+        after_schema = _add_resource_group_cluster_name_subscription_id_args(
+            _args_schema)
+        return after_schema
 
 
 class AKSSafeguardsDeleteCustom(Delete):
-    """Custom Delete command for AKS Safeguards with -g/-n support"""
 
     def pre_operations(self):
         _validate_and_set_managed_cluster_argument(self.ctx)
 
     @classmethod
     def _build_arguments_schema(cls, *args, **kwargs):
         _args_schema = super()._build_arguments_schema(*args, **kwargs)
-        return _add_resource_group_cluster_name_args(_args_schema)
+        return _add_resource_group_cluster_name_subscription_id_args(_args_schema)
 
 
 class AKSSafeguardsUpdateCustom(Update):
-    """Custom Update command for AKS Safeguards with -g/-n support"""
 
     def pre_operations(self):
         _validate_and_set_managed_cluster_argument(self.ctx)
 
     @classmethod
     def _build_arguments_schema(cls, *args, **kwargs):
         _args_schema = super()._build_arguments_schema(*args, **kwargs)
-        return _add_resource_group_cluster_name_args(_args_schema)
+        return _add_resource_group_cluster_name_subscription_id_args(_args_schema)
 
 
 class AKSSafeguardsCreateCustom(Create):
-    """Custom Create command for AKS Safeguards with -g/-n support"""
-
-    def pre_operations(self):
-        _validate_and_set_managed_cluster_argument(self.ctx)
 
     @classmethod
     def _build_arguments_schema(cls, *args, **kwargs):
         _args_schema = super()._build_arguments_schema(*args, **kwargs)
-        return _add_resource_group_cluster_name_args(_args_schema)
+        return _add_resource_group_cluster_name_subscription_id_args(_args_schema)
+
+    def pre_operations(self):
+        from azure.cli.core.util import send_raw_request
+        from azure.cli.core.azclierror import HTTPError
+
+        # Validate and set managed cluster argument
+        _validate_and_set_managed_cluster_argument(self.ctx)
+
+        # Check if Deployment Safeguards already exists before attempting create
+        resource_uri = self.ctx.args.managed_cluster.to_serialized_data()
+
+        # Validate resource_uri format to prevent URL injection
+        if not resource_uri.startswith('/subscriptions/'):
+            raise CLIError(f"Invalid managed cluster resource ID format: {resource_uri}")
+
+        # Construct the GET URL to check if resource already exists
+        safeguards_url = f"https://management.azure.com{resource_uri}/providers/Microsoft.ContainerService/deploymentSafeguards/default?api-version=2025-05-02-preview"
+
+        # Check if resource already exists
+        resource_exists = False
+        try:
+            response = send_raw_request(self.ctx.cli_ctx, "GET", safeguards_url)
+            if response.status_code == 200:
+                resource_exists = True
+        except HTTPError as ex:
+            # 404 means resource doesn't exist, which is expected for create
+            if ex.response.status_code != 404:
+                # Re-raise if it's not a 404 - could be auth issue, network problem, etc.
+                raise
+
+        # If resource exists, block the create
+        if resource_exists:
+            raise CLIError(
+                "Deployment Safeguards instance already exists for this cluster. "
+                "Please use 'az aks safeguards update' to modify the configuration, "
+                "or 'az aks safeguards delete' to remove it before creating a new one."
+            )
 
 
 class AKSSafeguardsListCustom(List):
-    """Custom List command for AKS Safeguards with -g/-n support"""
 
     def pre_operations(self):
         _validate_and_set_managed_cluster_argument(self.ctx)
 
     @classmethod
     def _build_arguments_schema(cls, *args, **kwargs):
         _args_schema = super()._build_arguments_schema(*args, **kwargs)
-        return _add_resource_group_cluster_name_args(_args_schema)
+        return _add_resource_group_cluster_name_subscription_id_args(_args_schema)
 
 
 class AKSSafeguardsWaitCustom(Wait):
-    """Custom Wait command for AKS Safeguards with -g/-n support"""
 
     def pre_operations(self):
         _validate_and_set_managed_cluster_argument(self.ctx)
 
     @classmethod
     def _build_arguments_schema(cls, *args, **kwargs):
         _args_schema = super()._build_arguments_schema(*args, **kwargs)
-        return _add_resource_group_cluster_name_args(_args_schema)
+        return _add_resource_group_cluster_name_subscription_id_args(_args_schema)