refactor: rewrite cfg.go
This commit is contained in:
		
							parent
							
								
									b872b73ebf
								
							
						
					
					
						commit
						6971098753
					
				
							
								
								
									
										4
									
								
								bot.go
								
								
								
								
							
							
						
						
									
										4
									
								
								bot.go
								
								
								
								
							| 
						 | 
				
			
			@ -24,7 +24,7 @@ func isFromAdmin(sender *tele.User) bool {
 | 
			
		|||
		return false
 | 
			
		||||
	}
 | 
			
		||||
 | 
			
		||||
	_, ok := config.AdminUIDs[sender.ID]
 | 
			
		||||
	_, ok := config.adminUidLookup[sender.ID]
 | 
			
		||||
	return ok
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
| 
						 | 
				
			
			@ -61,7 +61,7 @@ func initBot() (*tele.Bot, error) {
 | 
			
		|||
 | 
			
		||||
	adminGrp.Handle("/dig", handleDigCmd)
 | 
			
		||||
 | 
			
		||||
	// adminGrp.Handle("/test", handleTestCmd)
 | 
			
		||||
	// adminGrp.Handle("/test", testCmd)
 | 
			
		||||
 | 
			
		||||
	return b, nil
 | 
			
		||||
}
 | 
			
		||||
| 
						 | 
				
			
			
 | 
			
		|||
							
								
								
									
										122
									
								
								cfg.go
								
								
								
								
							
							
						
						
									
										122
									
								
								cfg.go
								
								
								
								
							| 
						 | 
				
			
			@ -1,66 +1,94 @@
 | 
			
		|||
package main
 | 
			
		||||
 | 
			
		||||
import (
 | 
			
		||||
	"os"
 | 
			
		||||
	"net"
 | 
			
		||||
	"strconv"
 | 
			
		||||
	"strings"
 | 
			
		||||
 | 
			
		||||
	"git.gensokyo.cafe/kkyy/tgbot_misaka_5882f7/utils"
 | 
			
		||||
	"github.com/go-errors/errors"
 | 
			
		||||
	"github.com/samber/lo"
 | 
			
		||||
	"github.com/ilyakaznacheev/cleanenv"
 | 
			
		||||
)
 | 
			
		||||
 | 
			
		||||
