wireguard-dashboard/component/wireguard.go

package component

import (
	"encoding/json"
	"errors"
	"fmt"
	"gitee.ltd/lxh/logger/log"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
	"gopkg.in/fsnotify/fsnotify.v1"
	"os"
	"strings"
	"time"
	"wireguard-ui/command"
	"wireguard-ui/config"
	"wireguard-ui/global/client"
	"wireguard-ui/model"
	"wireguard-ui/service"
	"wireguard-ui/template/render_data"
)

type WireguardComponent struct{}

func Wireguard() WireguardComponent {
	return WireguardComponent{}
}

// GetClients
// @description: 获取所有链接的客户端信息
// @receiver w
// @return peers
// @return err
func (w WireguardComponent) GetClients() (peers []wgtypes.Peer, err error) {
	device, err := client.WireguardClient.Devices()
	if err != nil {
		return
	}

	for _, v := range device {
		return v.Peers, nil
	}

	return
}

// GetClientByPublicKey
// @description: 根据公钥获取指定客户端信息
// @receiver w
// @return peer
// @return err
func (w WireguardComponent) GetClientByPublicKey(pk string) (peer *wgtypes.Peer, err error) {
	peers, err := w.GetClients()
	if err != nil {
		return
	}

	for _, v := range peers {
		if v.PublicKey.String() == pk {
			return &v, nil
		}
	}

	return
}

// GenerateClientFile
// @description: 生成客户端文件
// @receiver w
// @param clientInfo
// @param server
// @param setting
// @return filePath
// @return err
func (w WireguardComponent) GenerateClientFile(clientInfo *model.Client, server *render_data.Server, setting *render_data.ServerSetting) (filePath string, err error) {
	var keys render_data.Keys
	_ = json.Unmarshal([]byte(clientInfo.Keys), &keys)

	var serverDNS []string
	if clientInfo.UseServerDns == 1 {
		serverDNS = setting.DnsServer
	}

	// 处理一下数据
	execData := render_data.ClientConfig{
		PrivateKey:          keys.PrivateKey,
		IpAllocation:        clientInfo.IpAllocation,
		MTU:                 setting.MTU,
		DNS:                 strings.Join(serverDNS, ","),
		PublicKey:           server.PublicKey,
		PresharedKey:        keys.PresharedKey,
		AllowedIPS:          clientInfo.AllowedIps,
		Endpoint:            setting.EndpointAddress,
		ListenPort:          int(server.ListenPort),
		PersistentKeepalive: setting.PersistentKeepalive,
	}

	// 不同环境下处理文件路径
	var outPath = "/tmp/" + fmt.Sprintf("%s.conf", clientInfo.Name)
	var templatePath = "./template/wg.client.conf"
	if os.Getenv("GIN_MODE") != "release" {
		outPath = "E:\\Workspace\\Go\\wireguard-ui\\template\\tmp\\" + fmt.Sprintf("%s.conf", clientInfo.Name)
		templatePath = "E:\\Workspace\\Go\\wireguard-ui\\template\\conf\\wg.client.conf"
	}

	err = Template().Execute(templatePath, outPath, execData)
	if err != nil {
		return "", errors.New("文件渲染失败")
	}

	return outPath, nil
}

// ServerControl
// @description: 服务端控制
// @receiver w
// @return error
func (w WireguardComponent) ServerControl(filePath string) {
	if filePath == "" {
		data, err := service.Setting().GetWGSetForConfig()
		if err != nil {
			log.Errorf("获取服务端配置失败: %v", err.Error())
			return
		}
		filePath = data.ConfigFilePath
	}
	w.watchConfigFile(filePath, config.Config.Wireguard.RestartMode, config.Config.Wireguard.DelayTime)
}

// watchConfigFile
// @description: 监听并重新操作配置文件
// @receiver w
// @param filepath
func (w WireguardComponent) watchConfigFile(filepath string, mode string, delay int64) {
	go func() {
		watcher, err := fsnotify.NewWatcher()
		if err != nil {
			log.Errorf("创建文件监控失败: %v", err.Error())
			return
		}

		defer watcher.Close()
		done := make(chan bool)

		go func() {
			for {
				select {
				case event, ok := <-watcher.Events:
					if !ok {
						return
					}

					if event.Op == fsnotify.Write {
						switch mode {
						case "NOW":
							command.RestartWireguard(false, filepath)
						case "DELAY":
							time.Sleep(time.Duration(delay) * time.Second)
							command.RestartWireguard(true, filepath)
						}

					}

					// 打印监听事件
					log.Infof("监听事件是：%s", event.String())
				case _, ok := <-watcher.Errors:
					if !ok {
						return
					}
				}
			}
		}()

		if err = watcher.Add(filepath); err != nil {
			log.Errorf("添加[%s]监听失败: %v", filepath, err.Error())
			return
		}
		<-done
	}()
}