opsxdb/server/internal/aitools/policy.go

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()
}