speedporter/main.go

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

const (
	bandwidthScale       = 125000.0
	defaultListenPort    = 9090
	defaultListenAddr    = "0.0.0.0"
	defaultTestPeriod    = 5
	defaultSpeedtestExec = "/usr/bin/speedtest"
	listenPortEnvVar     = "SPEEDTEST_METRICS_LISTEN_PORT"
	listenAddrEnvVar     = "SPEEDTEST_METRICS_LISTEN_ADDR"
	testPeriodEnvVar     = "SPEEDTEST_PERIOD_MINS"
	speedtestExecEnvVar  = "SPEEDTEST_EXEC"
)

var (
	labels = []string{
		"interface_name",
		"interface_mac",
		"test_host",
	}
	uploadSpeed = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Name: "isp_upload_speed",
			Help: "The upload speed of the network",
		},
		labels,
	)
	downloadSpeed = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Name: "isp_download_speed",
			Help: "The download speed of the network",
		},
		labels,
	)
	pingLatency = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Name: "isp_ping_latency",
			Help: "The ping latency of the network",
		},
		labels,
	)
	packetLoss = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Name: "isp_packet_loss",
			Help: "The packet loss of the network",
		},
		labels,
	)

	listenPort     uint16
	listenAddr     string
	testPeriodMins uint
	speedtestExec  string
)

type SpeedTestResult struct {
	Type      string
	Timestamp time.Time
	Ping      struct {
		Jitter  float64
		Latency float64
		Low     float64
		High    float64
	}
	Download struct {
		Bandwidth int
		Bytes     int
		Elapsed   int
		Latency   struct {
			IQM    float64
			Low    float64
			High   float64
			Jitter float64
		}
	}
	Upload struct {
		Bandwidth int
		Bytes     int
		Elapsed   int
		Latency   struct {
			IQM    float64
			Low    float64
			High   float64
			Jitter float64
		}
	}
	PacketLoss float64
	ISP        string
	Interface  struct {
		InternalIP string
		Name       string
		MACAddr    string
		IsVPN      bool
		ExternalIP string
	}
	Server struct {
		ID       int
		Host     string
		Port     int16
		Name     string
		Location string
		Country  string
		IP       string
	}
	Result struct {
		ID        string
		URL       string
		Persisted bool
	}
}

func init() {
	prometheus.MustRegister(uploadSpeed)
	prometheus.MustRegister(downloadSpeed)
	prometheus.MustRegister(pingLatency)
	prometheus.MustRegister(packetLoss)

	if lpEV, ok := os.LookupEnv(listenPortEnvVar); !ok {
		listenPort = defaultListenPort
	} else {
		val, err := strconv.Atoi(lpEV)
		if err != nil {
			panic(fmt.Sprintf(
				"Invalid Listen Port Specified in Environment Variable: %s",
				listenPortEnvVar,
			))
		}
		listenPort = uint16(val)
		if int(listenPort) != val {
			panic(fmt.Sprintf(
				"Invalid Listen Port Specified in Environment Variable: %s=%d",
				listenPortEnvVar,
				val,
			))
		}
	}
	if tpEV, ok := os.LookupEnv(testPeriodEnvVar); !ok {
		testPeriodMins = defaultTestPeriod
	} else {
		val, err := strconv.Atoi(tpEV)
		if err != nil {
			panic(fmt.Sprintf(
				"Invalid Test Period Specified in Environment Variable: %s",
				testPeriodEnvVar,
			))
		}
		testPeriodMins = uint(val)
		if int(testPeriodMins) != val {
			panic(fmt.Sprintf(
				"Invalid Test Period Specified in Environment Variable: %s=%d",
				testPeriodEnvVar,
				val,
			))
		}
	}
	if laEV, ok := os.LookupEnv(listenAddrEnvVar); !ok {
		listenAddr = defaultListenAddr
	} else if net.ParseIP(laEV) == nil {
		panic(fmt.Sprintf(
			"Invalid Listen Address Specified in Environment Variable: %s=%s",
			listenAddrEnvVar,
			laEV,
		))
	} else {
		listenAddr = laEV
	}

	if seEV, ok := os.LookupEnv(speedtestExecEnvVar); !ok {
		speedtestExec = defaultSpeedtestExec
	} else {
		speedtestExec = seEV
	}

	cmd := exec.Command(speedtestExec, "-h")
	_, err := cmd.Output()
	if err != nil {
		panic(fmt.Sprintf("Error executing command: %v", err))
	}
}

func runTest() {
	cmd := exec.Command(speedtestExec, "-f", "json")
	output, err := cmd.Output()
	if err != nil {
		log.Printf("Error executing command: %v", err)
		return
	}

	var result SpeedTestResult
	err = json.Unmarshal(output, &result)
	if err != nil {
		log.Printf("Error parsing JSON: %v\n", err)
		return
	} else if result.Type != "result" {
		log.Printf("Did not receive a result: %s\n", result.Type)
		return
	}

	uploadBandwidth := float64(result.Upload.Bandwidth) / bandwidthScale
	downloadBandwidth := float64(result.Download.Bandwidth) / bandwidthScale

	labelVals := []string{
		result.Interface.Name,
		result.Interface.MACAddr,
		result.Server.Host,
	}

	uploadSpeed.WithLabelValues(labelVals...).Set(uploadBandwidth)
	downloadSpeed.WithLabelValues(labelVals...).Set(downloadBandwidth)
	pingLatency.WithLabelValues(labelVals...).Set(result.Ping.Latency)
	packetLoss.WithLabelValues(labelVals...).Set(result.PacketLoss)
}

func main() {
	go func() {
		for {
			runTest()
			time.Sleep(time.Duration(testPeriodMins) * time.Minute)
		}
	}()

	http.Handle("/metrics", promhttp.Handler())
	fmt.Printf("Starting server on %s:%d", listenAddr, listenPort)
	log.Fatal(http.ListenAndServe(fmt.Sprintf("%s:%d", listenAddr, listenPort), nil))
}