wireguard-dashboard/script/script.go

package script

import (
	"encoding/json"
	"gitee.ltd/lxh/logger/log"
	"github.com/spf13/cast"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
	"os"
	"wireguard-ui/component"
	"wireguard-ui/global/client"
	"wireguard-ui/global/constant"
	"wireguard-ui/http/param"
	"wireguard-ui/model"
	"wireguard-ui/service"
	"wireguard-ui/template/render_data"
	"wireguard-ui/utils"
)

type script struct{}

func New() script {
	return script{}
}

func (s script) Do() error {
	if err := s.migrate(); err != nil {
		return err
	}

	if err := s.createSuperAdmin(); err != nil {
		return err
	}

	if err := s.initServer(); err != nil {
		log.Errorf("初始化wg服务端失败: %v", err.Error())
	}

	return nil
}

// migrate
// @description: 生成数据库
// @receiver s
// @return error
func (s script) migrate() error {
	var ent = []any{
		new(model.User),
		new(model.Client),
		new(model.Setting),
		new(model.RequestLog),
	}

	return client.DB.AutoMigrate(ent...)
}

// createSuperAdmin
// @description: 创建超级管理员
// @receiver s
// @return error
func (s script) createSuperAdmin() error {
	var count int64
	if err := client.DB.Model(&model.User{}).
		Where("account = ?", "admin").
		Where("is_admin = ?", 1).Count(&count).Error; err != nil {
		return err
	}

	// 没有超管就创建一个
	if count > 0 {
		return nil
	}

	// 生成一下头像
	avatarPath, err := utils.Avatar().GenerateAvatar(true)
	if err != nil {
		log.Errorf("生成头像失败: %v", err.Error())
		return err
	}

	return service.User().CreateUser(&model.User{
		Avatar:   avatarPath,
		Nickname: "超级管理员",
		Account:  "admin",
		Contact:  "",
		Password: "admin123",
		IsAdmin:  constant.SuperAdmin,
		Status:   constant.Enabled,
	})
}

// initServer
// @description: 初始化wg的一些配置
// @receiver s
// @return error
func (s script) initServer() error {
	var count int64
	if err := client.DB.Model(&model.Setting{}).Where("code = ?", "WG_SERVER").Count(&count).Error; err != nil {
		return err
	}

	if count > 0 {
		return nil
	}

	// 初始化服务端的全局配置
	var data = map[string]any{
		"endpointAddress":     utils.Network().GetHostPublicIP(),
		"dnsServer":           []string{"10.25.8.1"},
		"MTU":                 1450,
		"persistentKeepalive": 15,
		"firewallMark":        "0xca6c",
		"table":               "",
		"configFilePath":      "/etc/wireguard/wg0.conf",
	}
	dataJ, _ := json.Marshal(data)

	globalSet := &model.Setting{
		Code:     "WG_SETTING",
		Data:     string(dataJ),
		Describe: "服务端全局配置",
	}
	if err := service.Setting().SetData(globalSet); err != nil {
		return err
	}

	// 生成密钥
	privateKey, err := wgtypes.GeneratePrivateKey()
	if err != nil {
		log.Errorf("生成密钥失败: %v", err.Error())
		return err
	}

	// 根据密钥生成公钥
	publicKey := privateKey.PublicKey()
	serverEnt := &param.SaveServer{
		IPScope:        []string{"10.25.8.1/24"},
		ListenPort:     51820,
		PrivateKey:     privateKey.String(),
		PublicKey:      publicKey.String(),
		PostUpScript:   constant.DefaultPostUpScript,
		PreDownScript:  constant.DefaultPreDownScript,
		PostDownScript: constant.DefaultPostDownScript,
	}

	serverJ, _ := json.Marshal(serverEnt)
	serverSet := &model.Setting{
		Code:     "WG_SERVER",
		Data:     string(serverJ),
		Describe: "服务端配置",
	}

	// 没有服务端，开始初始化
	if err = service.Setting().SetData(serverSet); err != nil {
		return err
	}

	// 处理一下要渲染到配置文件上的数据
	serverConfig := render_data.Server{
		Address:    serverEnt.IPScope,
		ListenPort: serverEnt.ListenPort,
		PrivateKey: serverEnt.PrivateKey,
		MTU:        cast.ToInt(data["MTU"]),
		PostUp:     serverEnt.PostUpScript,
		PreDown:    serverEnt.PreDownScript,
		PostDown:   serverEnt.PostDownScript,
		Table:      cast.ToString(data["table"]),
	}

	execData := map[string]any{
		"Server": serverConfig,
	}

	var templatePath, outFilePath string
	if os.Getenv("GIN_MODE") != "release" {
		templatePath = "E:\\Workspace\\Go\\wireguard-ui\\template\\conf\\wg.conf"
		outFilePath = "E:\\Workspace\\Go\\wireguard-ui\\template\\tmp\\wg0.conf"
	} else {
		templatePath = "./template/wg.conf"
		outFilePath = cast.ToString(data["configFilePath"])
	}

	// 先渲染模板
	if err = component.Template().Execute(templatePath, outFilePath, execData); err != nil {
		return nil
	}

	// 模板渲染成功，开始执行服务端控制
	component.Wireguard().ServerControl(outFilePath)
	return nil
}