adjusted package name to reflect repo name

rtt.go

@@ -0,0 +1,199 @@
+package main
+
+import (
+	"encoding/json"
+	"flag"
+	"fmt"
+	"github.com/montanaflynn/stats"
+	"github.com/prometheus-community/pro-bing"
+	"math"
+	"os"
+	"os/signal"
+	"sync"
+	"syscall"
+	"time"
+)
+
+type ContinuousPinger struct {
+	Pings []*probing.Packet
+	Lock  *sync.RWMutex
+}
+
+var (
+	Waybar              = flag.Bool("waybar", false, "output waybar json format")
+	PingCount           = flag.Int("count", 5, "how many pings to average")
+	PingWarningLimit    = flag.Int("warn", 50, "ping warning limit")
+	PingCritLimit       = flag.Int("crit", 100, "ping critical limit")
+	PacketLossWarnLimit = flag.Int("pwarn", 10, "packetloss warning limit")
+	PacketLossCritLimit = flag.Int("pcrit", 25, "packetloss critical limit")
+	Host                = flag.String("host", "google.com", "host to ping")
+	WarnColor           = flag.String("warnc", "ffbf00", "hex color for warning ping (without #)")
+	CritColor           = flag.String("critc", "ff5050", "hex color for critical ping (without #)")
+)
+
+type PingStats struct {
+	AvgRtt     time.Duration
+	StdDev     time.Duration
+	PacketLoss float64
+}
+
+//goland:noinspection ALL
+const (
+	ResetColor = "%{F-}"
+)
+
+type WaybarOut struct {
+	Class string `json:"class"`
+	Text  string `json:"text"`
+}
+
+func formatLinePolybar(stats *PingStats) {
+	if stats.PacketLoss >= 100.0 {
+		// fontawesome/forkawesome don't have a fitting (free) icon, so we use 🚫 from UTF-8 symbols/emoji
+		fmt.Printf("%%{F#%s}🚫\n", *CritColor)
+		return
+	}
+
+	rttColor := ResetColor
+	packetColor := ResetColor
+
+	switch {
+	case int(stats.AvgRtt.Milliseconds()) >= *PingWarningLimit:
+		rttColor = fmt.Sprintf("%%{F#%s}", *WarnColor)
+	case int(stats.AvgRtt.Milliseconds()) >= *PingCritLimit:
+		rttColor = fmt.Sprintf("%%{F#%s}", *CritColor)
+	}
+
+	switch {
+	case int(math.Round(stats.PacketLoss)) >= *PacketLossWarnLimit:
+		packetColor = fmt.Sprintf("%%{F#%s}", *WarnColor)
+	case int(math.Round(stats.PacketLoss)) >= *PacketLossCritLimit:
+		packetColor = fmt.Sprintf("%%{F#%s}", *CritColor)
+	}
+
+	fmt.Printf("%s\uE4E2 %d \uE43C %dms %s\uF1B2 %d%%\n", rttColor, stats.AvgRtt.Milliseconds(), stats.StdDev.Milliseconds(), packetColor, int(math.Round(stats.PacketLoss)))
+}
+
+func formatLineWaybar(stats *PingStats) {
+	res := new(WaybarOut)
+
+	res.Text = fmt.Sprintf("\uE4E2 %d \uE43C %dms \uF1B2 %d%%", stats.AvgRtt.Milliseconds(), stats.StdDev.Milliseconds(), int(math.Round(stats.PacketLoss)))
+	res.Class = "good"
+
+	switch {
+	case int(math.Round(stats.PacketLoss)) >= *PacketLossWarnLimit || int(stats.AvgRtt.Milliseconds()) >= *PingWarningLimit:
+		res.Class = "warning"
+		fallthrough
+	case int(math.Round(stats.PacketLoss)) >= *PacketLossCritLimit || int(stats.AvgRtt.Milliseconds()) >= *PingCritLimit:
+		res.Class = "critical"
+	}
+
+	jOut, err := json.Marshal(res)
+	if err != nil {
+		return
+	}
+
+	fmt.Println(string(jOut))
+}
+
+func (cp *ContinuousPinger) Push(packet *probing.Packet) {
+	cp.Lock.Lock()
+	defer cp.Lock.Unlock()
+	if len(cp.Pings) >= *PingCount {
+		cp.Pings = cp.Pings[:len(cp.Pings)-1]
+	}
+	cp.Pings = append([]*probing.Packet{packet}, cp.Pings...)
+}
+
+func (cp *ContinuousPinger) Stats(dSent, dReceive int) (*PingStats, error) {
+	cp.Lock.RLock()
+	defer cp.Lock.RUnlock()
+	var rttArray []float64
+	ps := new(PingStats)
+
+	var sum int64
+	var pkgsInArr int
+	for _, pkg := range cp.Pings {
+		if pkg != nil {
+			rttArray = append(rttArray, float64(pkg.Rtt))
+			sum += int64(pkg.Rtt)
+			pkgsInArr++
+		}
+	}
+	ps.AvgRtt = time.Duration(sum / int64(pkgsInArr))
+	ps.PacketLoss = math.Min(100, 100-(float64(dReceive)/float64(dSent)*100))
+	stdDev, err := stats.StandardDeviation(rttArray)
+	if err != nil {
+		return nil, err
+	}
+	ps.StdDev = time.Duration(stdDev)
+
+	return ps, nil
+}
+
+func main() {
+	flag.Parse()
+	killSignals := make(chan os.Signal, 1)
+	signal.Notify(killSignals, syscall.SIGINT, syscall.SIGTERM)
+
+	lp := new(ContinuousPinger)
+	lp.Pings = make([]*probing.Packet, *PingCount)
+	lp.Lock = new(sync.RWMutex)
+	ts := 0
+	tr := 0
+	lr := time.Now().UTC()
+
+	pinger, err := probing.NewPinger(*Host)
+	pinger.RecordRtts = false
+
+	pinger.OnRecv = func(pkt *probing.Packet) {
+		lp.Push(pkt)
+		pStats, err := lp.Stats(pinger.PacketsSent-ts, pinger.PacketsRecv-tr)
+		lr = time.Now().UTC()
+		ts = pinger.PacketsSent
+		tr = pinger.PacketsRecv
+		if err != nil {
+			fmt.Println(err)
+		} else if *Waybar {
+			formatLineWaybar(pStats)
+		} else {
+			formatLinePolybar(pStats)
+		}
+	}
+
+	go func() {
+		for {
+			if time.Since(lr).Milliseconds() > 1500 {
+				pStats, err := lp.Stats(pinger.PacketsSent-ts, pinger.PacketsRecv-tr)
+				ts = pinger.PacketsSent
+				tr = pinger.PacketsRecv
+				if err != nil {
+					fmt.Println(err)
+				} else if *Waybar {
+					formatLineWaybar(pStats)
+				} else {
+					formatLinePolybar(pStats)
+				}
+			}
+			time.Sleep(time.Second)
+		}
+	}()
+
+	go func() {
+		for {
+			if err = pinger.Run(); err != nil {
+				fmt.Println(err)
+				time.Sleep(time.Second)
+			}
+		}
+	}()
+
+killLoop:
+	for {
+		select {
+		case <-killSignals:
+			break killLoop
+		}
+	}
+	pinger.Stop()
+}