nginx-ui/agent/handler.go

package agent

import (
	"errors"
	"fmt"
	"github.com/labstack/gommon/log"
	"io"
	"net/http"
	"nginx-ui/server/constants"
	"nginx-ui/server/models"
	nginx2 "nginx-ui/server/nginx"
	"nginx-ui/server/utils"
	"os"
	"strings"
	"time"
)

var OnAgentMessageHandler func(agent *Agent, message *models.AgentData) (interface{}, error)

// OnNginxUpdated 不管是Agent注册，还是服务端修改Nginx信息，均会调用该方法
func OnNginxUpdated(c *Agent, data *models.AgentData) (interface{}, error) {
	nginx := &models.Nginx{}
	err := data.ReadData(nginx)
	if err != nil {
		return nil, err
	}
	c.Nginx = nginx
	return make(map[string]interface{}), nil

}

func OnServerConnected(c *Agent, data *models.AgentData) (interface{}, error) {
	fmt.Printf("OnServerConnected: %v,%v\n", c.Token, data)
	if c.Token == "" {
		log.Warn("missing token")
		return nil, nil
	}
	nginx := &models.Nginx{
		Name:     "Agent",
		IsServer: false,
		Proxy:    true,
		IsLocal:  true,
		Password: "",
		Remark:   "agent",
		IpAddr:   utils.GetNetIP(),
	}
	request := &models.AgentData{
		Type: models.RegisterNginxType,
		Data: nginx,
	}
	res, err := c.Send(request, 20*time.Second)
	if err != nil {
		log.Warn("register fail for: %v\n", err)
		return nil, nil
	}
	if res != nil {
		err = res.ReadData(nginx)
		if err != nil {
			log.Warn("register fail: %v\n", err)
			return nil, nil
		}
		c.Nginx = nginx
	}
	return nil, nil
}

// OnCMD 执行Nginx命令
func OnCMD(c *Agent, data *models.AgentData) (interface{}, error) {
	nginx := c.Nginx
	if nginx == nil {
		go func() {
			_, _ = OnServerConnected(c, data)
		}()
		return nil, errors.New("nginx agent not register")
	}
	request := &models.AgentCMD{}
	err := data.ReadData(request)
	if err != nil {
		return nil, err
	}

	instance := nginx2.NewLocalNginx(nginx)
	res, err := instance.Run(request.Cmd)
	return res, err
}

// OnSendFile 服务端复制文件到客户端
// 先使用Http接口下载文件，然后本地解压复制
func OnSendFile(c *Agent, data *models.AgentData) (interface{}, error) {
	nginx := c.Nginx
	if nginx == nil {
		go func() {
			_, _ = OnServerConnected(c, data)
		}()
		return nil, errors.New("nginx agent not register")
	}
	request := &models.AgentSendFile{}
	err := data.ReadData(request)
	if err != nil {
		log.Warn("send file fail for: %v\n", nginx)
		return nil, err
	}

	protocol := "http"
	if c.SSL == "Y" {
		protocol = "https"
	}
	url := fmt.Sprintf("%v://%v%v", protocol, c.Url, constants.DownloadFileUrl)
	url = strings.Replace(url, ":nginxId", fmt.Sprintf("%d", nginx.Id), 1)
	url = strings.Replace(url, ":fileName", request.FileName, 1)
	log.Printf("download file: %s\n", url)

	out, err := os.Create(request.Dst)
	if err != nil {
		log.Warn("create file fail for: %v\n", nginx)
		return nil, err
	}
	defer out.Close()

	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		log.Warn("create file fail for: %v\n", nginx)
		return nil, err
	}
	req.Header.Set("Token", c.Token)
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		log.Warn("download file fail for: %v\n", nginx)
		return nil, err
	}
	defer resp.Body.Close()
	_, err = io.Copy(out, resp.Body)
	if err != nil {
		log.Warn("download file fail for: %v\n", nginx)
		return nil, err
	}
	return make(map[string]interface{}), nil
}