ops-mcp/internal/tool/docker.go

package tool

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"os/exec"
	"regexp"
	"strings"
	"time"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
)

// DockerTool 提供安全的 Docker 容器信息查询工具。
// 使用 docker CLI 而非 SDK，避免跨平台编译问题，且更轻量。
type DockerTool struct {
	containerRegex *regexp.Regexp
	maxLogBytes    int64
	maxLogLines    int
	maxLogSince    time.Duration
}

type DockerOpts struct {
	ContainerPattern string
	MaxLogBytes      int64
	MaxLogLines      int
	MaxLogSince      time.Duration
}

func NewDockerTool(opts DockerOpts) (*DockerTool, error) {
	pattern := opts.ContainerPattern
	if pattern == "" {
		pattern = ".*"
	}

	re, err := regexp.Compile(pattern)
	if err != nil {
		return nil, fmt.Errorf("compile container pattern: %w", err)
	}

	return &DockerTool{
		containerRegex: re,
		maxLogBytes:    opts.MaxLogBytes,
		maxLogLines:    opts.MaxLogLines,
		maxLogSince:    opts.MaxLogSince,
	}, nil
}

func (d *DockerTool) Name() string        { return "docker" }
func (d *DockerTool) Description() string { return "Docker 容器日志和安全查询" }

func (d *DockerTool) Initialize(_ context.Context) error {
	if _, err := exec.LookPath("docker"); err != nil {
		return fmt.Errorf("docker CLI not found in PATH")
	}
	slog.Info("docker CLI found")
	return nil
}

func (d *DockerTool) Shutdown(_ context.Context) error { return nil }

func (d *DockerTool) HealthCheck(ctx context.Context) error {
	return d.dockerCmd(ctx, "version").Run()
}

func (d *DockerTool) Register(mcpServer *server.MCPServer) error {
	mcpServer.AddTool(mcp.NewTool("docker_ps",
		mcp.WithDescription("列出 Docker 容器"),
		mcp.WithString("filter", mcp.Description("按名称过滤（支持正则）")),
		mcp.WithBoolean("all", mcp.Description("是否包含已停止的容器，默认 false")),
	), d.handleList)

	mcpServer.AddTool(mcp.NewTool("docker_logs",
		mcp.WithDescription("获取容器日志（受白名单和大小限制保护）"),
		mcp.WithString("container", mcp.Required(), mcp.Description("容器名称或 ID")),
		mcp.WithNumber("tail", mcp.Description("返回最后 N 行，默认 100")),
		mcp.WithString("since", mcp.Description("从多久前开始，如 15m、1h，默认 15m")),
	), d.handleLogs)

	mcpServer.AddTool(mcp.NewTool("docker_inspect",
		mcp.WithDescription("查看容器详细信息（受白名单保护）"),
		mcp.WithString("container", mcp.Required(), mcp.Description("容器名称或 ID")),
	), d.handleInspect)

	return nil
}

// --- 安全校验 ---

func (d *DockerTool) validateContainer(name string) error {
	if !d.containerRegex.MatchString(name) {
		return fmt.Errorf("container %q not in allowed pattern", name)
	}
	return nil
}

// --- docker CLI 封装 ---

func (d *DockerTool) dockerCmd(ctx context.Context, args ...string) *exec.Cmd {
	return exec.CommandContext(ctx, "docker", args...)
}

func (d *DockerTool) dockerOutput(ctx context.Context, args ...string) ([]byte, error) {
	cmd := d.dockerCmd(ctx, args...)
	var stderr bytes.Buffer
	cmd.Stderr = &stderr

	output, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("docker %s: %v\n%s", strings.Join(args, " "), err, stderr.String())
	}
	return output, nil
}

// --- Handlers ---

