|
| 1 | +package commands |
| 2 | + |
| 3 | +import ( |
| 4 | + "fmt" |
| 5 | + "os" |
| 6 | + "strings" |
| 7 | + |
| 8 | + "github.com/f/mcptools/pkg/alias" |
| 9 | + "github.com/f/mcptools/pkg/client" |
| 10 | + "github.com/f/mcptools/pkg/guard" |
| 11 | + "github.com/spf13/cobra" |
| 12 | +) |
| 13 | + |
| 14 | +// Guard flags. |
| 15 | +const ( |
| 16 | + FlagAllow = "--allow" |
| 17 | + FlagAllowShort = "-a" |
| 18 | + FlagDeny = "--deny" |
| 19 | + FlagDenyShort = "-d" |
| 20 | +) |
| 21 | + |
| 22 | +var entityTypes = []string{ |
| 23 | + EntityTypeTool, |
| 24 | + EntityTypePrompt, |
| 25 | + EntityTypeRes, |
| 26 | +} |
| 27 | + |
| 28 | +// GuardCmd creates the guard command to filter tools, prompts, and resources. |
| 29 | +func GuardCmd() *cobra.Command { |
| 30 | + return &cobra.Command{ |
| 31 | + Use: "guard [--allow type:pattern] [--deny type:pattern] command args...", |
| 32 | + Short: "Filter tools, prompts, and resources using allow and deny patterns", |
| 33 | + Long: `Filter tools, prompts, and resources using allow and deny patterns. |
| 34 | +
|
| 35 | +Examples: |
| 36 | + mcp guard --allow tools:read_* --deny edit_*,write_*,create_* npx run @modelcontextprotocol/server-filesystem ~ |
| 37 | + mcp guard --allow prompts:system_* --deny tools:execute_* npx run @modelcontextprotocol/server-filesystem ~ |
| 38 | + mcp guard --allow tools:read_* fs # Using an alias |
| 39 | + |
| 40 | +Patterns can include wildcards: |
| 41 | + * matches any sequence of characters |
| 42 | + |
| 43 | +Entity types: |
| 44 | + tools: filter available tools |
| 45 | + prompts: filter available prompts |
| 46 | + resource: filter available resources`, |
| 47 | + DisableFlagParsing: true, |
| 48 | + SilenceUsage: true, |
| 49 | + Run: func(thisCmd *cobra.Command, args []string) { |
| 50 | + if len(args) == 1 && (args[0] == FlagHelp || args[0] == FlagHelpShort) { |
| 51 | + _ = thisCmd.Help() |
| 52 | + return |
| 53 | + } |
| 54 | + |
| 55 | + // Process and extract the allow and deny patterns |
| 56 | + allowPatterns, denyPatterns, cmdArgs := extractPatterns(args) |
| 57 | + |
| 58 | + // Process regular flags (format) |
| 59 | + parsedArgs := ProcessFlags(cmdArgs) |
| 60 | + |
| 61 | + // Print filtering info |
| 62 | + fmt.Fprintf(os.Stderr, "Guard filtering configuration:\n") |
| 63 | + for _, entityType := range entityTypes { |
| 64 | + if len(allowPatterns[entityType]) > 0 { |
| 65 | + fmt.Fprintf(os.Stderr, "Allowing %s matching: %s\n", entityType, strings.Join(allowPatterns[entityType], ", ")) |
| 66 | + } |
| 67 | + if len(denyPatterns[entityType]) > 0 { |
| 68 | + fmt.Fprintf(os.Stderr, "Denying %s matching: %s\n", entityType, strings.Join(denyPatterns[entityType], ", ")) |
| 69 | + } |
| 70 | + } |
| 71 | + |
| 72 | + // Check if we're using an alias for the server command |
| 73 | + if len(parsedArgs) == 1 { |
| 74 | + aliasName := parsedArgs[0] |
| 75 | + serverCmd, found := alias.GetServerCommand(aliasName) |
| 76 | + if found { |
| 77 | + fmt.Fprintf(os.Stderr, "Expanding alias '%s' to '%s'\n", aliasName, serverCmd) |
| 78 | + // Replace the alias with the actual command |
| 79 | + parsedArgs = client.ParseCommandString(serverCmd) |
| 80 | + } |
| 81 | + } |
| 82 | + |
| 83 | + // Verify we have a command to run |
| 84 | + if len(parsedArgs) == 0 { |
| 85 | + fmt.Fprintf(os.Stderr, "Error: command to execute is required\n") |
| 86 | + fmt.Fprintf(os.Stderr, "Example: mcp guard --allow tools:read_* npx -y @modelcontextprotocol/server-filesystem ~\n") |
| 87 | + os.Exit(1) |
| 88 | + } |
| 89 | + |
| 90 | + // Map our entity types to the guard proxy entity types |
| 91 | + guardAllowPatterns := map[string][]string{ |
| 92 | + "tool": allowPatterns[EntityTypeTool], |
| 93 | + "prompt": allowPatterns[EntityTypePrompt], |
| 94 | + "resource": allowPatterns[EntityTypeRes], |
| 95 | + } |
| 96 | + guardDenyPatterns := map[string][]string{ |
| 97 | + "tool": denyPatterns[EntityTypeTool], |
| 98 | + "prompt": denyPatterns[EntityTypePrompt], |
| 99 | + "resource": denyPatterns[EntityTypeRes], |
| 100 | + } |
| 101 | + |
| 102 | + // Run the guard proxy with the filtered environment |
| 103 | + fmt.Fprintf(os.Stderr, "Running command with filtered environment: %s\n", strings.Join(parsedArgs, " ")) |
| 104 | + if err := guard.RunFilterServer(guardAllowPatterns, guardDenyPatterns, parsedArgs); err != nil { |
| 105 | + fmt.Fprintf(os.Stderr, "Error: %v\n", err) |
| 106 | + os.Exit(1) |
| 107 | + } |
| 108 | + }, |
| 109 | + } |
| 110 | +} |
| 111 | + |
| 112 | +// extractPatterns processes arguments to extract allow and deny patterns. |
| 113 | +func extractPatterns(args []string) (map[string][]string, map[string][]string, []string) { |
| 114 | + allowPatterns := make(map[string][]string) |
| 115 | + denyPatterns := make(map[string][]string) |
| 116 | + |
| 117 | + // Initialize maps for all entity types |
| 118 | + for _, entityType := range entityTypes { |
| 119 | + allowPatterns[entityType] = []string{} |
| 120 | + denyPatterns[entityType] = []string{} |
| 121 | + } |
| 122 | + |
| 123 | + cmdArgs := []string{} |
| 124 | + i := 0 |
| 125 | + for i < len(args) { |
| 126 | + switch { |
| 127 | + case (args[i] == FlagAllow || args[i] == FlagAllowShort) && i+1 < len(args): |
| 128 | + // Process --allow flag |
| 129 | + patternsStr := args[i+1] |
| 130 | + processPatternString(patternsStr, allowPatterns) |
| 131 | + i += 2 |
| 132 | + case (args[i] == FlagDeny || args[i] == FlagDenyShort) && i+1 < len(args): |
| 133 | + // Process --deny flag |
| 134 | + patternsStr := args[i+1] |
| 135 | + processPatternString(patternsStr, denyPatterns) |
| 136 | + i += 2 |
| 137 | + default: |
| 138 | + // Not a flag we recognize, pass it along |
| 139 | + cmdArgs = append(cmdArgs, args[i]) |
| 140 | + i++ |
| 141 | + } |
| 142 | + } |
| 143 | + |
| 144 | + return allowPatterns, denyPatterns, cmdArgs |
| 145 | +} |
| 146 | + |
| 147 | +// processPatternString processes a comma-separated pattern string. |
| 148 | +func processPatternString(patternsStr string, patternMap map[string][]string) { |
| 149 | + patterns := strings.Split(patternsStr, ",") |
| 150 | + |
| 151 | + for _, pattern := range patterns { |
| 152 | + pattern = strings.TrimSpace(pattern) |
| 153 | + if pattern == "" { |
| 154 | + continue |
| 155 | + } |
| 156 | + |
| 157 | + parts := strings.SplitN(pattern, ":", 2) |
| 158 | + if len(parts) != 2 { |
| 159 | + // If no type specified, assume it's a tool pattern |
| 160 | + patternMap[EntityTypeTool] = append(patternMap[EntityTypeTool], pattern) |
| 161 | + continue |
| 162 | + } |
| 163 | + |
| 164 | + entityType := strings.ToLower(parts[0]) |
| 165 | + patternValue := parts[1] |
| 166 | + |
| 167 | + // Map entity type to known types |
| 168 | + switch entityType { |
| 169 | + case "tool", "tools": |
| 170 | + patternMap[EntityTypeTool] = append(patternMap[EntityTypeTool], patternValue) |
| 171 | + case "prompt", "prompts": |
| 172 | + patternMap[EntityTypePrompt] = append(patternMap[EntityTypePrompt], patternValue) |
| 173 | + case "resource", "resources", "res": |
| 174 | + patternMap[EntityTypeRes] = append(patternMap[EntityTypeRes], patternValue) |
| 175 | + default: |
| 176 | + // Unknown entity type, treat as tool pattern |
| 177 | + patternMap[EntityTypeTool] = append(patternMap[EntityTypeTool], pattern) |
| 178 | + } |
| 179 | + } |
| 180 | +} |
0 commit comments