initial changes for stack-hci vmconnect command

src/stack-hci/azext_stack_hci/commands.py

@@ -19,3 +19,8 @@ def load_command_table(self, _):  # pylint: disable=unused-argument
         self.command_table['stack-hci cluster create'] = ClusterCreate(loader=self)
         self.command_table['stack-hci cluster identity assign'] = IdentityAssign(loader=self)
         self.command_table['stack-hci cluster identity remove'] = IdentityRemove(loader=self)
+
+    with self.command_group('stack-hci vmconnect'):
+        from azext_stack_hci.custom import VmConnectEnable, VmConnectDisable
+        self.command_table['stack-hci vmconnect enable'] = VmConnectEnable(loader=self)
+        self.command_table['stack-hci vmconnect disable'] = VmConnectDisable(loader=self)

src/stack-hci/azext_stack_hci/custom.py

@@ -75,3 +75,146 @@ 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)}")