custom.py

# --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
#
# Code generated by aaz-dev-tools
# --------------------------------------------------------------------------------------------

# pylint: disable=too-many-lines
# pylint: disable=too-many-statements, protected-access
from azext_stack_hci.aaz.latest.stack_hci.cluster import List as _ClusterList
from azext_stack_hci.aaz.latest.stack_hci.cluster import Create as _ClusterCreate
from azext_stack_hci.aaz.latest.stack_hci.cluster.identity import Assign as _IdentityAssign
from azext_stack_hci.aaz.latest.stack_hci.cluster.identity import Remove as _IdentityRemove


class ClusterList(_ClusterList):
    def post_operations(self):
        # next_link cannot be accessed
        if self.ctx.vars.instance.next_link.to_serialized_data():
            self.ctx.vars.instance.next_link = None


class ClusterCreate(_ClusterCreate):
    @classmethod
    def _build_arguments_schema(cls, *args, **kwargs):
        from azure.cli.core.aaz import AAZBoolArg
        args_schema = super()._build_arguments_schema(*args, **kwargs)
        args_schema.mi_system_assigned = AAZBoolArg(
            options=["--mi-system-assigned"],
            help="Enable system assigned identity"
        )
        args_schema.identity._registered = False
        return args_schema

    def pre_operations(self):
        args = self.ctx.args
        if args.mi_system_assigned:
            args.identity.type = "SystemAssigned"


class IdentityAssign(_IdentityAssign):
    @classmethod
    def _build_arguments_schema(cls, *args, **kwargs):
        from azure.cli.core.aaz import AAZBoolArg
        args_schema = super()._build_arguments_schema(*args, **kwargs)
        args_schema.system_assigned = AAZBoolArg(
            options=["--system-assigned"],
            help="Enable system assigned identity"
        )
        args_schema.type._registered = False
        args_schema.type._required = False
        args_schema.user_assigned_identities._registered = False
        return args_schema

    def pre_operations(self):
        args = self.ctx.args
        if args.system_assigned:
            args.type = "SystemAssigned"


class IdentityRemove(_IdentityRemove):
    @classmethod
    def _build_arguments_schema(cls, *args, **kwargs):
        from azure.cli.core.aaz import AAZBoolArg
        args_schema = super()._build_arguments_schema(*args, **kwargs)
        args_schema.system_assigned = AAZBoolArg(
            options=["--system-assigned"],
            help="Enable system assigned identity"
        )
        args_schema.type._registered = False
        args_schema.user_assigned_identities._registered = False
        return args_schema

    def pre_instance_update(self, instance):
        args = self.ctx.args
        if args.system_assigned:
            args.type = 'None'


class VmConnectEnable(_VmConnectEnable):
    @classmethod
    def _build_arguments_schema(cls, *args, **kwargs):
        from azure.cli.core.aaz import AAZResourceGroupNameArg, AAZStrArg
        args_schema = super()._build_arguments_schema(*args, **kwargs)
        args_schema.cluster_name = AAZStrArg(
            options=["--cluster-name", "-n"],
            help="The name of the cluster.",
            required=True
        )
        args_schema.resource_group = AAZResourceGroupNameArg(
            options=["--resource-group", "-g"],
            help="Name of resource group.",
            required=True
        )
        args_schema.vm_name = AAZStrArg(
            options=["--vm-name"],
            help="The name of the virtual machine.",
            required=True
        )
        return args_schema

    def __call__(self, cmd, **kwargs):
        from azure.cli.core.commands.client_factory import get_subscription_id
        from azure.cli.core.util import send_raw_request
        import json

        cluster_name = kwargs.get('cluster_name')
        resource_group = kwargs.get('resource_group')
        vm_name = kwargs.get('vm_name')
        subscription_id = get_subscription_id(cmd.cli_ctx)

        # Construct the REST API path
        path = (
            f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/"
            "providers/Microsoft.AzureStackHCI/clusters/"
            f"{cluster_name}/jobs/VmConnectDeprovision"
        )
        # API version
        api_version = "2023-12-01-preview"
        url = f"https://management.azure.com{path}?api-version={api_version}"
        # Default payload with VM name
        payload = {
            "properties": {
                "jobType": "VmConnectProvision",
                "deploymentMode": "Deploy",
                "vmConnectProvisionJobDetails": [
                    {
                        "vmName": vm_name
                    }
                ]
            }
        }

        # Make the REST API call
        try:
            response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))
            if response.content:
                return response.json()
            else:
                return {
                    "message": (
                        f"VM Connect provision job initiated successfully for VM: {vm_name}"
                    )
                }
        except Exception as e:
            from azure.cli.core.util import CLIError
            raise CLIError(f"Failed to enable VM Connect for VM '{vm_name}': {str(e)}")


class VmConnectDisable(_VmConnectDisable):
    @classmethod
    def _build_arguments_schema(cls, *args, **kwargs):
        from azure.cli.core.aaz import AAZResourceGroupNameArg, AAZStrArg
        args_schema = super()._build_arguments_schema(*args, **kwargs)
        args_schema.cluster_name = AAZStrArg(
            options=["--cluster-name", "-n"],
            help="The name of the cluster.",
            required=True
        )
        args_schema.resource_group = AAZResourceGroupNameArg(
            options=["--resource-group", "-g"],
            help="Name of resource group.",
            required=True
        )
        args_schema.vm_name = AAZStrArg(
            options=["--vm-name"],
            help="The name of the virtual machine.",
            required=True
        )
        return args_schema

    def __call__(self, cmd, **kwargs):
        from azure.cli.core.commands.client_factory import get_subscription_id
        from azure.cli.core.util import send_raw_request
        import json

        cluster_name = kwargs.get('cluster_name')
        resource_group = kwargs.get('resource_group')
        vm_name = kwargs.get('vm_name')

        subscription_id = get_subscription_id(cmd.cli_ctx)

        # Construct the REST API path for deprovision
        path = (
            f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/"
            "providers/Microsoft.AzureStackHCI/clusters/"
            f"{cluster_name}/jobs/VmConnectDeprovision"
        )

        # API version
        api_version = "2023-12-01-preview"
        url = f"https://management.azure.com{path}?api-version={api_version}"

        # Payload for VM Connect deprovision
        payload = {
            "properties": {
                "jobType": "VmConnectRemove",
                "deploymentMode": "Deploy",
                "vmConnectRemoveJobDetails": [
                    {
                        "vmName": vm_name
                    }
                ]
            }
        }

        # Make the REST API call
        try:
            response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))
            if response.content:
                return response.json()
            else:
                return {
                    "message": (
                        f"VM Connect provision job initiated successfully for VM: {vm_name}"
                    )
                }
        except Exception as e:
            from azure.cli.core.util import CLIError
            raise CLIError(f"Failed to disable VM Connect for VM '{vm_name}': {str(e)}")