qq_bot/handler/restart/restart.go

package restart

import (
	"errors"
	"log"
	"net/http"
	"os"
	"os/exec"
	"sync"
	"time"

	"git.lxtend.com/qqbot/config"
	"git.lxtend.com/qqbot/constants"
	"git.lxtend.com/qqbot/handler"
	"git.lxtend.com/qqbot/model"
	docker "git.lxtend.com/qqbot/service/exec"
	"git.lxtend.com/qqbot/util"
	"github.com/gin-gonic/gin"
)

var hasVaildBuild = true

// 避免被并发请求构建或重启
var stuffMutex sync.Mutex

func init() {
	handler.RegisterHandler("/重启bot", restart, constants.LEVEL_ADMIN)
	handler.RegisterHelpInform("/重启bot", "热更新", "重启bot")
	handler.RegisterHandler("/构建bot", build, constants.LEVEL_ADMIN)
	handler.RegisterHelpInform("/构建bot", "热更新", "构建bot")
	handler.RegisterHandler("/构建重启", buildAndRestart, constants.LEVEL_ADMIN)
	handler.RegisterHelpInform("/构建重启", "热更新", "构建并重启")
	handler.RegisterHandler("/gitpull", pullCode, constants.LEVEL_USER)
	handler.RegisterHelpInform("/gitpull", "热更新", "拉取最新代码")
	handler.RegisterHandler("/shutdown", shutdown, constants.LEVEL_ADMIN)
	handler.RegisterHelpInform("/shutdown", "热更新", "结束程序")
}

func restart(msg model.Message) *model.Reply {
	if !hasVaildBuild {
		return &model.Reply{
			ReplyMsg:       "上次构建失败，请先成功构建再部署",
			ReferOriginMsg: true,
			FromMsg:        msg,
		}
	}
	go restartProgram()
	return &model.Reply{
		ReplyMsg:       "重启中...",
		ReferOriginMsg: true,
		FromMsg:        msg,
	}
}

func shutdown(msg model.Message) *model.Reply {
	go shutdownProgram()
	return &model.Reply{
		ReplyMsg:       "关闭中...",
		ReferOriginMsg: true,
		FromMsg:        msg,
	}
}

func build(msg model.Message) *model.Reply {
	err := buildBot()
	if err != nil {
		return &model.Reply{
			ReplyMsg:       "构建失败,报错如下\n" + err.Error(),
			ReferOriginMsg: true,
			FromMsg:        msg,
		}
	}
	return &model.Reply{
		ReplyMsg:       "构建成功",
		ReferOriginMsg: true,
		FromMsg:        msg,
	}
}

func buildAndRestart(msg model.Message) *model.Reply {
	err := buildBot()
	if err != nil {
		return &model.Reply{
			ReplyMsg:       "构建失败,报错如下\n" + err.Error(),
			ReferOriginMsg: true,
			FromMsg:        msg,
		}
	}
	if !hasVaildBuild {
		return &model.Reply{
			ReplyMsg:       "构建失败，请先成功构建再部署",
			ReferOriginMsg: true,
			FromMsg:        msg,
		}
	}
	go restartProgram()
	return &model.Reply{
		ReplyMsg:       "构建成功,重启中...",
		ReferOriginMsg: true,
		FromMsg:        msg,
	}
}

func pullCode(msg model.Message) *model.Reply {
	err := util.GitPull()
	if err != nil {
		return &model.Reply{
			ReplyMsg:       "拉取代码失败,报错如下\n" + err.Error(),
			ReferOriginMsg: true,
			FromMsg:        msg,
		}
	}
	return &model.Reply{
		ReplyMsg:       "拉取代码成功",
		ReferOriginMsg: true,
		FromMsg:        msg,
	}
}

func restartProgram() error {
	log.Println("重启程序...")
	docker.DockerContainer.RemoveContainer()
	time.Sleep(500 * time.Millisecond)
	path, err := os.Executable()
	if err != nil {
		return err
	}
	workDir, err := os.Getwd()
	if err != nil {
		return err
	}
	cmd := exec.Command(path)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Dir = workDir
	if err := cmd.Start(); err != nil {
		return err
	}
	os.Exit(0)
	return nil
}

func shutdownProgram() error {
	log.Println("关闭程序...")
	docker.DockerContainer.RemoveContainer()
	time.Sleep(500 * time.Millisecond)
	os.Exit(0)
	return nil
}

func buildBot() error {
	hasVaildBuild = false
	workDir, err := os.Getwd()
	if err != nil {
		return err
	}
	cmd := exec.Command("go", "mod", "tidy")
	cmd.Dir = workDir
	output, err := cmd.CombinedOutput()
	if err != nil {
		return errors.New(string(output) + err.Error())
	}

	cmd = exec.Command("go", "build", "-o", "qqbot")
	cmd.Dir = workDir

	output, err = cmd.CombinedOutput()
	if err != nil {
		return errors.New(string(output) + err.Error())
	}
	hasVaildBuild = true
	return nil
}

func PullCodeHandler(c *gin.Context) {
	if !stuffMutex.TryLock() {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "有其他请求正在运行，请稍后再试"})
		return
	}
	defer stuffMutex.Unlock()
	err := util.GitPull()
	log.Println("拉取代码...")
	if err != nil {
		log.Println(err)
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.Status(http.StatusOK)
	// c.JSON(http.StatusOK, gin.H{"message": "拉取代码成功"})
}

func BuildBotHandler(c *gin.Context) {
	if !stuffMutex.TryLock() {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "有其他请求正在运行，请稍后再试"})
		return
	}
	defer stuffMutex.Unlock()
	err := buildBot()
	log.Println("构建程序...")
	if err != nil {
		log.Println(err)
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.Status(http.StatusOK)
	// c.JSON(http.StatusOK, gin.H{"message": "构建成功"})
}

func RestartBotHandler(c *gin.Context) {
	token := c.Request.Header.Get("Authorization")
	if token != config.ConfigManager.GetConfig().WebServer.Token {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"})
		return
	}
	if !stuffMutex.TryLock() {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "有其他请求正在运行，请稍后再试"})
		return
	}
	defer stuffMutex.Unlock()
	err := restartProgram()
	log.Println("重启程序...")
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
	}
	c.Status(http.StatusOK)
	// c.JSON(http.StatusOK, gin.H{"message": "重启成功"})
}

func AllInOneHandler(c *gin.Context) {
	PullCodeHandler(c)
	BuildBotHandler(c)
	RestartBotHandler(c)
	if c.Writer.Status() != http.StatusOK {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "unauthorized"})
		return
	}
	c.JSON(http.StatusOK, gin.H{"message": "全部操作成功"})
}