From e6be1c8f28a185dd5b6c821a979644d83471b747 Mon Sep 17 00:00:00 2001 From: wangjian Date: Tue, 28 Feb 2023 09:56:20 +0800 Subject: [PATCH] first commit --- config/cmd/server.go | 133 ++++++++++++++++++++++++++++ config/config.go | 62 +++++++++++++ config/config.yaml | 21 +++++ go.mod | 61 +++++++++++++ internal/monitor/monitor.go | 210 ++++++++++++++++++++++++++++++++++++++++++++ internal/monitor/myip.go | 90 +++++++++++++++++++ main.go | 27 ++++++ model/node.go | 56 ++++++++++++ pkg/utils/http.go | 92 +++++++++++++++++++ pkg/utils/utils.go | 66 ++++++++++++++ pkg/utils/utils_test.go | 41 +++++++++ 11 files changed, 859 insertions(+) create mode 100644 config/cmd/server.go create mode 100644 config/config.go create mode 100644 config/config.yaml create mode 100644 go.mod create mode 100644 internal/monitor/monitor.go create mode 100644 internal/monitor/myip.go create mode 100644 main.go create mode 100644 model/node.go create mode 100644 pkg/utils/http.go create mode 100644 pkg/utils/utils.go create mode 100644 pkg/utils/utils_test.go diff --git a/config/cmd/server.go b/config/cmd/server.go new file mode 100644 index 0000000..fc0f81d --- /dev/null +++ b/config/cmd/server.go @@ -0,0 +1,133 @@ +package cmd + +import ( + "environmentCaptureAgent/config" + "environmentCaptureAgent/internal/monitor" + "fmt" + "git.hpds.cc/pavement/hpds_node" + "go.uber.org/zap" + "os" + "os/signal" + "syscall" + "time" + + "git.hpds.cc/Component/logging" + "github.com/spf13/cobra" +) + +var ( + ConfigFileFlag string = "./config/config.yaml" + logger *logging.Logger +) + +func must(err error) { + if err != nil { + fmt.Fprint(os.Stderr, err) + os.Exit(1) + } +} + +func NewStartCmd() *cobra.Command { + cmd := &cobra.Command{ + Use: "start", + Short: "Start hpds environment capture agent", + Run: func(cmd *cobra.Command, args []string) { + var ( + cfg *config.AgentConfig + err error + ) + //ctx, cancel := context.WithCancel(context.Background()) + //defer cancel() + configFileFlag, err := cmd.Flags().GetString("c") + if err != nil { + fmt.Println("get local config err: ", err) + return + } + + must(err) + cfg, err = config.ParseConfigByFile(configFileFlag) + must(err) + logger = LoadLoggerConfig(cfg.Logging) + exitChannel := make(chan os.Signal) + defer close(exitChannel) + + // 退出信号监听 + go func(c chan os.Signal) { + signal.Notify(c, syscall.SIGINT, syscall.SIGTERM) + }(exitChannel) + ap := hpds_node.NewAccessPoint( + cfg.Name, + hpds_node.WithMqAddr(fmt.Sprintf("%s:%d", cfg.Node.Host, cfg.Node.Port)), + hpds_node.WithCredential(cfg.Node.Token), + ) + err = ap.Connect() + must(err) + //defer ap.Close() + //for _, v := range cfg.Funcs { + // ap.SetDataTag(v.DataTag) + //} + ap.SetDataTag(18) + node := monitor.GetHost() + byteNode := node.ToByte() + _ = generateAndSendData(ap, byteNode) + + ticker := time.NewTicker(time.Duration(cfg.Delay) * time.Second) + count := 0 + //c := cron.New() + //spec := fmt.Sprintf("*/%d * * * * *", cfg.Delay) + //_, err = c.AddFunc(spec, func() { + // stat := monitor.GetState().ToByte() + // _ = generateAndSendData(ap, stat) + // logger.With( + // zap.String("agent", "发送状态信息"), + // ) + //}) + //must(err) + //c.Start() + for { + select { + case <-ticker.C: + stat := monitor.GetState(node.NodeName).ToByte() + go func() { + _ = generateAndSendData(ap, stat) + }() + case errs := <-exitChannel: + count++ + if count > 3 { + logger.With( + zap.String("agent", "服务退出"), + ).Info(errs.String()) + return + } + } + } + }, + } + cmd.Flags().StringVar(&ConfigFileFlag, "c", "./config/config.yaml", "The configuration file path") + return cmd +} + +func generateAndSendData(stream hpds_node.AccessPoint, data []byte) error { + _, err := stream.Write(data) + if err != nil { + return err + } + time.Sleep(1000 * time.Millisecond) + return nil +} + +func LoadLoggerConfig(opt config.LogOptions) *logging.Logger { + return logging.NewLogger( + logging.SetPath(opt.Path), + logging.SetPrefix(opt.Prefix), + logging.SetDevelopment(opt.Development), + logging.SetDebugFileSuffix(opt.DebugFileSuffix), + logging.SetWarnFileSuffix(opt.WarnFileSuffix), + logging.SetErrorFileSuffix(opt.ErrorFileSuffix), + logging.SetInfoFileSuffix(opt.InfoFileSuffix), + logging.SetMaxAge(opt.MaxAge), + logging.SetMaxBackups(opt.MaxBackups), + logging.SetMaxSize(opt.MaxSize), + logging.SetLevel(logging.LogLevel["debug"]), + ) +} diff --git a/config/config.go b/config/config.go new file mode 100644 index 0000000..31c201d --- /dev/null +++ b/config/config.go @@ -0,0 +1,62 @@ +package config + +import ( + "bytes" + "os" + + "github.com/spf13/viper" +) + +type AgentConfig struct { + Name string `yaml:"name,omitempty"` + Mode string `yaml:"mode,omitempty"` + Delay int `yaml:"delay"` + Logging LogOptions `yaml:"logging"` + Node HpdsNode `yaml:"node,omitempty"` + Funcs []FuncConfig `yaml:"functions,omitempty"` +} + +type FuncConfig struct { + Name string `yaml:"name"` + DataTag uint8 `yaml:"dataTag"` +} + +type HpdsNode struct { + Host string `yaml:"host"` + Port int `yaml:"port"` + Token string `yaml:"token,omitempty"` +} + +type LogOptions struct { + Path string `yaml:"path" json:"path" toml:"path"` // 文件保存地方 + Prefix string `yaml:"prefix" json:"prefix" toml:"prefix"` // 日志文件前缀 + ErrorFileSuffix string `yaml:"errorFileSuffix" json:"errorFileSuffix" toml:"errorFileSuffix"` // error日志文件后缀 + WarnFileSuffix string `yaml:"warnFileSuffix" json:"warnFileSuffix" toml:"warnFileSuffix"` // warn日志文件后缀 + InfoFileSuffix string `yaml:"infoFileSuffix" json:"infoFileSuffix" toml:"infoFileSuffix"` // info日志文件后缀 + DebugFileSuffix string `yaml:"debugFileSuffix" json:"debugFileSuffix" toml:"debugFileSuffix"` // debug日志文件后缀 + Level string `yaml:"level" json:"level" toml:"level"` // 日志等级 + MaxSize int `yaml:"maxSize" json:"maxSize" toml:"maxSize"` // 日志文件大小（M） + MaxBackups int `yaml:"maxBackups" json:"maxBackups" toml:"maxBackups"` // 最多存在多少个切片文件 + MaxAge int `yaml:"maxAge" json:"maxAge" toml:"maxAge"` // 保存的最大天数 + Development bool `yaml:"development" json:"development" toml:"development"` // 是否是开发模式 +} + +func ParseConfigByFile(path string) (cfg *AgentConfig, err error) { + buffer, err := os.ReadFile(path) + if err != nil { + return nil, err + } + return load(buffer) +} + +func load(buf []byte) (cfg *AgentConfig, err error) { + cViper := viper.New() + cViper.SetConfigType("yaml") + cfg = new(AgentConfig) + cViper.ReadConfig(bytes.NewBuffer(buf)) + err = cViper.Unmarshal(cfg) + if err != nil { + return nil, err + } + return +} diff --git a/config/config.yaml b/config/config.yaml new file mode 100644 index 0000000..72cdfba --- /dev/null +++ b/config/config.yaml @@ -0,0 +1,21 @@ +name: capture-agent +mode: dev +delay: 15 +node: + host: 114.55.236.153 + port: 9188 + token: 06d36c6f5705507dae778fdce90d0767 +logging: + path: ./logs + prefix: capture-agent + errorFileSuffix: error.log + warnFileSuffix: warn.log + infoFileSuffix: info.log + debugFileSuffix: debug.log + maxSize: 100 + maxBackups: 3000 + maxAge: 30 + development: true +functions: + - name: capture-agent + dataTag: 18 \ No newline at end of file diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..cff24ea --- /dev/null +++ b/go.mod @@ -0,0 +1,61 @@ +module environmentCaptureAgent + +go 1.19 + +require ( + git.hpds.cc/Component/logging v0.0.0-20230106105738-e378e873921b + git.hpds.cc/pavement/hpds_node v0.0.0-20221023053316-37f7ba99eab3 + github.com/Erope/goss v0.0.0-20211230093305-df3c03fd1ed4 + github.com/robfig/cron/v3 v3.0.1 + github.com/shirou/gopsutil/v3 v3.23.1 + github.com/spf13/cobra v1.6.1 + github.com/spf13/viper v1.15.0 + github.com/stretchr/testify v1.8.1 +) + +require ( + git.hpds.cc/Component/mq_coder v0.0.0-20221010064749-174ae7ae3340 // indirect + git.hpds.cc/Component/network v0.0.0-20221012021659-2433c68452d5 // indirect + github.com/davecgh/go-spew v1.1.1 // indirect + github.com/fsnotify/fsnotify v1.6.0 // indirect + github.com/go-ole/go-ole v1.2.6 // indirect + github.com/go-task/slim-sprig v0.0.0-20210107165309-348f09dbbbc0 // indirect + github.com/golang/mock v1.6.0 // indirect + github.com/hashicorp/hcl v1.0.0 // indirect + github.com/inconshreveable/mousetrap v1.0.1 // indirect + github.com/lucas-clemente/quic-go v0.29.1 // indirect + github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect + github.com/magiconair/properties v1.8.7 // indirect + github.com/marten-seemann/qtls-go1-18 v0.1.2 // indirect + github.com/marten-seemann/qtls-go1-19 v0.1.0 // indirect + github.com/matoous/go-nanoid/v2 v2.0.0 // indirect + github.com/mitchellh/mapstructure v1.5.0 // indirect + github.com/nxadm/tail v1.4.8 // indirect + github.com/onsi/ginkgo v1.16.4 // indirect + github.com/pelletier/go-toml/v2 v2.0.6 // indirect + github.com/pkg/errors v0.9.1 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect + github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect + github.com/spf13/afero v1.9.3 // indirect + github.com/spf13/cast v1.5.0 // indirect + github.com/spf13/jwalterweatherman v1.1.0 // indirect + github.com/spf13/pflag v1.0.5 // indirect + github.com/subosito/gotenv v1.4.2 // indirect + github.com/tklauser/go-sysconf v0.3.11 // indirect + github.com/tklauser/numcpus v0.6.0 // indirect + github.com/yusufpapurcu/wmi v1.2.2 // indirect + go.uber.org/atomic v1.9.0 // indirect + go.uber.org/multierr v1.8.0 // indirect + go.uber.org/zap v1.23.0 // indirect + golang.org/x/crypto v0.0.0-20220525230936-793ad666bf5e // indirect + golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e // indirect + golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4 // indirect + golang.org/x/net v0.4.0 // indirect + golang.org/x/sys v0.4.0 // indirect + golang.org/x/text v0.5.0 // indirect + golang.org/x/tools v0.1.12 // indirect + gopkg.in/ini.v1 v1.67.0 // indirect + gopkg.in/natefinch/lumberjack.v2 v2.0.0 // indirect + gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/internal/monitor/monitor.go b/internal/monitor/monitor.go new file mode 100644 index 0000000..2395c72 --- /dev/null +++ b/internal/monitor/monitor.go @@ -0,0 +1,210 @@ +package monitor + +import ( + "environmentCaptureAgent/model" + "fmt" + "github.com/shirou/gopsutil/v3/disk" + "os/exec" + "regexp" + "runtime" + "strconv" + "strings" + "syscall" + "time" + + "github.com/Erope/goss" + "github.com/shirou/gopsutil/v3/cpu" + "github.com/shirou/gopsutil/v3/host" + "github.com/shirou/gopsutil/v3/load" + "github.com/shirou/gopsutil/v3/mem" + "github.com/shirou/gopsutil/v3/net" + "github.com/shirou/gopsutil/v3/process" +) + +var ( + netInSpeed, netOutSpeed, netInTransfer, netOutTransfer, lastUpdateNetStats uint64 + cachedBootTime time.Time + expectDiskFsTypes = []string{ + "apfs", "ext4", "ext3", "ext2", "f2fs", "reiserfs", "jfs", "btrfs", + "fuseblk", "zfs", "simfs", "ntfs", "fat32", "exfat", "xfs", "fuse.rclone", + } + getMacDiskNo = regexp.MustCompile(`\/dev\/disk(\d)s.*`) + Version string +) + +func GetHost() *model.Node { + hi, _ := host.Info() + var cpuType string + if hi.VirtualizationSystem != "" { + cpuType = "Virtual" + } else { + cpuType = "Physical" + } + cpuModelCount := make(map[string]int) + ci, _ := cpu.Info() + for i := 0; i < len(ci); i++ { + cpuModelCount[ci[i].ModelName]++ + } + var cpus []string + for model, count := range cpuModelCount { + cpus = append(cpus, fmt.Sprintf("%s %d %s Core", model, count, cpuType)) + } + mv, _ := mem.VirtualMemory() + diskTotal, _ := getDiskTotalAndUsed() + + var swapMemTotal uint64 + if runtime.GOOS == "windows" { + ms, _ := mem.SwapMemory() + swapMemTotal = ms.Total + } else { + swapMemTotal = mv.SwapTotal + } + + if cachedBootTime.IsZero() { + cachedBootTime = time.Unix(int64(hi.BootTime), 0) + } + + return &model.Node{ + NodeName: hi.HostID, + Platform: hi.OS, + PlatformVersion: hi.PlatformVersion, + CPU: cpus, + MemTotal: mv.Total, + SwapTotal: swapMemTotal, + DiskTotal: diskTotal, + Arch: hi.KernelArch, + Virtualization: hi.VirtualizationSystem, + BootTime: hi.BootTime, + IP: cachedIP, + CountryCode: strings.ToLower(cachedCountry), + Version: Version, + } +} + +func GetState(nodeName string) *model.NodeState { + procs, _ := process.Pids() + + mv, _ := mem.VirtualMemory() + + var swapMemUsed uint64 + if runtime.GOOS == "windows" { + // gopsutil 在 Windows 下不能正确取 swap + ms, _ := mem.SwapMemory() + swapMemUsed = ms.Used + } else { + swapMemUsed = mv.SwapTotal - mv.SwapFree + } + + var cpuPercent float64 + cp, err := cpu.Percent(0, false) + if err == nil { + cpuPercent = cp[0] + } + + _, diskUsed := getDiskTotalAndUsed() + loadStat, _ := load.Avg() + + var tcpConnCount, udpConnCount uint64 + + ssErr := true + if runtime.GOOS == "linux" { + tcpStat, errTcp := goss.ConnectionsWithProtocol(syscall.IPPROTO_TCP) + udpStat, errUdp := goss.ConnectionsWithProtocol(syscall.IPPROTO_UDP) + if errTcp == nil && errUdp == nil { + ssErr = false + tcpConnCount = uint64(len(tcpStat)) + udpConnCount = uint64(len(udpStat)) + } + } + if ssErr { + conns, _ := net.Connections("all") + for i := 0; i < len(conns); i++ { + switch conns[i].Type { + case syscall.SOCK_STREAM: + tcpConnCount++ + case syscall.SOCK_DGRAM: + udpConnCount++ + } + } + } + + return &model.NodeState{ + NodeName: nodeName, + CPU: cpuPercent, + MemUsed: mv.Total - mv.Available, + SwapUsed: swapMemUsed, + DiskUsed: diskUsed, + NetInTransfer: netInTransfer, + NetOutTransfer: netOutTransfer, + NetInSpeed: netInSpeed, + NetOutSpeed: netOutSpeed, + Uptime: uint64(time.Since(cachedBootTime).Seconds()), + Load1: loadStat.Load1, + Load5: loadStat.Load5, + Load15: loadStat.Load15, + TcpConnCount: tcpConnCount, + UdpConnCount: udpConnCount, + ProcessCount: uint64(len(procs)), + } +} + +func getDiskTotalAndUsed() (total uint64, used uint64) { + diskList, _ := disk.Partitions(false) + devices := make(map[string]string) + countedDiskForMac := make(map[string]struct{}) + for _, d := range diskList { + fsType := strings.ToLower(d.Fstype) + // 不统计 K8s 的虚拟挂载点：https://github.com/shirou/gopsutil/issues/1007 + if devices[d.Device] == "" && isListContainsStr(expectDiskFsTypes, fsType) && !strings.Contains(d.Mountpoint, "/var/lib/kubelet") { + devices[d.Device] = d.Mountpoint + } + } + for device, mountPath := range devices { + diskUsageOf, _ := disk.Usage(mountPath) + // 这里是针对 Mac 机器的处理，https://github.com/giampaolo/psutil/issues/1509 + matches := getMacDiskNo.FindStringSubmatch(device) + if len(matches) == 2 { + if _, has := countedDiskForMac[matches[1]]; !has { + countedDiskForMac[matches[1]] = struct{}{} + total += diskUsageOf.Total + } + } else { + total += diskUsageOf.Total + } + used += diskUsageOf.Used + } + + // Fallback 到这个方法,仅统计根路径,适用于OpenVZ之类的. + if runtime.GOOS == "linux" { + if total == 0 && used == 0 { + cmd := exec.Command("df") + out, err := cmd.CombinedOutput() + if err == nil { + s := strings.Split(string(out), "\n") + for _, c := range s { + info := strings.Fields(c) + if len(info) == 6 { + if info[5] == "/" { + total, _ = strconv.ParseUint(info[1], 0, 64) + used, _ = strconv.ParseUint(info[2], 0, 64) + // 默认获取的是1K块为单位的. + total = total * 1024 + used = used * 1024 + } + } + } + } + } + } + + return +} + +func isListContainsStr(list []string, str string) bool { + for i := 0; i < len(list); i++ { + if strings.Contains(str, list[i]) { + return true + } + } + return false +} diff --git a/internal/monitor/myip.go b/internal/monitor/myip.go new file mode 100644 index 0000000..edc469d --- /dev/null +++ b/internal/monitor/myip.go @@ -0,0 +1,90 @@ +package monitor + +import ( + "encoding/json" + "environmentCaptureAgent/pkg/utils" + "fmt" + "io/ioutil" + "net/http" + "strings" + "time" +) + +type geoIP struct { + CountryCode string `json:"country_code,omitempty"` + IP string `json:"ip,omitempty"` + Query string `json:"query,omitempty"` +} + +var ( + geoIPApiList = []string{ + "https://api.ip.sb/geoip", + "https://ip.seeip.org/geoip", + "https://ipapi.co/json", + "https://freegeoip.app/json/", + "http://ip-api.com/json/", + "https://extreme-ip-lookup.com/json/", + } + cachedIP, cachedCountry string + httpClientV4 = utils.NewSingleStackHTTPClient(time.Second*20, time.Second*5, time.Second*10, false) + httpClientV6 = utils.NewSingleStackHTTPClient(time.Second*20, time.Second*5, time.Second*10, true) +) + +func UpdateIP() { + for { + ipv4 := fetchGeoIP(geoIPApiList, false) + ipv6 := fetchGeoIP(geoIPApiList, true) + if ipv4.IP == "" && ipv6.IP == "" { + time.Sleep(time.Minute) + continue + } + if ipv4.IP == "" || ipv6.IP == "" { + cachedIP = fmt.Sprintf("%s%s", ipv4.IP, ipv6.IP) + } else { + cachedIP = fmt.Sprintf("%s/%s", ipv4.IP, ipv6.IP) + } + if ipv4.CountryCode != "" { + cachedCountry = ipv4.CountryCode + } else if ipv6.CountryCode != "" { + cachedCountry = ipv6.CountryCode + } + time.Sleep(time.Minute * 30) + } +} + +func fetchGeoIP(servers []string, isV6 bool) geoIP { + var ip geoIP + var resp *http.Response + var err error + for i := 0; i < len(servers); i++ { + if isV6 { + resp, err = httpClientV6.Get(servers[i]) + } else { + resp, err = httpClientV4.Get(servers[i]) + } + if err == nil { + body, err := ioutil.ReadAll(resp.Body) + if err != nil { + continue + } + resp.Body.Close() + err = json.Unmarshal(body, &ip) + if err != nil { + continue + } + if ip.IP == "" && ip.Query != "" { + ip.IP = ip.Query + } + // 没取到 v6 IP + if isV6 && !strings.Contains(ip.IP, ":") { + continue + } + // 没取到 v4 IP + if !isV6 && !strings.Contains(ip.IP, ".") { + continue + } + return ip + } + } + return ip +} diff --git a/main.go b/main.go new file mode 100644 index 0000000..106a933 --- /dev/null +++ b/main.go @@ -0,0 +1,27 @@ +package main + +import ( + "environmentCaptureAgent/config/cmd" + "fmt" + "github.com/spf13/cobra" + "os" +) + +var ( + rootCmd = &cobra.Command{ + Use: "hpds_environment_capture_agent", + Long: "hpds_agent is a environment capture agent", + Version: "0.1", + } +) + +func init() { + rootCmd.AddCommand(cmd.NewStartCmd()) +} +func main() { + + if err := rootCmd.Execute(); err != nil { + fmt.Fprint(os.Stderr, err.Error()) + os.Exit(1) + } +} diff --git a/model/node.go b/model/node.go new file mode 100644 index 0000000..6321745 --- /dev/null +++ b/model/node.go @@ -0,0 +1,56 @@ +package model + +import "encoding/json" + +// Node 节点信息 +type Node struct { + NodeName string `json:"nodeName"` + Platform string `json:"platform,omitempty"` + PlatformVersion string `json:"platformVersion,omitempty"` + CPU []string `json:"cpu,omitempty"` + MemTotal uint64 `json:"memTotal,omitempty"` + DiskTotal uint64 `json:"diskTotal,omitempty"` + SwapTotal uint64 `json:"swapTotal,omitempty"` + Arch string `json:"arch,omitempty"` + Virtualization string `json:"virtualization,omitempty"` + BootTime uint64 `json:"bootTime,omitempty"` + IP string `json:"ip"` + CountryCode string `json:"countryCode,omitempty"` + Version string `json:"version,omitempty"` +} + +func (n Node) ToByte() []byte { + data, err := json.Marshal(n) + if err != nil { + return []byte("") + } + return data +} + +// NodeState 节点状态信息 +type NodeState struct { + NodeName string `json:"nodeName"` + CPU float64 `json:"cpu,omitempty"` + MemUsed uint64 `json:"memUsed,omitempty"` + SwapUsed uint64 `json:"swapUsed,omitempty"` + DiskUsed uint64 `json:"diskUsed,omitempty"` + NetInTransfer uint64 `json:"netInTransfer,omitempty"` + NetOutTransfer uint64 `json:"netOutTransfer,omitempty"` + NetInSpeed uint64 `json:"netInSpeed,omitempty"` + NetOutSpeed uint64 `json:"netOutSpeed,omitempty"` + Uptime uint64 `json:"uptime,omitempty"` + Load1 float64 `json:"load1,omitempty"` + Load5 float64 `json:"load5,omitempty"` + Load15 float64 `json:"load15,omitempty"` + TcpConnCount uint64 `json:"tcpConnCount,omitempty"` + UdpConnCount uint64 `json:"udpConnCount,omitempty"` + ProcessCount uint64 `json:"processCount,omitempty"` +} + +func (ns NodeState) ToByte() []byte { + data, err := json.Marshal(ns) + if err != nil { + return []byte("") + } + return data +} diff --git a/pkg/utils/http.go b/pkg/utils/http.go new file mode 100644 index 0000000..06d946a --- /dev/null +++ b/pkg/utils/http.go @@ -0,0 +1,92 @@ +package utils + +import ( + "context" + "errors" + "net" + "net/http" + "strings" + "time" +) + +func NewSingleStackHTTPClient(httpTimeout, dialTimeout, keepAliveTimeout time.Duration, ipv6 bool) *http.Client { + dialer := &net.Dialer{ + Timeout: dialTimeout, + KeepAlive: keepAliveTimeout, + } + + transport := &http.Transport{ + Proxy: http.ProxyFromEnvironment, + ForceAttemptHTTP2: false, + DialContext: func(ctx context.Context, network string, addr string) (net.Conn, error) { + ip, err := resolveIP(addr, ipv6) + if err != nil { + return nil, err + } + return dialer.DialContext(ctx, network, ip) + }, + } + + return &http.Client{ + Transport: transport, + Timeout: httpTimeout, + } +} + +func resolveIP(addr string, ipv6 bool) (string, error) { + url := strings.Split(addr, ":") + + dnsServers := []string{"[2606:4700:4700::1001]", "[2001:4860:4860::8844]", "[2400:3200::1]", "[2400:3200:baba::1]"} + if !ipv6 { + dnsServers = []string{"1.0.0.1", "8.8.4.4", "223.5.5.5", "223.6.6.6"} + } + + res, err := net.LookupIP(url[0]) + if err != nil { + for i := 0; i < len(dnsServers); i++ { + r := &net.Resolver{ + PreferGo: true, + Dial: func(ctx context.Context, network, address string) (net.Conn, error) { + d := net.Dialer{ + Timeout: time.Second * 10, + } + return d.DialContext(ctx, "udp", dnsServers[i]+":53") + }, + } + res, err = r.LookupIP(context.Background(), "ip", url[0]) + if err == nil { + break + } + } + } + + if err != nil { + return "", err + } + + var ipv4Resolved, ipv6Resolved bool + + for i := 0; i < len(res); i++ { + ip := res[i].String() + if strings.Contains(ip, ".") && !ipv6 { + ipv4Resolved = true + url[0] = ip + break + } + if strings.Contains(ip, ":") && ipv6 { + ipv6Resolved = true + url[0] = "[" + ip + "]" + break + } + } + + if ipv6 && !ipv6Resolved { + return "", errors.New("the AAAA record not resolved") + } + + if !ipv6 && !ipv4Resolved { + return "", errors.New("the A record not resolved") + } + + return strings.Join(url, ":"), nil +} diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go new file mode 100644 index 0000000..1c4590c --- /dev/null +++ b/pkg/utils/utils.go @@ -0,0 +1,66 @@ +package utils + +import ( + "crypto/md5" // #nosec + "encoding/hex" + "math/rand" + "os" + "regexp" + "time" + "unsafe" +) + +const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" +const ( + letterIdxBits = 6 // 6 bits to represent a letter index + letterIdxMask = 1<= 0; { + if remain == 0 { + cache, remain = src.Int63(), letterIdxMax + } + if idx := int(cache & letterIdxMask); idx < len(letterBytes) { + b[i] = letterBytes[idx] + i-- + } + cache >>= letterIdxBits + remain-- + } + + return *(*string)(unsafe.Pointer(&b)) //#nosec +} + +func MD5(plantext string) string { + hash := md5.New() // #nosec + hash.Write([]byte(plantext)) + return hex.EncodeToString(hash.Sum(nil)) +} + +func IsWindows() bool { + return os.PathSeparator == '\\' && os.PathListSeparator == ';' +} + +var ipv4Re = regexp.MustCompile(`(\d*\.).*(\.\d*)`) + +func ipv4Desensitize(ipv4Addr string) string { + return ipv4Re.ReplaceAllString(ipv4Addr, "$1****$2") +} + +var ipv6Re = regexp.MustCompile(`(\w*:\w*:).*(:\w*:\w*)`) + +func ipv6Desensitize(ipv6Addr string) string { + return ipv6Re.ReplaceAllString(ipv6Addr, "$1****$2") +} + +func IPDesensitize(ipAddr string) string { + ipAddr = ipv4Desensitize(ipAddr) + ipAddr = ipv6Desensitize(ipAddr) + return ipAddr +} diff --git a/pkg/utils/utils_test.go b/pkg/utils/utils_test.go new file mode 100644 index 0000000..688a703 --- /dev/null +++ b/pkg/utils/utils_test.go @@ -0,0 +1,41 @@ +package utils + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +type testSt struct { + input string + output string +} + +func TestNotification(t *testing.T) { + cases := []testSt{ + { + input: "103.80.236.249/d5ce:d811:cdb8:067a:a873:2076:9521:9d2d", + output: "103.****.249/d5ce:d811:****:9521:9d2d", + }, + { + input: "3.80.236.29/d5ce::cdb8:067a:a873:2076:9521:9d2d", + output: "3.****.29/d5ce::****:9521:9d2d", + }, + { + input: "3.80.236.29/d5ce::cdb8:067a:a873:2076::9d2d", + output: "3.****.29/d5ce::****::9d2d", + }, + { + input: "3.80.236.9/d5ce::cdb8:067a:a873:2076::9d2d", + output: "3.****.9/d5ce::****::9d2d", + }, + { + input: "3.80.236.9/d5ce::cdb8:067a:a873:2076::9d2d", + output: "3.****.9/d5ce::****::9d2d", + }, + } + + for _, c := range cases { + assert.Equal(t, IPDesensitize(c.input), c.output) + } +}