flnxi-Wu/flnxi-agent-go/internal/docker/client.go

package docker

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"time"

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

type Client struct {
	cli *client.Client
}

func NewClient() (*Client, error) {
	cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
	if err != nil {
		return nil, err
	}
	return &Client{cli: cli}, nil
}

func (c *Client) Close() error {
	return c.cli.Close()
}

// PullImage 拉取镜像
func (c *Client) PullImage(ctx context.Context, image string, authStr string) error {
	return c.PullImageWithProgress(ctx, image, authStr, nil)
}

// PullImageWithProgress 拉取镜像并报告进度
func (c *Client) PullImageWithProgress(ctx context.Context, image string, authStr string, onProgress func(current, total int64, status string)) error {
	opts := types.ImagePullOptions{}
	if authStr != "" {
		opts.RegistryAuth = authStr
	}

	reader, err := c.cli.ImagePull(ctx, image, opts)
	if err != nil {
		return err
	}
	defer reader.Close()

	// 解析 Docker 拉取进度
	decoder := json.NewDecoder(reader)
	layerProgress := make(map[string]int64)
	layerTotal := make(map[string]int64)

	for {
		var event struct {
			Status         string `json:"status"`
			ID             string `json:"id"`
			Progress       string `json:"progress"`
			ProgressDetail struct {
				Current int64 `json:"current"`
				Total   int64 `json:"total"`
			} `json:"progressDetail"`
		}

		if err := decoder.Decode(&event); err != nil {
			if err == io.EOF {
				break
			}
			continue
		}

		if event.ID != "" && event.ProgressDetail.Total > 0 {
			layerProgress[event.ID] = event.ProgressDetail.Current
			layerTotal[event.ID] = event.ProgressDetail.Total
		}

		if onProgress != nil {
			var current, total int64
			for id, t := range layerTotal {
				total += t
				current += layerProgress[id]
			}
			status := event.Status
			if event.ID != "" {
				status = event.Status + " " + event.ID
			}
			onProgress(current, total, status)
		}
	}

	return nil
}

// ListContainers 列出容器
func (c *Client) ListContainers(ctx context.Context, nameFilter string) ([]types.Container, error) {
	opts := types.ContainerListOptions{All: true}
	if nameFilter != "" {
		opts.Filters = filters.NewArgs(filters.Arg("name", nameFilter))
	}
	return c.cli.ContainerList(ctx, opts)
}

// ContainerInfo 容器详细信息
type ContainerInfo struct {
	ID         string
	ImageID    string
	State      string
	Config     *container.Config
	HostConfig *container.HostConfig
}

// GetContainer 获取容器信息
func (c *Client) GetContainer(ctx context.Context, name string) (*ContainerInfo, error) {
	containers, err := c.ListContainers(ctx, name)
	if err != nil {
		return nil, err
	}
	for _, cont := range containers {
		for _, n := range cont.Names {
			if n == "/"+name || n == name {
				// 获取详细信息
				inspect, err := c.cli.ContainerInspect(ctx, cont.ID)
				if err != nil {
					return nil, err
				}
				return &ContainerInfo{
					ID:         cont.ID,
					ImageID:    cont.ImageID,
					State:      cont.State,
					Config:     inspect.Config,
					HostConfig: inspect.HostConfig,
				}, nil
			}
		}
	}
	return nil, nil
}

// StopContainer 停止容器
func (c *Client) StopContainer(ctx context.Context, containerID string) error {
	timeout := 30
	return c.cli.ContainerStop(ctx, containerID, container.StopOptions{Timeout: &timeout})
}

// RemoveContainer 删除容器
func (c *Client) RemoveContainer(ctx context.Context, containerID string) error {
	return c.cli.ContainerRemove(ctx, containerID, types.ContainerRemoveOptions{Force: true})
}

// StartContainer 启动容器
func (c *Client) StartContainer(ctx context.Context, containerID string) error {
	return c.cli.ContainerStart(ctx, containerID, types.ContainerStartOptions{})
}

// RestartContainer 重启容器
func (c *Client) RestartContainer(ctx context.Context, containerID string) error {
	timeout := 30
	return c.cli.ContainerRestart(ctx, containerID, container.StopOptions{Timeout: &timeout})
}

// CreateAndStartContainer 创建并启动容器（如果已存在则跳过创建，直接启动）
func (c *Client) CreateAndStartContainer(ctx context.Context, config *container.Config, hostConfig *container.HostConfig, name string) (string, error) {
	// 检查容器是否已存在
	existing, err := c.GetContainer(ctx, name)
	if err != nil {
		return "", err
	}

	if existing != nil {
		// 容器已存在，确保它在运行
		if existing.State != "running" {
			if err := c.cli.ContainerStart(ctx, existing.ID, types.ContainerStartOptions{}); err != nil {
				return "", err
			}
		}
		return existing.ID, nil
	}

	// 容器不存在，创建新容器
	resp, err := c.cli.ContainerCreate(ctx, config, hostConfig, nil, nil, name)
	if err != nil {
		return "", err
	}

	if err := c.cli.ContainerStart(ctx, resp.ID, types.ContainerStartOptions{}); err != nil {
		return "", err
	}

	return resp.ID, nil
}

