rtt-polybar-waybar/rtt.go

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++
		}
	}

	if pkgsInArr == 0 {
		return nil, fmt.Errorf("no packets received")
	}

	ps.AvgRtt = time.Duration(sum / int64(pkgsInArr))
	ps.PacketLoss = math.Min(100, math.Max(100-(float64(dReceive)/float64(dSent)*100), 0))
	if math.IsNaN(ps.PacketLoss) || math.IsInf(ps.PacketLoss, 0) {
		ps.PacketLoss = 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()

	var pinger *probing.Pinger
	var err error
	for pinger, err = probing.NewPinger(*Host); err != nil; {
		time.Sleep(2 * time.Second)
	}

	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
		switch {
		case err != nil:
			fmt.Println(err)
		case *Waybar:
			formatLineWaybar(pStats)
		default:
			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
				switch {
				case err != nil:
					fmt.Println(err)
				case *Waybar:
					formatLineWaybar(pStats)
				default:
					formatLinePolybar(pStats)
				}
			}
			time.Sleep(time.Second)
		}
	}()

	go func() {
		for {
			if err = pinger.Run(); err != nil {
				fmt.Println(err)
				time.Sleep(time.Second)
			}
		}
	}()

	for range killSignals {
		break
	}
	pinger.Stop()
}