forge/internal/cmd/make.es.go

package cmd

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"

	"gitea.loveuer.com/yizhisec/pkg3/logger"
	"github.com/spf13/cobra"
	"yizhisec.com/hsv2/forge/internal/controller/maker"
	"yizhisec.com/hsv2/forge/internal/opt"
	"yizhisec.com/hsv2/forge/pkg/archiver"
	"yizhisec.com/hsv2/forge/pkg/resource"
)

func makeES() *cobra.Command {
	const (
		dockerFile = `
FROM docker-mirror.yizhisec.com/library/alpine:3.22.2

WORKDIR /data

RUN sed -i 's/dl-cdn.alpinelinux.org/mirrors.tuna.tsinghua.edu.cn/g' /etc/apk/repositories && apk add curl

ENV TZ=Asia/Shanghai

COPY plugins /data/plugins
COPY create_index.sh /data/create_index.sh
RUN chmod +x /data/create_index.sh
`
	)

	var (
		makeHelper bool
		memRate    int
		storage    int
	)

	_cmd := &cobra.Command{
		Use:     "es",
		Aliases: []string{"elastic", "elasticsearch"},
		RunE: func(cmd *cobra.Command, args []string) error {
			var (
				err               error
				output            []byte
				location          = filepath.Join(opt.Cfg.Make.Dir, "dependency", "elastic")
				helperTarLocation = filepath.Join(opt.Cfg.Make.Dir, "dependency", "image", "es-init-helper.alpine-3.22.2.tar")
			)

			logger.Info("MakeES: 开始构建 elasticsearch(es) 依赖...")

			if err := os.MkdirAll(location, 0755); err != nil {
				logger.Error("MakeES: 创建目录失败")
				logger.Debug("MakeES: 创建目录失败: %v", err)
				return err
			}

			if makeHelper {
				tmp := filepath.Join(os.TempDir(), "make-es")
				logger.Info("MakeES: 重新构建 helper 镜像...(%s)", tmp)

				if err = os.RemoveAll(tmp); err != nil {
					logger.Error("MakeES: 删除临时目录失败")
					logger.Debug("MakeES: 删除临时目录失败: %v", err)
					return err
				}

				if err = os.MkdirAll(tmp, 0755); err != nil {
					logger.Error("MakeES: 创建临时目录失败")
					logger.Debug("MakeES: 创建临时目录失败: %v", err)
					return err
				}

				if err = os.WriteFile(filepath.Join(tmp, "Dockerfile"), []byte(dockerFile), 0644); err != nil {
					logger.Error("MakeES: 写入 Dockerfile 失败")
					logger.Debug("MakeES: 写入 Dockerfile 失败: %v", err)
					return err
				}

				if err = archiver.DownloadAndExtract(
					cmd.Context(),
					"https://artifactory.yizhisec.com:443/artifactory/filestore/hsv3/dependency/elasticsearch/v2-es-plugins.tar.gz",
					tmp,
					archiver.WithInsecureSkipVerify(),
				); err != nil {
					logger.Error("MakeES: 下载 es 插件失败")
					logger.Debug("MakeES: 下载 es 插件失败: %v", err)
					return err
				}

				if err = os.WriteFile(filepath.Join(tmp, "create_index.sh"), resource.BashESInit, 0644); err != nil {
					logger.Error("MakeES: 写入 create_index.sh 文件失败")
					logger.Debug("MakeES: 写入 create_index.sh 文件失败: %v", err)
					return err
				}

				helpCmd := exec.CommandContext(cmd.Context(), "docker", "build", "-t",
					"hub.yizhisec.com/hybridscope/v2/es-init-helper:alpine-3.22.2",
					"-f", filepath.Join(tmp, "Dockerfile"),
					tmp,
				)

				if output, err = helpCmd.CombinedOutput(); err != nil {
					logger.Error("MakeES: 重新构建 helper 镜像失败")
					logger.Debug("MakeES: 重新构建 helper 镜像失败: %v", err)
					logger.Debug("MakeES: 重新构建 helper 镜像失败: %s", string(output))
					return err
				}

				logger.Info("MakeES: 重新构建 helper 镜像成功!!!")
			}

			mk := maker.NewMaker()
			if err = mk.Image(
				cmd.Context(),
				"hub.yizhisec.com/hybridscope/v2/es-init-helper:alpine-3.22.2",
				maker.WithImageSave(helperTarLocation),
			); err != nil {
				logger.Error("MakeES: 准备 es init helper 镜像失败")
				logger.Debug("MakeES: 准备 es init helper 镜像失败: %v", err)
				return err
			}

			bs := []byte(fmt.Sprintf(resource.YAMLES,
				memRate, memRate*2,
				memRate, memRate,
				memRate*2, memRate*2,
				storage,
			))

			if err = os.WriteFile(filepath.Join(opt.Cfg.Make.Dir, "dependency", "elastic", "es.yaml"), bs, 0644); err != nil {
				logger.Error("MakeES: 写入 es.yaml 文件失败")
				logger.Debug("MakeES: 写入 es.yaml 文件失败: %v", err)
				return err
			}

			if err = os.WriteFile(filepath.Join(opt.Cfg.Make.Dir, "dependency", "elastic", "kibana.yaml"), resource.YAMLKibana, 0644); err != nil {
				logger.Error("MakeES: 写入 es.yaml 文件失败")
				logger.Debug("MakeES: 写入 es.yaml 文件失败: %v", err)
				return err
			}

			logger.Info("MakeES: 构建 elastic(es) 依赖成功!!!")
			return nil
		},
	}

	_cmd.Flags().BoolVar(&makeHelper, "make-helper", false, "重新构建 helper 镜像")
	_cmd.Flags().IntVar(&memRate, "mem-rate", 1, "内存倍率(从 1G 开始)")
	_cmd.Flags().IntVar(&storage, "storage", 100, "存储空间(单位: G)")

	return _cmd
}