type Config struct {
 | 
			
		||||
	AdminUIDs          map[int64]struct{}
 | 
			
		||||
	TGBotToken         string
 | 
			
		||||
	TGAnnounceCommands bool
 | 
			
		||||
type ConfigDef struct {
 | 
			
		||||
	AdminUIDs          []int64 `env:"TG_ADMIN_UIDS"`
 | 
			
		||||
	TGBotToken         string  `env:"TG_TOKEN" env-required:""`
 | 
			
		||||
	TGAnnounceCommands bool    `env:"TG_ANNOUNCE_CMDS"`
 | 
			
		||||
 | 
			
		||||
	WatchedInterface       string
 | 
			
		||||
	MonthlyTrafficLimitGiB int
 | 
			
		||||
	// API
 | 
			
		||||
	APIMasterKey string `env:"TG_API_MASTER_KEY" env-required:""`
 | 
			
		||||
	APIListen    string `env:"TG_API_LISTEN" env-default:"127.0.0.1:8080"`
 | 
			
		||||
 | 
			
		||||
	// Traffic usage
 | 
			
		||||
	WatchedInterface       string `env:"TG_WATCHED_INTERFACE"`
 | 
			
		||||
	MonthlyTrafficLimitGiB int    `env:"TG_MONTHLY_TRAFFIC_LIMIT_GIB" env-default:"1000"`
 | 
			
		||||
 | 
			
		||||
	// Parsed fields
 | 
			
		||||
	adminUidLookup map[int64]struct{}
 | 
			
		||||
	apiListenAddr  net.Addr
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
var config *Config
 | 
			
		||||
 | 
			
		||||
func LoadCfg() error {
 | 
			
		||||
	cfg := Config{}
 | 
			
		||||
 | 
			
		||||
	token := os.Getenv("TG_TOKEN")
 | 
			
		||||
	if token == "" {
 | 
			
		||||
		return errors.New("TG_TOKEN env not set")
 | 
			
		||||
func (c *ConfigDef) Parse() error {
 | 
			
		||||
	c.adminUidLookup = make(map[int64]struct{}, len(c.AdminUIDs))
 | 
			
		||||
	for _, u := range c.AdminUIDs {
 | 
			
		||||
		if u <= 0 {
 | 
			
		||||
			return errors.New("invalid admin UID: " + strconv.FormatInt(u, 10))
 | 
			
		||||
		}
 | 
			
		||||
		c.adminUidLookup[u] = struct{}{}
 | 
			
		||||
	}
 | 
			
		||||
	cfg.TGBotToken = token
 | 
			
		||||
 | 
			
		||||
	adminUIDsEnv := os.Getenv("TG_ADMIN_UIDS")
 | 
			
		||||
	adminUIDs := lo.FilterMap(strings.Split(adminUIDsEnv, ","), func(s string, _ int) (string, bool) {
 | 
			
		||||
		trimmed := strings.TrimSpace(s)
 | 
			
		||||
		return trimmed, trimmed != ""
 | 
			
		||||
	})
 | 
			
		||||
	cfg.AdminUIDs = make(map[int64]struct{}, len(adminUIDs))
 | 
			
		||||
	for _, uidStr := range adminUIDs {
 | 
			
		||||
		uid, err := strconv.ParseInt(uidStr, 10, 64)
 | 
			
		||||
	// validate API listen
 | 
			
		||||
	ip, port, err := net.SplitHostPort(c.APIListen)
 | 
			
		||||
	if err != nil {
 | 
			
		||||
		return errors.WrapPrefix(err, "invalid API listen address", 0)
 | 
			
		||||
	}
 | 
			
		||||
	if tIP := net.ParseIP(ip); tIP == nil {
 | 
			
		||||
		return errors.New("invalid API listen address: " + ip)
 | 
			
		||||
	}
 | 
			
		||||
	if tPort, err := strconv.Atoi(port); err != nil || tPort <= 0 || tPort > 65535 {
 | 
			
		||||
		return errors.New("invalid API listen port: " + port)
 | 
			
		||||
	}
 | 
			
		||||
 | 
			
		||||
	if c.WatchedInterface == "" {
 | 
			
		||||
		tIf, err := utils.FindDefaultAdapter()
 | 
			
		||||
		if err != nil {
 | 
			
		||||
			return errors.New("invalid admin UID: " + uidStr)
 | 
			
		||||
			return errors.WrapPrefix(err, "Watched interface is not defined and failed to find a default one.", 0)
 | 
			
		||||
		}
 | 
			
		||||
		cfg.AdminUIDs[uid] = struct{}{}
 | 
			
		||||
		c.WatchedInterface = tIf.Name
 | 
			
		||||
		logger.Infof("Watching default interface %s", c.WatchedInterface)
 | 
			
		||||
	}
 | 
			
		||||
 | 
			
		||||
	announceCmdsEnv := os.Getenv("TG_ANNOUNCE_CMDS")
 | 
			
		||||
	if !lo.Contains([]string{"", "no", "false", "0"}, strings.ToLower(announceCmdsEnv)) {
 | 
			
		||||
		cfg.TGAnnounceCommands = true
 | 
			
		||||
	}
 | 
			
		||||
 | 
			
		||||
	cfg.WatchedInterface = "eth0"
 | 
			
		||||
	if iface := os.Getenv("TG_WATCHED_INTERFACE"); iface != "" {
 | 
			
		||||
		cfg.WatchedInterface = iface
 | 
			
		||||
	}
 | 
			
		||||
 | 
			
		||||
	cfg.MonthlyTrafficLimitGiB = 1000
 | 
			
		||||
	if trafficLimitStr := os.Getenv("TG_MONTHLY_TRAFFIC_LIMIT_GIB"); trafficLimitStr != "" {
 | 
			
		||||
		var err error
 | 
			
		||||
		if cfg.MonthlyTrafficLimitGiB, err = strconv.Atoi(trafficLimitStr); err != nil {
 | 
			
		||||
			return errors.New("invalid traffic limit: " + trafficLimitStr)
 | 
			
		||||
		}
 | 
			
		||||
	}
 | 
			
		||||
 | 
			
		||||
	config = &cfg
 | 
			
		||||
	return nil
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
func (c *ConfigDef) IsAdmin(id int64) bool {
 | 
			
		||||
	_, ok := c.adminUidLookup[id]
 | 
			
		||||
	return ok
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
func (c *ConfigDef) GetAPIListenAddr() net.Addr {
 | 
			
		||||
	return c.apiListenAddr
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
func configFromEnv() (*ConfigDef, error) {
 | 
			
		||||
	cfg := ConfigDef{}
 | 
			
		||||
	err := cleanenv.ReadEnv(&cfg)
 | 
			
		||||
	if err != nil {
 | 
			
		||||
		return nil, errors.Wrap(err, 0)
 | 
			
		||||
	}
 | 
			
		||||
	return &cfg, nil
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
var config *ConfigDef
 | 
			
		||||
 | 
			
		||||
func LoadCfg() error {
 | 
			
		||||
	cfg, err := configFromEnv()
 | 
			
		||||
	if err != nil {
 | 
			
		||||
		return errors.WrapPrefix(err, "failed to load config from environment", 0)
 | 
			
		||||
	}
 | 
			
		||||
	if err = cfg.Parse(); err != nil {
 | 
			
		||||
		return errors.WrapPrefix(err, "failed to parse config", 0)
 | 
			
		||||
	}
 | 
			
		||||
 | 
			
		||||
	config = cfg
 | 
			
		||||
	return nil
 | 
			
		||||
}
 | 
			
		||||
| 
						 | 
				
			
			
 | 
			
		|||
| 
						 | 
				
			
			@ -0,0 +1,53 @@
 | 
			
		|||
package utils
 | 
			
		||||
 | 
			
		||||
import (
 | 
			
		||||
	"net"
 | 
			
		||||
 | 
			
		||||
	"github.com/go-errors/errors"
 | 
			
		||||
)
 | 
			
		||||
 | 
			
		||||
func getLocalIP() (net.IP, error) {
 | 
			
		||||
	conn, err := net.Dial("udp", "1.0.0.1:53")
 | 
			
		||||
	if err != nil {
 | 
			
		||||
		return nil, err
 | 
			
		||||
	}
 | 
			
		||||
	defer conn.Close()
 | 
			
		||||
 | 
			
		||||
	host, _, err := net.SplitHostPort(conn.LocalAddr().String())
 | 
			
		||||
	if err != nil {
 | 
			
		||||
		return nil, errors.Wrap(err, 0)
 | 
			
		||||
	}
 | 
			
		||||
	ip := net.ParseIP(host)
 | 
			
		||||
	if ip == nil {
 | 
			
		||||
		return nil, errors.Errorf("failed to parse IP address %q", host)
 | 
			
		||||
	}
 | 
			
		||||
	return ip, nil
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
func FindDefaultAdapter() (*net.Interface, error) {
 | 
			
		||||
	localIP, err := getLocalIP()
 | 
			
		||||
	if err != nil {
 | 
			
		||||
		return nil, errors.WrapPrefix(err, "failed to get local address", 0)
 | 
			
		||||
	}
 | 
			
		||||
 | 
			
		||||
	iflist, err := net.Interfaces()
 | 
			
		||||
	if err != nil {
 | 
			
		||||
		return nil, errors.WrapPrefix(err, "failed to get interface list", 0)
 | 
			
		||||
	}
 | 
			
		||||
	for _, iface := range iflist {
 | 
			
		||||
		addrs, err := iface.Addrs()
 | 
			
		||||
		if err != nil {
 | 
			
		||||
			continue
 | 
			
		||||
		}
 | 
			
		||||
		for _, addr := range addrs {
 | 
			
		||||
			ipnet, ok := addr.(*net.IPNet)
 | 
			
		||||
			if !ok {
 | 
			
		||||
				continue
 | 
			
		||||
			}
 | 
			
		||||
			if localIP.Equal(ipnet.IP) {
 | 
			
		||||
				return &iface, nil
 | 
			
		||||
			}
 | 
			
		||||
		}
 | 
			
		||||
	}
 | 
			
		||||
	return nil, errors.New("failed to find default adapter")
 | 
			
		||||
}
 | 
			
		||||
		Loading…
	
		Reference in New Issue