BlackBean/internal/agent/tools_docker.go

package agent

import (
	"context"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"sort"
	"strconv"
	"strings"

	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/image"
	"github.com/docker/docker/client"
)

// Docker 工具：通过 Docker API 操作本地或远程 Docker。
// socket 地址通过 DockerConfigStore 动态读取，用户在 Web 页面配置后即时生效。

func (t *Toolset) dockerClient() (*client.Client, error) {
	socket := t.dockerSocket()
	cli, err := client.NewClientWithOpts(
		client.WithHost(socket),
		client.WithAPIVersionNegotiation(),
	)
	if err != nil {
		return nil, fmt.Errorf("创建 Docker 客户端失败: %w", err)
	}
	return cli, nil
}

func (t *Toolset) dockerPS(ctx context.Context, params map[string]any) ToolResult {
	all, _ := params["all"].(bool)
	cli, err := t.dockerClient()
	if err != nil {
		return ToolResult{Success: false, Output: err.Error()}
	}
	defer cli.Close()

	ctx, cancel := context.WithTimeout(ctx, t.Timeout)
	defer cancel()
	list, err := cli.ContainerList(ctx, container.ListOptions{All: all})
	if err != nil {
		return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
	}
	if len(list) == 0 {
		return ToolResult{Success: true, Output: "没有找到容器"}
	}

	sort.Slice(list, func(i, j int) bool {
		return list[i].Created > list[j].Created
	})
	var out strings.Builder
	fmt.Fprintf(&out, "%-14s %-26s %-20s %-12s %s\n", "CONTAINER ID", "NAME", "IMAGE", "STATUS", "PORTS")
	for _, c := range list {
		name := strings.TrimPrefix(strings.Join(c.Names, ","), "/")
		fmt.Fprintf(&out, "%-14s %-26s %-20s %-12s %s\n",
			c.ID[:min(12, len(c.ID))], name, c.Image, c.Status, dockerPorts(c.Ports))
	}
	return ToolResult{Success: true, Output: t.truncate(out.String())}
}

func (t *Toolset) dockerImages(ctx context.Context, params map[string]any) ToolResult {
	cli, err := t.dockerClient()
	if err != nil {
		return ToolResult{Success: false, Output: err.Error()}
	}
	defer cli.Close()

	ctx, cancel := context.WithTimeout(ctx, t.Timeout)
	defer cancel()
	list, err := cli.ImageList(ctx, image.ListOptions{All: true})
	if err != nil {
		return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
	}
	if len(list) == 0 {
		return ToolResult{Success: true, Output: "没有找到镜像"}
	}

	sort.Slice(list, func(i, j int) bool {
		return list[i].Created > list[j].Created
	})
	var out strings.Builder
	fmt.Fprintf(&out, "%-22s %-28s %-14s %s\n", "REPOSITORY", "TAG", "IMAGE ID", "SIZE")
	for _, img := range list {
		if len(img.RepoTags) == 0 {
			continue
		}
		for _, tag := range img.RepoTags {
			repo, imageTag, _ := strings.Cut(tag, ":")
			id := img.ID
			if len(id) > 12 {
				id = id[7:19]
			}
			fmt.Fprintf(&out, "%-22s %-28s %-14s %s\n", repo, imageTag, id, formatBytes(img.Size))
		}
	}
	return ToolResult{Success: true, Output: t.truncate(out.String())}
}

func (t *Toolset) dockerLogs(ctx context.Context, params map[string]any) ToolResult {
	containerID, _ := params["container"].(string)
	containerID = strings.TrimSpace(containerID)
	if containerID == "" {
		return ToolResult{Success: false, Output: "container 不能为空"}
	}
	tail := intParam(params, "tail")
	if tail <= 0 {
		tail = 100
	}

	cli, err := t.dockerClient()
	if err != nil {
		return ToolResult{Success: false, Output: err.Error()}
	}
	defer cli.Close()

	ctx, cancel := context.WithTimeout(ctx, t.Timeout)
	defer cancel()
	resp, err := cli.ContainerLogs(ctx, containerID, container.LogsOptions{
		ShowStdout: true,
		ShowStderr: true,
		Tail:       strconv.Itoa(tail),
	})
	if err != nil {
		return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
	}
	defer resp.Close()

	raw, err := io.ReadAll(resp)
	if err != nil {
		return ToolResult{Success: false, Output: "读取 Docker 日志失败: " + err.Error()}
	}
	return ToolResult{Success: true, Output: t.truncate(demuxDockerLogs(raw))}
}

// demuxDockerLogs 解析 Docker 日志流：多路复用流（stdout/stderr 分帧）按帧解出，
// 普通文本流（TTY 容器）直接返回。
func demuxDockerLogs(data []byte) string {
	if len(data) >= 8 && isMultiplexedFrame(data[0:8]) {
		var out strings.Builder
		for i := 0; i+8 <= len(data); {
			frameSize := int(binary.BigEndian.Uint32(data[i+4 : i+8]))
			i += 8
			if frameSize < 0 || i+frameSize > len(data) {
				break
			}
			out.Write(data[i : i+frameSize])
			i += frameSize
		}
		return out.String()
	}
	return string(data)
}

