ProjectOctopus/agent-wdd/config/Config.go

package config

import (
	"agent-wdd/log"
	"agent-wdd/utils"
	"fmt"
	"os"
	"os/exec"
	"runtime"
	"strings"

	"github.com/spf13/viper"
	"gopkg.in/yaml.v3"
)

var WddConfigFilePath = "/usr/local/etc/wdd/agent-wdd-config.yaml"

// ConfigCache 配置缓存
var ConfigCache = &Config{
	TimeStamp:     "",
	ModifiedTimes: 0,
	Agent: Agent{
		OS:      OS{},
		Network: Network{},
		CPU:     CPU{},
		Mem:     Memory{},
		Swap:    Swap{},
		Disks:   []Disk{},
	},
}

func init() {
	// 根据运行的操作系统不同, 修改 WddConfigFilePath 的位置
	if runtime.GOOS == "windows" {
		homedir, _ := os.UserHomeDir()
		WddConfigFilePath = homedir + "\\agent-wdd\\agent-wdd-config.yaml"
	}
}

type Config struct {
	TimeStamp     string `yaml:"timestamp"`
	ModifiedTimes int    `yaml:"modifiedTimes"`
	Agent         Agent  `yaml:"agent"`
}

type Agent struct {
	OS      OS      `yaml:"os"`
	Network Network `yaml:"network"`
	CPU     CPU     `yaml:"cpu"`
	Mem     Memory  `yaml:"mem"`
	Swap    Swap    `yaml:"swap"`
	Disks   []Disk  `yaml:"disks"`
}

type OS struct {
	Hostname      string `mapstructure:"hostname" yaml:"hostname" comment:"主机名"`
	OsName        string `mapstructure:"os_name" yaml:"os_name" comment:"操作系统名称"`
	OsFamily      string `mapstructure:"os_family" yaml:"os_family" comment:"操作系统家族"`
	OsVersion     string `mapstructure:"os_version" yaml:"os_version" comment:"操作系统版本"`
	OsType        string `mapstructure:"os_type" yaml:"os_type" comment:"操作系统类型"`
	Kernel        string `mapstructure:"kernel" yaml:"kernel" comment:"内核版本"`
	Arch          string `mapstructure:"arch" yaml:"arch" comment:"架构"`
	IsUbuntuType  bool   `mapstructure:"is_ubuntu_type" yaml:"is_ubuntu_type" comment:"是否是ubuntu类型的操作系统"`
	PackInit      bool   `mapstructure:"pack_init" yaml:"pack_init" comment:"是否初始化，ubuntu需要"`
	OSReleaseCode string `mapstructure:"os_release_code" yaml:"os_release_code" comment:"主机操作系统的发行版代号， focal之类的"`
}

type Network struct {
	Internet   int         `yaml:"internet" comment:"是否能够上网 外网为9 国内为7 不能上网为1"`
	Public     PublicInfo  `yaml:"public"`
	Interfaces []Interface `yaml:"interfaces"`
}

type PublicInfo struct {
	IPv4     string `yaml:"ipv4"`
	IPv6     string `yaml:"ipv6"`
	Country  string `yaml:"country"`
	City     string `yaml:"city"`
	ASN      string `yaml:"asn"`
	Timezone string `yaml:"timezone"`
}

type Interface struct {
	Name string `yaml:"name"`
	IPv4 string `yaml:"ipv4"`
	IPv6 string `yaml:"ipv6"`
	MAC  string `yaml:"mac"`
	MTU  int    `yaml:"mtu"`
}

type CPU struct {
	Cores      int    `yaml:"cores"`
	Brand      string `yaml:"brand"`
	Mhz        string `yaml:"mhz"`
	Arch       string `yaml:"arch"`
	Virt       string `yaml:"virt"`
	Hypervisor string `yaml:"hypervisor"`
}

type Memory struct {
	Size  string `yaml:"size"`
	Type  string `yaml:"type"`
	Speed int    `yaml:"speed"`
}

type Swap struct {
	Open bool   `yaml:"open"`
	Size string `yaml:"size"`
}

type Disk struct {
	Path    string `yaml:"path"`
	Size    string `yaml:"size"`
	Usage   string `yaml:"usage"`
	Percent string `yaml:"percent"`
}

func InitConfig() {

	//  检查配置文件是否存在
	if !utils.FileOrFolderExists(WddConfigFilePath) {
		utils.AppendContentToFile("", WddConfigFilePath)
	}

	v := viper.New()
	v.SetConfigFile(WddConfigFilePath)
	v.SetConfigType("yaml")

	if err := v.ReadInConfig(); err != nil {
		log.Error("读取配置文件失败: %w", err)
		panic(err)

	}

	if err := v.Unmarshal(&ConfigCache); err != nil {
		log.Error("解析配置失败: %w", err)
		panic(err)
	}

}

// 写入配置文件
func SaveConfig() {

	// 每次写入新的时间
	ConfigCache.TimeStamp = utils.CurrentTimeString()

	// 每次增加修改文件的次数计数
	ConfigCache.ModifiedTimes += 1

	data, err := yaml.Marshal(&ConfigCache)
	if err != nil {
		log.Error("YAML序列化失败: %w", err)
	}

	if err := os.WriteFile(WddConfigFilePath, data, 0644); err != nil {
		log.Error("写入文件失败: %w", err)
	}
}

// 归一化配置-重命名主机名
func (c *Config) NormalizeConfig() {
	// 重命名主机名

	log.Info("归一化主机配置")

	// 重新读取配置
	InitConfig()

	// 主机名称应该为 City(格式为首字母大写)-amd64-内网IP的最后一段(格式化为2位)-公网IPv4(如果公网IPv4为空，则使用内网IPv4; ip的格式为127-0-0-1)

	// 获取城市(格式为首字母大写)
	city := strings.Title(ConfigCache.Agent.Network.Public.City)

	// 获取公网IPv4 ip的格式为127-0-0-1
	ipInfo := ConfigCache.Agent.Network.Public.IPv4
	if ipInfo == "" {
		ipInfo = ConfigCache.Agent.Network.Interfaces[0].IPv4
	} else {
		// 可以获取到公网IPv4, 修改IP的格式
		innerIpv4 := ConfigCache.Agent.Network.Interfaces[0].IPv4
		ipSegments := strings.Split(innerIpv4, ".")
		prefix := fmt.Sprintf("%02s", ipSegments[len(ipSegments)-1])
		ipInfo = prefix + "." + ipInfo
	}

	ipInfo = strings.ReplaceAll(ipInfo, ".", "-")

	// 获取架构
	arch := ConfigCache.Agent.CPU.Arch

	uniformHostname := city + "-" + arch + "-" + ipInfo

	// 重命名主机名
	log.Info("重命名主机名: %s", uniformHostname)
	cmd := exec.Command("hostnamectl", "set-hostname", uniformHostname)

	// 执行命令
	cmd.Run()

	// 更新配置
	ConfigCache.Agent.OS.Hostname = uniformHostname

	// 更新配置
	SaveConfig()
}