Merge pull request #153 from Azure/guwe/merge-vm-tools

Feat: support more vm/vmss operations

Author: feiskyer

internal/components/compute/azcommands.go

@@ -1,57 +1,203 @@
 package compute
 
 import (
-	"github.com/Azure/aks-mcp/internal/utils"
+	"fmt"
+	"slices"
+
+	"github.com/Azure/aks-mcp/internal/config"
 	"github.com/mark3labs/mcp-go/mcp"
 )
 
-// ComputeCommand defines a specific az vmss command to be registered as a tool
-type ComputeCommand struct {
-	Name        string
-	Description string
-	ArgsExample string // Example of command arguments
-}
+// ComputeOperationType defines the type of compute operation
+type ComputeOperationType string
+
+// ResourceType defines the compute resource type
+type ResourceType string
+
+const (
+	// VM operations - safe operations only
+	OpVMShow            ComputeOperationType = "show"
+	OpVMList            ComputeOperationType = "list"
+	OpVMStart           ComputeOperationType = "start"
+	OpVMStop            ComputeOperationType = "stop"
+	OpVMRestart         ComputeOperationType = "restart"
+	OpVMGetInstanceView ComputeOperationType = "get-instance-view"
+	OpVMRunCommand      ComputeOperationType = "run-command"
+
+	// VMSS operations - only safe operations for AKS-managed VMSS
+	OpVMSSShow            ComputeOperationType = "show"
+	OpVMSSList            ComputeOperationType = "list"
+	OpVMSSRestart         ComputeOperationType = "restart"
+	OpVMSSReimage         ComputeOperationType = "reimage"
+	OpVMSSGetInstanceView ComputeOperationType = "get-instance-view"
+	OpVMSSRunCommand      ComputeOperationType = "run-command"
+
+	// Resource types
+	ResourceTypeVM   ResourceType = "vm"
+	ResourceTypeVMSS ResourceType = "vmss"
+)
+
+// generateToolDescription creates a tool description based on access level
+func generateToolDescription(accessLevel string) string {
+	baseDesc := `Unified tool for managing Azure Virtual Machines (VMs) and Virtual Machine Scale Sets (VMSS) using Azure CLI.
+
+IMPORTANT: VM/VMSS resources are managed by AKS. Write operations should be used carefully and only for debugging purposes.
+
+Use resource_type="vm" for single virtual machines or resource_type="vmss" for virtual machine scale sets.
+
+Available operation values:`
 
-// RegisterAzComputeCommand registers a specific az vmss command as an MCP tool
-func RegisterAzComputeCommand(cmd ComputeCommand) mcp.Tool {
-	// Convert spaces to underscores for valid tool name
-	commandName := cmd.Name
-	validToolName := utils.ReplaceSpacesWithUnderscores(commandName)
+	// Add operations by access level
+	desc := baseDesc + "\n"
 
-	description := "Run " + cmd.Name + " command: " + cmd.Description + "."
+	// Basic operations for all access levels
+	desc += "- show: Get details of a VM/VMSS\n"
+	desc += "- list: List VMs/VMSS in subscription or resource group\n"
+	desc += "- get-instance-view: Get runtime status\n"
 
-	// Add example if available, with proper punctuation
-	if cmd.ArgsExample != "" {
-		description += "\nExample: `" + cmd.ArgsExample + "`"
+	// Management operations for readwrite/admin
+	if accessLevel == "readwrite" || accessLevel == "admin" {
+		desc += "- start: Start VM\n"
+		desc += "- stop: Stop VM\n"
+		desc += "- restart: Restart VM/VMSS instances\n"
+		desc += "- run-command: Execute commands remotely on VM/VMSS instances\n"
+		desc += "- reimage: Reimage VMSS instances (VM not supported for reimage)\n"
 	}
 
-	return mcp.NewTool(validToolName,
+	// Note: All destructive operations (create, delete, deallocate, update, resize, scale)
+	// have been removed for AKS environment safety
+
+	// Examples
+	desc += "\nEXAMPLES:\n"
+	desc += `List VMSS: operation="list", resource_type="vmss", args="--resource-group myRG"` + "\n"
+	desc += `Show VMSS: operation="show", resource_type="vmss", args="--name myVMSS --resource-group myRG"` + "\n"
+	desc += `List VMs: operation="list", resource_type="vm", args="--resource-group myRG"` + "\n"
+
+	if accessLevel == "readwrite" || accessLevel == "admin" {
+		desc += `Restart VMSS: operation="restart", resource_type="vmss", args="--name myVMSS --resource-group myRG"` + "\n"
+		desc += `Reimage VMSS: operation="reimage", resource_type="vmss", args="--name myVMSS --resource-group myRG"` + "\n"
+		desc += `Run command on VM: operation="run-command", resource_type="vm", args="--name myVM --resource-group myRG --command-id RunShellScript --scripts 'echo hello'"` + "\n"
+		desc += `Run command on VMSS: operation="run-command", resource_type="vmss", args="--name myVMSS --resource-group myRG --command-id RunShellScript --scripts 'hostname' --instance-id 0"` + "\n"
+	}
+
+	return desc
+}
+
+// RegisterAzComputeOperations registers the unified compute operations tool
+func RegisterAzComputeOperations(cfg *config.ConfigData) mcp.Tool {
+	description := generateToolDescription(cfg.AccessLevel)
+
+	return mcp.NewTool("az_compute_operations",
 		mcp.WithDescription(description),
+		mcp.WithString("operation",
+			mcp.Required(),
+			mcp.Description("Operation to perform. Common operations: list, show, start, stop, restart, deallocate, run-command, scale, etc."),
+		),
+		mcp.WithString("resource_type",
+			mcp.Required(),
+			mcp.Description("Resource type: 'vm' (single virtual machine) or 'vmss' (virtual machine scale set)"),
+		),
 		mcp.WithString("args",
 			mcp.Required(),
-			mcp.Description("Arguments for the `"+cmd.Name+"` command"),
+			mcp.Description("Azure CLI arguments: '--resource-group myRG' (required for most operations), '--name myVM' (for specific resources), '--new-capacity 3' (for scaling)"),
 		),
 	)
 }
 
-// GetReadOnlyVmssCommands returns all read-only az vmss commands
-func GetReadOnlyVmssCommands() []ComputeCommand {
-	return []ComputeCommand{
-		// No read-only commands for now
+// GetOperationAccessLevel returns the required access level for an operation
+func GetOperationAccessLevel(operation string) string {
+	readOnlyOps := []string{
+		string(OpVMShow), string(OpVMList), string(OpVMGetInstanceView),
+		string(OpVMSSShow), string(OpVMSSList), string(OpVMSSGetInstanceView),
+	}
+
+	readWriteOps := []string{
+		// VM operations - safe operations only
+		string(OpVMStart), string(OpVMStop), string(OpVMRestart), string(OpVMRunCommand),
+		// VMSS operations - only safe operations for AKS-managed VMSS
+		string(OpVMSSRestart), string(OpVMSSReimage), string(OpVMSSRunCommand),
+	}
+
+	// No admin operations - all unsafe operations removed
+	adminOps := []string{}
+
+	if slices.Contains(readOnlyOps, operation) {
+		return "readonly"
+	}
+
+	if slices.Contains(readWriteOps, operation) {
+		return "readwrite"
+	}
+
+	if slices.Contains(adminOps, operation) {
+		return "admin"
 	}
+
+	return "unknown"
 }
 
-// GetReadWriteVmssCommands returns all read-write az vmss commands
-func GetReadWriteVmssCommands() []ComputeCommand {
-	return []ComputeCommand{
-		// Run command execution
-		{Name: "az vmss run-command invoke", Description: "Execute a command on instances of a Virtual Machine Scale Set", ArgsExample: "--name myVMSS --resource-group myResourceGroup --command-id RunShellScript --scripts 'echo Hello World' --instance-id 0"},
+// ValidateOperationAccess checks if the operation is allowed for the given access level
+func ValidateOperationAccess(operation string, cfg *config.ConfigData) error {
+	requiredLevel := GetOperationAccessLevel(operation)
+
+	switch requiredLevel {
+	case "admin":
+		if cfg.AccessLevel != "admin" {
+			return fmt.Errorf("operation '%s' requires admin access level", operation)
+		}
+	case "readwrite":
+		if cfg.AccessLevel != "readwrite" && cfg.AccessLevel != "admin" {
+			return fmt.Errorf("operation '%s' requires readwrite or admin access level", operation)
+		}
+	case "readonly":
+		// All access levels can perform readonly operations
+	case "unknown":
+		return fmt.Errorf("unknown operation: %s", operation)
 	}
+
+	return nil
 }
 
-// GetAdminVmssCommands returns all admin az vmss commands
-func GetAdminVmssCommands() []ComputeCommand {
-	return []ComputeCommand{
-		// No admin commands for now
+// MapOperationToCommand maps an operation and resource type to its corresponding az command
+func MapOperationToCommand(operation string, resourceType string) (string, error) {
+	// Validate resource type
+	if resourceType != string(ResourceTypeVM) && resourceType != string(ResourceTypeVMSS) {
+		return "", fmt.Errorf("invalid resource type: %s (must be 'vm' or 'vmss')", resourceType)
+	}
+
+	commandMap := map[string]map[string]string{
+		string(ResourceTypeVM): {
+			// Safe VM operations only
+			string(OpVMShow):            "az vm show",
+			string(OpVMList):            "az vm list",
+			string(OpVMStart):           "az vm start",
+			string(OpVMStop):            "az vm stop",
+			string(OpVMRestart):         "az vm restart",
+			string(OpVMGetInstanceView): "az vm get-instance-view",
+			string(OpVMRunCommand):      "az vm run-command invoke",
+		},
+		string(ResourceTypeVMSS): {
+			// Read-only operations
+			string(OpVMSSShow):            "az vmss show",
+			string(OpVMSSList):            "az vmss list",
+			string(OpVMSSGetInstanceView): "az vmss get-instance-view",
+			// Safe operations for AKS-managed VMSS
+			string(OpVMSSRestart):    "az vmss restart",
+			string(OpVMSSReimage):    "az vmss reimage",
+			string(OpVMSSRunCommand): "az vmss run-command invoke",
+			// Removed unsafe operations: create, delete, start, stop, deallocate, scale, update
+		},
 	}
+
+	resourceCommands, exists := commandMap[resourceType]
+	if !exists {
+		return "", fmt.Errorf("unsupported resource type: %s", resourceType)
+	}
+
+	cmd, exists := resourceCommands[operation]
+	if !exists {
+		return "", fmt.Errorf("unsupported operation '%s' for resource type '%s'", operation, resourceType)
+	}
+
+	return cmd, nil
 }

internal/components/compute/executor.go

@@ -0,0 +1,137 @@
+package compute
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/Azure/aks-mcp/internal/command"
+	"github.com/Azure/aks-mcp/internal/config"
+	"github.com/Azure/aks-mcp/internal/security"
+)
+
+// ComputeOperationsExecutor handles execution of compute operations
+type ComputeOperationsExecutor struct{}
+
+// NewComputeOperationsExecutor creates a new ComputeOperationsExecutor
+func NewComputeOperationsExecutor() *ComputeOperationsExecutor {
+	return &ComputeOperationsExecutor{}
+}
+
+// Execute handles the compute operations
+func (e *ComputeOperationsExecutor) Execute(params map[string]interface{}, cfg *config.ConfigData) (string, error) {
+	// Parse operation parameter
+	operation, ok := params["operation"].(string)
+	if !ok {
+		return "", fmt.Errorf("missing or invalid 'operation' parameter. Common operations: list, show, start, stop, restart, run-command, reimage. Example: operation=\"list\"")
+	}
+
+	// Parse resource_type parameter
+	resourceType, ok := params["resource_type"].(string)
+	if !ok {
+		return "", fmt.Errorf("missing or invalid 'resource_type' parameter. Must be 'vm' (Virtual Machine) or 'vmss' (Virtual Machine Scale Set). Example: resource_type=\"vm\"")
+	}
+
+	// Parse args parameter
+	args, ok := params["args"].(string)
+	if !ok {
+		args = ""
+	}
+
+	// Validate access for this operation
+	if err := ValidateOperationAccess(operation, cfg); err != nil {
+		// Enhance access error with suggestions
+		requiredLevel := GetOperationAccessLevel(operation)
+		return "", fmt.Errorf("%v. Your current access level is '%s', but this operation requires '%s' access. Contact your administrator to request higher access", err, cfg.AccessLevel, requiredLevel)
+	}
+
+	// Map operation to Azure CLI command
+	baseCommand, err := MapOperationToCommand(operation, resourceType)
+	if err != nil {
+		// Provide helpful suggestions for invalid operations
+		validOps := getSuggestedOperations(resourceType, cfg.AccessLevel)
+		return "", fmt.Errorf("%v. Valid operations for %s with %s access: %s", err, resourceType, cfg.AccessLevel, validOps)
+	}
+
+	// Build full command
+	fullCommand := baseCommand
+	if args != "" {
+		fullCommand += " " + args
+	}
+
+	// Validate the command against security settings
+	validator := security.NewValidator(cfg.SecurityConfig)
+	err = validator.ValidateCommand(fullCommand, security.CommandTypeAz)
+	if err != nil {
+		return "", err
+	}
+
+	// Extract binary name and arguments from command
+	cmdParts := strings.Fields(fullCommand)
+	if len(cmdParts) == 0 {
+		return "", fmt.Errorf("empty command")
+	}
+
+	// Use the first part as the binary name
+	binaryName := cmdParts[0]
+
+	// The rest of the command becomes the arguments
+	cmdArgs := ""
+	if len(cmdParts) > 1 {
+		cmdArgs = strings.Join(cmdParts[1:], " ")
+	}
+
+	// If the command is not an az command, return an error
+	if binaryName != "az" {
+		return "", fmt.Errorf("command must start with 'az'")
+	}
+
+	// Execute the command
+	process := command.NewShellProcess(binaryName, cfg.Timeout)
+	result, err := process.Run(cmdArgs)
+	if err != nil {
+		// Provide helpful error messages for common issues
+		errorMsg := fmt.Sprintf("Azure CLI command failed: %v", err)
+
+		// Add contextual help based on the operation
+		switch operation {
+		case "list":
+			errorMsg += "\nTip: For listing resources, try without --name parameter, or verify --resource-group exists"
+		case "show":
+			errorMsg += "\nTip: Verify the resource name and resource group are correct and the resource exists"
+		case "start", "stop", "restart":
+			errorMsg += "\nTip: Verify the resource exists and check if it's already in the desired state"
+		case "reimage":
+			errorMsg += "\nTip: Verify the VMSS name is correct and the instances are ready for reimaging"
+		case "run-command":
+			errorMsg += "\nTip: Ensure the resource is running and the command syntax is correct. Use --command-id RunShellScript for shell commands"
+		}
+
+		return "", fmt.Errorf("%s\nExecuted command: %s", errorMsg, fullCommand)
+	}
+
+	return strings.TrimSpace(result), nil
+}
+
+// getSuggestedOperations returns a helpful list of valid operations for the given resource type and access level
+func getSuggestedOperations(resourceType, accessLevel string) string {
+	var operations []string
+
+	// Read-only operations (available to all access levels)
+	operations = append(operations, "list", "show", "get-instance-view")
+
+	// Read-write operations
+	if accessLevel == "readwrite" || accessLevel == "admin" {
+		switch resourceType {
+		case "vm":
+			// Only safe VM operations
+			operations = append(operations, "start", "stop", "restart", "run-command")
+		case "vmss":
+			// Only safe operations for AKS-managed VMSS
+			operations = append(operations, "restart", "reimage", "run-command")
+		}
+	}
+
+	// No admin operations - all unsafe operations removed for AKS safety
+
+	return strings.Join(operations, ", ")
+}

internal/components/compute/registry.go

@@ -28,25 +28,3 @@ func RegisterAKSVMSSInfoTool() mcp.Tool {
 		),
 	)
 }
-
-// VMSS Command Registration Functions
-
-// RegisterAzVmssCommand registers a specific az vmss command as an MCP tool
-func RegisterAzVmssCommand(cmd ComputeCommand) mcp.Tool {
-	return RegisterAzComputeCommand(cmd)
-}
-
-// GetReadOnlyVmssAzCommands returns all read-only az vmss commands
-func GetReadOnlyVmssAzCommands() []ComputeCommand {
-	return GetReadOnlyVmssCommands()
-}
-
-// GetReadWriteVmssAzCommands returns all read-write az vmss commands
-func GetReadWriteVmssAzCommands() []ComputeCommand {
-	return GetReadWriteVmssCommands()
-}
-
-// GetAdminVmssAzCommands returns all admin az vmss commands
-func GetAdminVmssAzCommands() []ComputeCommand {
-	return GetAdminVmssCommands()
-}