package aitools import ( "encoding/json" "fmt" "os" "path/filepath" "strings" ) // Mode 控制读写权限;写操作仍须确认。 const ( ModeReadOnly = "read_only" ModeReadWrite = "read_write" ) // Policy ~/.opsxdb/mcp.json type Policy struct { AllowedConnectionIDs []string `json:"allowed_connection_ids"` Mode string `json:"mode"` // read_only | read_write // AllowDangerous 为 true 时,在 mode=read_write 且已确认后允许 DROP/TRUNCATE/ALTER 等危险 SQL。 AllowDangerous bool `json:"allow_dangerous"` DefaultLimit int `json:"default_limit,omitempty"` MaxLimit int `json:"max_limit,omitempty"` } func DefaultPolicy() Policy { return Policy{ AllowedConnectionIDs: nil, Mode: ModeReadWrite, AllowDangerous: true, DefaultLimit: 100, MaxLimit: 500, } } func opsxdbDir() (string, error) { home, err := os.UserHomeDir() if err != nil { return "", err } return filepath.Join(home, ".opsxdb"), nil } func PolicyPath() (string, error) { dir, err := opsxdbDir() if err != nil { return "", err } return filepath.Join(dir, "mcp.json"), nil } // LoadPolicy 读 mcp.json;缺省文件返回 DefaultPolicy(白名单空=拒绝全部连接)。 func LoadPolicy() (Policy, error) { p := DefaultPolicy() path, err := PolicyPath() if err != nil { return p, err } b, err := os.ReadFile(path) if err != nil { if os.IsNotExist(err) { return p, nil } return p, err } // 已有文件沿用显式配置;旧版缺少危险 SQL 字段时仍视为禁用。 p.Mode = ModeReadOnly p.AllowDangerous = false if err := json.Unmarshal(b, &p); err != nil { return DefaultPolicy(), fmt.Errorf("mcp.json: %w", err) } p.Mode = strings.ToLower(strings.TrimSpace(p.Mode)) if p.Mode == "" { p.Mode = ModeReadOnly } if p.DefaultLimit <= 0 { p.DefaultLimit = 100 } if p.MaxLimit <= 0 { p.MaxLimit = 500 } if p.DefaultLimit > p.MaxLimit { p.DefaultLimit = p.MaxLimit } return p, nil } func (p Policy) IsReadOnly() bool { return p.Mode != ModeReadWrite } // WritesEnabled P1:仅当 mcp.json mode=read_write 时允许写路径(仍须确认令牌)。 func (p Policy) WritesEnabled() bool { return strings.EqualFold(strings.TrimSpace(p.Mode), ModeReadWrite) } func (p Policy) AllowConnection(id string) error { id = strings.TrimSpace(id) if id == "" { return fmt.Errorf("connection_id required") } if len(p.AllowedConnectionIDs) == 0 { return fmt.Errorf("connection not allowed: whitelist empty (configure allowed_connection_ids in ~/.opsxdb/mcp.json)") } for _, a := range p.AllowedConnectionIDs { if a == id { return nil } } return fmt.Errorf("connection not allowed: %s not in allowed_connection_ids", id) } func (p Policy) ClampLimit(n int) int { if n <= 0 { n = p.DefaultLimit } if n > p.MaxLimit { n = p.MaxLimit } return n } // SavePolicy 原子写入 ~/.opsxdb/mcp.json(先写临时文件再 rename)。 func SavePolicy(p Policy) error { p.Mode = strings.ToLower(strings.TrimSpace(p.Mode)) if p.Mode != ModeReadWrite { p.Mode = ModeReadOnly } if p.DefaultLimit <= 0 { p.DefaultLimit = 100 } if p.MaxLimit <= 0 { p.MaxLimit = 500 } if p.DefaultLimit > p.MaxLimit { p.DefaultLimit = p.MaxLimit } if p.AllowedConnectionIDs == nil { p.AllowedConnectionIDs = []string{} } // 去空白、去重保序 seen := map[string]struct{}{} clean := make([]string, 0, len(p.AllowedConnectionIDs)) for _, id := range p.AllowedConnectionIDs { id = strings.TrimSpace(id) if id == "" { continue } if _, ok := seen[id]; ok { continue } seen[id] = struct{}{} clean = append(clean, id) } p.AllowedConnectionIDs = clean path, err := PolicyPath() if err != nil { return err } dir := filepath.Dir(path) if err := os.MkdirAll(dir, 0o700); err != nil { return err } b, err := json.MarshalIndent(p, "", " ") if err != nil { return err } b = append(b, '\n') tmp := path + ".tmp" if err := os.WriteFile(tmp, b, 0o600); err != nil { return err } return os.Rename(tmp, path) } // AllowsDangerousDDL 危险 SQL(DROP/TRUNCATE/…)是否在确认后放行。 func (p Policy) AllowsDangerousDDL() bool { return p.AllowDangerous && p.WritesEnabled() }