omada-bandwidth/main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"git.tordarus.net/tordarus/channel"
	"git.tordarus.net/tordarus/envvars"
	"git.tordarus.net/tordarus/gmath"
	omada "git.tordarus.net/tordarus/omada-api"
	omadamodel "git.tordarus.net/tordarus/omada-api/model"
	"git.tordarus.net/tordarus/slices"
)

var ( // flags
	FlagSiteNames = envvars.StringSlice("SITES", ",", []string{})

	FlagOmadaURL          = envvars.String("OMADA_URL", "http://localhost:8088")
	FlagOmadaID           = envvars.String("OMADA_ID", "")
	FlagOmadaClientID     = envvars.String("OMADA_CLIENT_ID", "")
	FlagOmadaClientSecret = envvars.String("OMADA_CLIENT_SECRET", "")
	FlagOmadaUsername     = envvars.String("OMADA_USERNAME", "")
	FlagOmadaPassword     = envvars.String("OMADA_PASSWORD", "")
	FlagRefreshInterval   = envvars.Duration("REFRESH_INTERVAL", 30*time.Second)

	FlagHttpInterface = envvars.String("HTTP_INTERFACE", "")
	FlagHttpPort      = envvars.Uint16("HTTP_PORT", 8080)
)

var MinRefreshInterval = 30 * time.Second

var CurrentTraffic = &TrafficStats{}

func main() {
	FlagRefreshInterval = gmath.Max(FlagRefreshInterval, MinRefreshInterval)

	api, err := omada.NewApi(omada.ApiConfig{
		BasePath:     FlagOmadaURL,
		OmadaID:      FlagOmadaID,
		ClientID:     FlagOmadaClientID,
		ClientSecret: FlagOmadaClientSecret,
		Username:     FlagOmadaUsername,
		Password:     FlagOmadaPassword,
	})

	if err != nil {
		panic(err)
	}

	go api.AutoRefresh(context.Background())

	go func() {
		http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			if err := json.NewEncoder(w).Encode(CurrentTraffic); err != nil {
				w.WriteHeader(http.StatusInternalServerError)
			}
		})

		http.ListenAndServe(fmt.Sprintf("%s:%d", FlagHttpInterface, FlagHttpPort), nil)
	}()

	sites := slices.Filter(channel.ToSlice(api.GetSites()), FilterSitesByName(FlagSiteNames...))

	ticker := time.NewTicker(FlagRefreshInterval)

	CalculateSiteTraffic(api, sites, 0)

	lastTick := time.Now()
	for now := range ticker.C {
		CurrentTraffic = CalculateSiteTraffic(api, sites, now.Sub(lastTick))
		fmt.Println(CurrentTraffic)
		lastTick = now
	}
}

func CalculateSiteTraffic(api *omada.Api, sites []*omadamodel.Site, duration time.Duration) *TrafficStats {
	siteTraffics := map[string]*SiteTraffic{}

	for _, site := range sites {
		trafficByClient := map[string]TrafficRate{}

		clients := channel.ToSlice(api.GetClients(site.ID))
		for _, client := range clients {
			traffic := CalculateClientTraffic(site, client, duration)
			trafficByClient[client.Name] = traffic
		}

		siteTraffics[site.Name] = NewSiteTraffic(trafficByClient)
	}

	return NewTrafficStats(siteTraffics)
}

func FilterSitesByName(allowedSiteNames ...string) func(site *omadamodel.Site) bool {
	if len(allowedSiteNames) == 0 {
		return func(site *omadamodel.Site) bool { return true }
	}

	siteNames := slices.ToStructMap(allowedSiteNames)

	return func(site *omadamodel.Site) bool {
		_, ok := siteNames[site.Name]
		return ok
	}
}