func (d *DockerTool) handleList(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	args := getArgs(req)
	filterName, _ := args["filter"].(string)
	showAll, _ := args["all"].(bool)

	dockerArgs := []string{"ps", "--format", "{{json .}}"}
	if showAll {
		dockerArgs = append(dockerArgs, "-a")
	}

	output, err := d.dockerOutput(ctx, dockerArgs...)
	if err != nil {
		return mcp.NewToolResultError(err.Error()), nil
	}

	type dockerPsRow struct {
		Names     string `json:"Names"`
		Image     string `json:"Image"`
		Status    string `json:"Status"`
		CreatedAt string `json:"CreatedAt"`
		Ports     string `json:"Ports"`
	}

	var result []dockerPsRow
	for _, line := range strings.Split(strings.TrimSpace(string(output)), "\n") {
		if line == "" {
			continue
		}
		var row dockerPsRow
		if err := json.Unmarshal([]byte(line), &row); err != nil {
			continue
		}

		// 应用名称过滤
		if filterName != "" {
			re, err := regexp.Compile(filterName)
			if err != nil {
				continue
			}
			if !re.MatchString(row.Names) {
				continue
			}
		}

		// 白名单检查
		if err := d.validateContainer(row.Names); err != nil {
			continue
		}

		result = append(result, row)
	}

	data, _ := json.MarshalIndent(result, "", "  ")
	return mcp.NewToolResultText(fmt.Sprintf("```json\n%s\n```", string(data))), nil
}

func (d *DockerTool) handleLogs(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	args := getArgs(req)
	containerName, _ := args["container"].(string)

	if err := d.validateContainer(containerName); err != nil {
		return mcp.NewToolResultError(err.Error()), nil
	}

	tail := "100"
	if t, ok := args["tail"].(float64); ok {
		tail = fmt.Sprintf("%d", int(t))
	}

	since := "15m"
	if s, ok := args["since"].(string); ok && s != "" {
		dur, err := time.ParseDuration(s)
		if err == nil && d.maxLogSince > 0 && dur > d.maxLogSince {
			since = d.maxLogSince.String()
		} else {
			since = s
		}
	}

	output, err := d.dockerOutput(ctx, "logs", "--tail", tail, "--since", since, "--timestamps", containerName)
	if err != nil {
		// docker logs 返回非 0 时容器可能不存在
		return mcp.NewToolResultError(fmt.Sprintf("docker logs failed: %v\nOutput: %s", err, string(output))), nil
	}

	// 截断字节数
	text := string(output)
	if len(text) > int(d.maxLogBytes) {
		text = text[len(text)-int(d.maxLogBytes):]
	}

	// 截断行数
	lines := strings.Split(text, "\n")
	if len(lines) > d.maxLogLines {
		lines = lines[len(lines)-d.maxLogLines:]
	}

	return mcp.NewToolResultText(strings.Join(lines, "\n")), nil
}

func (d *DockerTool) handleInspect(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	args := getArgs(req)
	containerName, _ := args["container"].(string)

	if err := d.validateContainer(containerName); err != nil {
		return mcp.NewToolResultError(err.Error()), nil
	}

	output, err := d.dockerOutput(ctx, "inspect", containerName)
	if err != nil {
		return mcp.NewToolResultError(err.Error()), nil
	}

	// 只提取安全字段
	var inspects []struct {
		Name  string `json:"Name"`
		ID    string `json:"Id"`
		Image string `json:"Image"`
		State struct {
			Status    string `json:"Status"`
			Running   bool   `json:"Running"`
			StartedAt string `json:"StartedAt"`
			Pid       int    `json:"Pid"`
		} `json:"State"`
		Created string `json:"Created"`
		Config  struct {
			Image string `json:"Image"`
			Env   []string `json:"Env"`
		} `json:"Config"`
		Mounts []struct {
			Source      string `json:"Source"`
			Destination string `json:"Destination"`
			Mode        string `json:"Mode"`
		} `json:"Mounts"`
	}

	if err := json.Unmarshal(output, &inspects); err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("parse inspect: %v", err)), nil
	}

	if len(inspects) == 0 {
		return mcp.NewToolResultError("container not found"), nil
	}

	insp := inspects[0]
	// 不暴露 Env（含密钥），只暴露安全信息
	safe := map[string]any{
		"name":    strings.TrimPrefix(insp.Name, "/"),
		"id":      insp.ID[:12],
		"image":   insp.Config.Image,
		"state":   insp.State,
		"created": insp.Created,
		"mounts":  insp.Mounts,
	}

	data, _ := json.MarshalIndent(safe, "", "  ")
	return mcp.NewToolResultText(fmt.Sprintf("```json\n%s\n```", string(data))), nil
}