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'


def vmconnect_enable(cmd, cluster_name, resource_group, vm_name):
    """
    Enable VM Connect for a virtual machine in an Azure Stack HCI cluster.
    """
    from azure.cli.core.commands.client_factory import get_subscription_id
    from azure.cli.core.util import send_raw_request
    import json

    subscription_id = get_subscription_id(cmd.cli_ctx)
    path = (
        f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/"
        f"providers/Microsoft.AzureStackHCI/clusters/{cluster_name}/jobs/VmConnectProvision"
    )
    api_version = "2023-12-01-preview"
    url = f"https://management.azure.com{path}?api-version={api_version}"
    payload = {
        "properties": {
            "jobType": "VmConnectProvision",
            "deploymentMode": "Deploy",
            "vmConnectProvisionJobDetails": [
                {"vmName": vm_name}
            ]
        }
    }

    try:
        response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))
        if response.content:
            return response.json()
        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)}")


def vmconnect_disable(cmd, cluster_name, resource_group, vm_name):
    """
    Disable VM Connect for a virtual machine in an Azure Stack HCI cluster.
    """
    from azure.cli.core.commands.client_factory import get_subscription_id
    from azure.cli.core.util import send_raw_request
    import json

    subscription_id = get_subscription_id(cmd.cli_ctx)
    path = (
        f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/"
        f"providers/Microsoft.AzureStackHCI/clusters/{cluster_name}/jobs/VmConnectRemove"
    )
    api_version = "2023-12-01-preview"
    url = f"https://management.azure.com{path}?api-version={api_version}"
    payload = {
        "properties": {
            "jobType": "VmConnectRemove",
            "deploymentMode": "Deploy",
            "vmConnectRemoveJobDetails": [
                {"vmName": vm_name}
            ]
        }
    }

    try:
        response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))
        if response.content:
            return response.json()
        return {"message": f"VM Connect remove 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)}")