func isMultiplexedFrame(header []byte) bool {
	streamType := header[0]
	if streamType != 0 && streamType != 1 && streamType != 2 {
		return false
	}
	return header[1] == 0 && header[2] == 0 && header[3] == 0
}

func (t *Toolset) dockerInspect(ctx context.Context, params map[string]any) ToolResult {
	containerID, _ := params["container"].(string)
	containerID = strings.TrimSpace(containerID)
	if containerID == "" {
		return ToolResult{Success: false, Output: "container 不能为空"}
	}

	cli, err := t.dockerClient()
	if err != nil {
		return ToolResult{Success: false, Output: err.Error()}
	}
	defer cli.Close()

	ctx, cancel := context.WithTimeout(ctx, t.Timeout)
	defer cancel()
	inspect, err := cli.ContainerInspect(ctx, containerID)
	if err != nil {
		return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
	}

	state := inspect.State
	if state == nil {
		state = &container.State{}
	}
	summary := map[string]any{
		"id":            inspect.ID,
		"name":          strings.TrimPrefix(inspect.Name, "/"),
		"image":         inspect.Config.Image,
		"status":        state.Status,
		"running":       state.Running,
		"exit_code":     state.ExitCode,
		"restart_count": inspect.RestartCount,
		"created":       inspect.Created,
		"command":       inspect.Config.Cmd,
		"env_count":     len(inspect.Config.Env),
		"ports":         inspect.NetworkSettings.Ports,
		"network_mode":  string(inspect.HostConfig.NetworkMode),
		"ip_address":    inspect.NetworkSettings.IPAddress,
	}
	data, _ := json.MarshalIndent(summary, "", "  ")
	return ToolResult{Success: true, Output: t.truncate(string(data))}
}

func (t *Toolset) dockerExec(ctx context.Context, params map[string]any) ToolResult {
	containerID, _ := params["container"].(string)
	containerID = strings.TrimSpace(containerID)
	command, _ := params["command"].(string)
	command = strings.TrimSpace(command)
	if containerID == "" || command == "" {
		return ToolResult{Success: false, Output: "container 和 command 不能为空"}
	}

	cli, err := t.dockerClient()
	if err != nil {
		return ToolResult{Success: false, Output: err.Error()}
	}
	defer cli.Close()

	ctx, cancel := context.WithTimeout(ctx, t.Timeout)
	defer cancel()
	execCfg := container.ExecOptions{
		AttachStdout: true,
		AttachStderr: true,
		Cmd:          []string{"sh", "-c", command},
	}
	created, err := cli.ContainerExecCreate(ctx, containerID, execCfg)
	if err != nil {
		return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
	}
	attach, err := cli.ContainerExecAttach(ctx, created.ID, container.ExecAttachOptions{})
	if err != nil {
		return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
	}
	defer attach.Close()

	raw, err := io.ReadAll(attach.Reader)
	if err != nil {
		return ToolResult{Success: false, Output: "读取命令输出失败: " + err.Error()}
	}
	return ToolResult{Success: true, Output: t.truncate(demuxDockerLogs(raw))}
}

func (t *Toolset) dockerStart(ctx context.Context, params map[string]any) ToolResult {
	containerID, _ := params["container"].(string)
	containerID = strings.TrimSpace(containerID)
	if containerID == "" {
		return ToolResult{Success: false, Output: "container 不能为空"}
	}

	cli, err := t.dockerClient()
	if err != nil {
		return ToolResult{Success: false, Output: err.Error()}
	}
	defer cli.Close()

	ctx, cancel := context.WithTimeout(ctx, t.Timeout)
	defer cancel()
	if err := cli.ContainerStart(ctx, containerID, container.StartOptions{}); err != nil {
		return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
	}
	return ToolResult{Success: true, Output: "容器 " + containerID + " 已启动"}
}

func (t *Toolset) dockerStop(ctx context.Context, params map[string]any) ToolResult {
	containerID, _ := params["container"].(string)
	containerID = strings.TrimSpace(containerID)
	if containerID == "" {
		return ToolResult{Success: false, Output: "container 不能为空"}
	}
	timeoutSeconds := intParam(params, "timeout")

	cli, err := t.dockerClient()
	if err != nil {
		return ToolResult{Success: false, Output: err.Error()}
	}
	defer cli.Close()

	ctx, cancel := context.WithTimeout(ctx, t.Timeout)
	defer cancel()
	var timeout *int
	if timeoutSeconds > 0 {
		timeout = &timeoutSeconds
	}
	if err := cli.ContainerStop(ctx, containerID, container.StopOptions{Timeout: timeout}); err != nil {
		return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()}
	}
	return ToolResult{Success: true, Output: "容器 " + containerID + " 已停止"}
}

func dockerPorts(ports []container.Port) string {
	if len(ports) == 0 {
		return ""
	}
	var parts []string
	for _, p := range ports {
		host := ""
		if p.PublicPort > 0 {
			host = fmt.Sprintf("%d->", p.PublicPort)
		}
		parts = append(parts, fmt.Sprintf("%s%d/%s", host, p.PrivatePort, p.Type))
	}
	return strings.Join(parts, ", ")
}

func formatBytes(size int64) string {
	const unit = 1024
	if size < unit {
		return fmt.Sprintf("%dB", size)
	}
	div, exp := int64(unit), 0
	for n := size / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f%cB", float64(size)/float64(div), "KMGTPE"[exp])
}