From 01c38c26d0006887c26b3779ef90668fcc5519ba Mon Sep 17 00:00:00 2001 From: xbingW Date: Wed, 9 Apr 2025 19:16:05 +0800 Subject: [PATCH] feat: whitelist --- mcp_go/README.md | 20 ++++---- ...reate_blacklist_rule.go => create_rule.go} | 8 +-- mcp_go/internal/tools/init.go | 1 + .../tools/rule/create_blacklist_rule.go | 2 +- .../tools/rule/create_whitelist_rule.go | 50 +++++++++++++++++++ 5 files changed, 67 insertions(+), 14 deletions(-) rename mcp_go/internal/api/rule/{create_blacklist_rule.go => create_rule.go} (71%) create mode 100644 mcp_go/internal/tools/rule/create_whitelist_rule.go diff --git a/mcp_go/README.md b/mcp_go/README.md index 6ad64cd..d31aa6f 100644 --- a/mcp_go/README.md +++ b/mcp_go/README.md @@ -99,10 +99,11 @@ go run main.go ### Application Management -- **create_application** - Create protection application - - `domain`: Application domain (string, required) - - `port`: Listening port (number, required) - - `upstream`: Upstream address (string, required) +- **create_application** + +### Rule Management +- **create_blacklist_rule** +- **create_whitelist_rule** For more API details, please refer to the [API Documentation](https://master.safeline-ce.staging.dev.in.chaitin.net:9443/swagger/index.html#). @@ -169,12 +170,14 @@ func (t *ToolName) Execute(ctx context.Context, params ToolParams) (result ToolR } ``` -3. **Create API Implementation** +3. **[Optional]Create API Implementation** + +If you need to use some APIs that have not been implemented yet, you need to create corresponding files in the api directory for implementation - Create same directory structure under `internal/api` - - File name should match tool file + - File name should match tool func - Example: `internal/api/app/create_application.go` -4. **API Implementation Template** +**API Implementation Template** ```go package app @@ -202,7 +205,7 @@ func APIName(ctx context.Context, req *RequestType) (ResultType, error) { return resp.Data, nil } ``` -5. **Tool Registration (init.go)** +4. **Tool Registration (init.go)** The tool registration file `internal/tools/init.go` is used to centrally manage all tool registrations - Register all tools uniformly in the `init()` function @@ -221,7 +224,6 @@ The tool registration file `internal/tools/init.go` is used to centrally manage 1. **Naming Conventions** - Use lowercase letters and underscores for tool names - File names should match tool names - - API implementation files should have the same name as tool files 2. **Directory Organization** - Divide directories by functional modules (e.g., app, rule, etc.) diff --git a/mcp_go/internal/api/rule/create_blacklist_rule.go b/mcp_go/internal/api/rule/create_rule.go similarity index 71% rename from mcp_go/internal/api/rule/create_blacklist_rule.go rename to mcp_go/internal/api/rule/create_rule.go index 5c63715..6eb8843 100644 --- a/mcp_go/internal/api/rule/create_blacklist_rule.go +++ b/mcp_go/internal/api/rule/create_rule.go @@ -7,7 +7,7 @@ import ( "github.com/chaitin/SafeLine/mcp_server/pkg/errors" ) -type CreateBlacklistRuleRequest struct { +type CreateRuleRequest struct { Name string `json:"name"` IP []string `json:"ip"` IsEnabled bool `json:"is_enabled"` @@ -15,8 +15,8 @@ type CreateBlacklistRuleRequest struct { Action int `json:"action"` } -// CreateBlacklistRule Create new blacklist rule -func CreateBlacklistRule(ctx context.Context, req *CreateBlacklistRuleRequest) (int64, error) { +// CreateRule Create new rule +func CreateRule(ctx context.Context, req *CreateRuleRequest) (int64, error) { if req == nil { return 0, errors.New("request is required") } @@ -24,7 +24,7 @@ func CreateBlacklistRule(ctx context.Context, req *CreateBlacklistRuleRequest) ( var resp api.Response[int64] err := api.Service().Post(ctx, "/api/open/policy", req, &resp) if err != nil { - return 0, errors.Wrap(err, "failed to create blacklist rule") + return 0, errors.Wrap(err, "failed to create policy rule") } if resp.Err != nil { diff --git a/mcp_go/internal/tools/init.go b/mcp_go/internal/tools/init.go index 8331f47..8c10199 100644 --- a/mcp_go/internal/tools/init.go +++ b/mcp_go/internal/tools/init.go @@ -8,4 +8,5 @@ import ( func init() { AppendTool(&app.CreateApp{}) AppendTool(&rule.CreateBlacklistRule{}) + AppendTool(&rule.CreateWhitelistRule{}) } diff --git a/mcp_go/internal/tools/rule/create_blacklist_rule.go b/mcp_go/internal/tools/rule/create_blacklist_rule.go index 29a4ca5..2e04a29 100644 --- a/mcp_go/internal/tools/rule/create_blacklist_rule.go +++ b/mcp_go/internal/tools/rule/create_blacklist_rule.go @@ -28,7 +28,7 @@ func (t *CreateBlacklistRule) Validate(params CreateBlacklistRuleParams) error { } func (t *CreateBlacklistRule) Execute(ctx context.Context, params CreateBlacklistRuleParams) (int64, error) { - id, err := rule.CreateBlacklistRule(ctx, &rule.CreateBlacklistRuleRequest{ + id, err := rule.CreateRule(ctx, &rule.CreateRuleRequest{ Name: params.Name, IP: params.IP, IsEnabled: true, diff --git a/mcp_go/internal/tools/rule/create_whitelist_rule.go b/mcp_go/internal/tools/rule/create_whitelist_rule.go new file mode 100644 index 0000000..16fd66f --- /dev/null +++ b/mcp_go/internal/tools/rule/create_whitelist_rule.go @@ -0,0 +1,50 @@ +package rule + +import ( + "context" + + "github.com/chaitin/SafeLine/mcp_server/internal/api" + "github.com/chaitin/SafeLine/mcp_server/internal/api/rule" +) + +type CreateWhitelistRule struct{} + +type CreateWhitelistRuleParams struct { + Name string `json:"name" desc:"name" required:"true"` + IP []string `json:"ip" desc:"ip" required:"false"` +} + +func (t *CreateWhitelistRule) Name() string { + return "create_whitelist_rule" +} + +func (t *CreateWhitelistRule) Description() string { + return "create a new whitelist rule" +} + +func (t *CreateWhitelistRule) Validate(params CreateWhitelistRuleParams) error { + return nil +} + +func (t *CreateWhitelistRule) Execute(ctx context.Context, params CreateWhitelistRuleParams) (int64, error) { + id, err := rule.CreateRule(ctx, &rule.CreateRuleRequest{ + Name: params.Name, + IP: params.IP, + IsEnabled: true, + Action: int(api.PolicyRuleActionAllow), + Pattern: [][]api.Pattern{ + { + { + K: api.KeySrcIP, + Op: api.OpEq, + V: params.IP, + SubK: "", + }, + }, + }, + }) + if err != nil { + return 0, err + } + return id, nil +}