Sanitize CLI command in telemetry to prevent secret leakage (#335)

gossion · claude · web-flow · commit 67cfcb570787 · 2026-03-11T01:33:52.000Z
Strip flags and arguments from CLI commands before passing to
TrackToolInvocation, keeping only the "az &lt;group&gt; &lt;subcommand&gt;" prefix.
This prevents secrets (--client-secret, tokens) from leaking into
OTLP/Application Insights and reduces telemetry cardinality.

Co-authored-by: Claude &lt;noreply@anthropic.com&gt;
diff --git a/internal/components/azapi/handler.go b/internal/components/azapi/handler.go
@@ -3,6 +3,7 @@ package azapi
 import (
 	"context"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/Azure/aks-mcp/internal/config"
@@ -11,6 +12,21 @@ import (
 	"github.com/mark3labs/mcp-go/mcp"
 )
 
+// sanitizeCliCommand extracts only the "az <group> <subcommand>" prefix from a
+// CLI command, stripping all flags and arguments. This prevents secrets and
+// high-cardinality values from leaking into telemetry.
+func sanitizeCliCommand(cmd string) string {
+	tokens := strings.Fields(cmd)
+	var kept []string
+	for _, t := range tokens {
+		if strings.HasPrefix(t, "-") {
+			break
+		}
+		kept = append(kept, t)
+	}
+	return strings.Join(kept, " ")
+}
+
 func AzApiHandler(azClient azcli.Client, cfg *config.ConfigData) func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
 	return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
 		args, ok := req.Params.Arguments.(map[string]interface{})
@@ -48,7 +64,7 @@ func AzApiHandler(azClient azcli.Client, cfg *config.ConfigData) func(ctx contex
 			errMsg := fmt.Sprintf("failed to execute Azure CLI command: %v", err)
 			logger.Errorf("AzApiHandler: %s", errMsg)
 			if cfg.TelemetryService != nil {
-				cfg.TelemetryService.TrackToolInvocation(ctx, req.Params.Name, cliCommand, false)
+				cfg.TelemetryService.TrackToolInvocation(ctx, req.Params.Name, sanitizeCliCommand(cliCommand), false)
 			}
 			return mcp.NewToolResultError(errMsg), nil
 		}
@@ -57,15 +73,15 @@ func AzApiHandler(azClient azcli.Client, cfg *config.ConfigData) func(ctx contex
 			errMsg := fmt.Sprintf("Azure CLI command failed: %s", result.Error)
 			logger.Errorf("AzApiHandler: %s", errMsg)
 			if cfg.TelemetryService != nil {
-				cfg.TelemetryService.TrackToolInvocation(ctx, req.Params.Name, cliCommand, false)
+				cfg.TelemetryService.TrackToolInvocation(ctx, req.Params.Name, sanitizeCliCommand(cliCommand), false)
 			}
 			return mcp.NewToolResultError(errMsg), nil
 		}
 
 		logger.Debugf("AzApiHandler: Command completed successfully")
 
 		if cfg.TelemetryService != nil {
-			cfg.TelemetryService.TrackToolInvocation(ctx, req.Params.Name, cliCommand, true)
+			cfg.TelemetryService.TrackToolInvocation(ctx, req.Params.Name, sanitizeCliCommand(cliCommand), true)
 		}
 
 		return mcp.NewToolResultText(string(result.Output)), nil
diff --git a/internal/components/azapi/handler_test.go b/internal/components/azapi/handler_test.go
@@ -296,3 +296,51 @@ func TestAzApiHandler_NilTelemetryService(t *testing.T) {
 		t.Fatal("expected success result")
 	}
 }
+
+func TestSanitizeCliCommand(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    string
+		expected string
+	}{
+		{
+			name:     "simple command without flags",
+			input:    "az group list",
+			expected: "az group list",
+		},
+		{
+			name:     "command with flags",
+			input:    "az aks show --name foo --resource-group bar",
+			expected: "az aks show",
+		},
+		{
+			name:     "command with secrets",
+			input:    "az login --service-principal --client-secret xxx",
+			expected: "az login",
+		},
+		{
+			name:     "deep subcommand with flags",
+			input:    "az network vnet subnet show --name sub1 --vnet-name vnet1",
+			expected: "az network vnet subnet show",
+		},
+		{
+			name:     "empty string",
+			input:    "",
+			expected: "",
+		},
+		{
+			name:     "command with short flags",
+			input:    "az group list -o table",
+			expected: "az group list",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			got := sanitizeCliCommand(tc.input)
+			if got != tc.expected {
+				t.Errorf("sanitizeCliCommand(%q) = %q, want %q", tc.input, got, tc.expected)
+			}
+		})
+	}
+}
