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