// WaitForHealthy 等待容器健康
func (c *Client) WaitForHealthy(ctx context.Context, containerID string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		inspect, err := c.cli.ContainerInspect(ctx, containerID)
		if err != nil {
			return err
		}

		if inspect.State.Health != nil && inspect.State.Health.Status == "healthy" {
			return nil
		}

		if inspect.State.Running {
			// 如果没有健康检查，只要运行就算成功
			if inspect.State.Health == nil {
				return nil
			}
		}

		time.Sleep(2 * time.Second)
	}
	return nil
}

// GetDockerInfo 获取 Docker 信息
func (c *Client) GetDockerInfo(ctx context.Context) (types.Info, error) {
	return c.cli.Info(ctx)
}

// CreateNetworkIfNotExists 创建网络（如果不存在）
func (c *Client) CreateNetworkIfNotExists(ctx context.Context, name string) error {
	networks, err := c.cli.NetworkList(ctx, types.NetworkListOptions{
		Filters: filters.NewArgs(filters.Arg("name", name)),
	})
	if err != nil {
		return err
	}

	for _, n := range networks {
		if n.Name == name {
			return nil // 网络已存在
		}
	}

	_, err = c.cli.NetworkCreate(ctx, name, types.NetworkCreate{
		Driver: "bridge",
	})
	return err
}

// RemoveImage 删除镜像
func (c *Client) RemoveImage(ctx context.Context, imageID string) error {
	_, err := c.cli.ImageRemove(ctx, imageID, types.ImageRemoveOptions{Force: false, PruneChildren: true})
	return err
}

// PruneImages 清理悬空镜像
func (c *Client) PruneImages(ctx context.Context) error {
	_, err := c.cli.ImagesPrune(ctx, filters.NewArgs(filters.Arg("dangling", "true")))
	return err
}

// ExecInContainer 在容器内执行命令
func (c *Client) ExecInContainer(ctx context.Context, containerName string, cmd []string) error {
	// 获取容器
	cont, err := c.GetContainer(ctx, containerName)
	if err != nil || cont == nil {
		return fmt.Errorf("容器 %s 不存在", containerName)
	}

	// 创建 exec 实例
	execConfig := types.ExecConfig{
		Cmd:          cmd,
		AttachStdout: true,
		AttachStderr: true,
	}
	execResp, err := c.cli.ContainerExecCreate(ctx, cont.ID, execConfig)
	if err != nil {
		return err
	}

	// 执行命令
	if err := c.cli.ContainerExecStart(ctx, execResp.ID, types.ExecStartCheck{}); err != nil {
		return err
	}

	// 等待执行完成
	for {
		inspect, err := c.cli.ContainerExecInspect(ctx, execResp.ID)
		if err != nil {
			return err
		}
		if !inspect.Running {
			if inspect.ExitCode != 0 {
				return fmt.Errorf("命令执行失败，退出码: %d", inspect.ExitCode)
			}
			break
		}
		time.Sleep(100 * time.Millisecond)
	}

	return nil
}

// ExecOnHost 在主机上执行命令（通过临时容器）
func (c *Client) ExecOnHost(ctx context.Context, cmd string) error {
	// 使用可配置的 alpine 镜像，默认走官方镜像，不再绑定旧 Harbor
	alpineImage := os.Getenv("ALPINE_IMAGE")
	if alpineImage == "" {
		alpineImage = "alpine:3.21"
	}

	// 先拉取镜像
	if err := c.PullImage(ctx, alpineImage, ""); err != nil {
		return fmt.Errorf("拉取 alpine 镜像失败: %w", err)
	}

	cfg := &container.Config{
		Image: alpineImage,
		Cmd:   []string{"sh", "-c", cmd},
	}
	hcfg := &container.HostConfig{
		Binds:      []string{"/:/host"},
		AutoRemove: true,
	}

	// 修改命令，使其在 /host 目录下执行
	cfg.Cmd = []string{"sh", "-c", "cd /host && " + cmd}

	resp, err := c.cli.ContainerCreate(ctx, cfg, hcfg, nil, nil, "")
	if err != nil {
		return err
	}

	if err := c.cli.ContainerStart(ctx, resp.ID, types.ContainerStartOptions{}); err != nil {
		return err
	}

	// 等待容器完成
	statusCh, errCh := c.cli.ContainerWait(ctx, resp.ID, container.WaitConditionNotRunning)
	select {
	case err := <-errCh:
		if err != nil {
			return err
		}
	case status := <-statusCh:
		if status.StatusCode != 0 {
			return fmt.Errorf("命令执行失败，退出码: %d", status.StatusCode)
		}
	}

	return nil
}