package model import ( "io/ioutil" "os" "strconv" "strings" "github.com/spf13/viper" "gopkg.in/yaml.v2" ) const ( ConfigTypeGitHub = "github" ConfigTypeGitee = "gitee" ) const ( ConfigCoverAll = iota ConfigCoverIgnoreAll ) type AgentConfig struct { HardDrivePartitionAllowlist []string NICAllowlist map[string]bool v *viper.Viper } // Read 从给定的文件目录加载配置文件 func (c *AgentConfig) Read(path string) error { c.v = viper.New() c.v.SetConfigFile(path) err := c.v.ReadInConfig() if err != nil { return err } err = c.v.Unmarshal(c) if err != nil { return err } return nil } func (c *AgentConfig) Save() error { data, err := yaml.Marshal(c) if err != nil { return err } return ioutil.WriteFile(c.v.ConfigFileUsed(), data, os.ModePerm) } // Config 站点配置 type Config struct { Debug bool // debug模式开关 Site struct { Brand string // 站点名称 CookieName string // 浏览器 Cookie 名称 Theme string CustomCode string ViewPassword string // 前台查看密码 } Oauth2 struct { Type string Admin string // 管理员用户名列表 ClientID string ClientSecret string } HTTPPort uint GRPCPort uint GRPCHost string ProxyGRPCPort uint TLS bool EnableIPChangeNotification bool IPChangeNotificationTag string EnablePlainIPInNotification bool // IP变更提醒 Cover uint8 // 覆盖范围(0:提醒未被 IgnoredIPNotification 包含的所有服务器; 1:仅提醒被 IgnoredIPNotification 包含的服务器;) IgnoredIPNotification string // 特定服务器IP(多个服务器用逗号分隔) v *viper.Viper IgnoredIPNotificationServerIDs map[uint64]bool // [ServerID] -> bool(值为true代表当前ServerID在特定服务器列表内) } // Read 读取配置文件并应用 func (c *Config) Read(path string) error { c.v = viper.New() c.v.SetConfigFile(path) err := c.v.ReadInConfig() if err != nil { return err } err = c.v.Unmarshal(c) if err != nil { return err } if c.Site.Theme == "" { c.Site.Theme = "default" } if c.GRPCPort == 0 { c.GRPCPort = 5555 } if c.EnableIPChangeNotification && c.IPChangeNotificationTag == "" { c.IPChangeNotificationTag = "default" } c.updateIgnoredIPNotificationID() return nil } // updateIgnoredIPNotificationID 更新用于判断服务器ID是否属于特定服务器的map func (c *Config) updateIgnoredIPNotificationID() { c.IgnoredIPNotificationServerIDs = make(map[uint64]bool) splitedIDs := strings.Split(c.IgnoredIPNotification, ",") for i := 0; i < len(splitedIDs); i++ { id, _ := strconv.ParseUint(splitedIDs[i], 10, 64) if id > 0 { c.IgnoredIPNotificationServerIDs[id] = true } } } // Save 保存配置文件 func (c *Config) Save() error { c.updateIgnoredIPNotificationID() data, err := yaml.Marshal(c) if err != nil { return err } return ioutil.WriteFile(c.v.ConfigFileUsed(), data, os.ModePerm) }