tf2stats/main.go

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"hash/fnv"
	"io"
	"net/http"
	"os/exec"
	"strings"
	"time"
	"unicode"

	"github.com/influxdata/line-protocol/v2/lineprotocol"
	log "github.com/sirupsen/logrus"
)

type SteamAPIResponse struct {
	Response struct {
		PlayerCount int `json:"player_count"`
		Result      int `json:"result"`
	} `json:"response"`
}

type TF2Stat struct {
	Name       string  `json:"name"`
	Map        string  `json:"map"`
	Password   bool    `json:"password"`
	MaxPlayers int     `json:"maxplayers"`
	NumPlayers int     `json:"numplayers"`
	QueryPort  int     `json:"queryPort"`
	Ping       int     `json:"ping"`
	Error      *string `json:"error,omitempty"`
}

type TF2Server string

var (
	server = flag.String("server", "", "comma-separated list of addresses to query")
	hashes = flag.String("translate", "", "comma-separated list of maps to hash")
)

func main() {
	flag.Parse()

	if *hashes != "" {
		maps := strings.Split(*hashes, ",")
		for _, tf2map := range maps {
			h := fnv.New32a()
			h.Write([]byte(tf2map))
			mHash := h.Sum32()

			fmt.Printf("%d\t%s\n", mHash, tf2map)
		}

		return
	}

	servers := strings.Split(*server, ",")
	var tServer []TF2Server
	for _, srv := range servers {
		tServer = append(tServer, TF2Server(strings.TrimSpace(srv)))
	}

	totalStat := new(SteamAPIResponse)
	resp, err := http.Get("https://api.steampowered.com/ISteamUserStats/GetNumberOfCurrentPlayers/v1?appid=440")
	if err != nil {
		log.Debugf("error retrieving tf2 player count: %v", err)
	}
	defer resp.Body.Close()
	bResp, err := io.ReadAll(resp.Body)
	if err != nil {
		log.Debugf("error reading tf2 player count response: %v", err)
	}
	err = json.Unmarshal(bResp, &totalStat)
	if err != nil {
		log.Debugf("error parsing tf2 player count response: %v", err)
	}

	if totalStat != nil {
		var globalEnc lineprotocol.Encoder
		globalEnc.StartLine("tf2stats")
		globalEnc.AddField("players", lineprotocol.IntValue(int64(totalStat.Response.PlayerCount)))
		globalEnc.EndLine(time.Now())
		fmt.Print(string(globalEnc.Bytes()))
	}

	for _, ip := range tServer {
		tStat, err := ip.Stats()
		if err != nil {
			log.Debugf("error retrieving stats for %s: %v", ip, err)
			continue
		}

		cleanName := strings.Map(func(r rune) rune {
			if unicode.IsPrint(r) {
				return r
			}
			return -1
		}, tStat.Name)

		var enc lineprotocol.Encoder
		enc.StartLine("tf2server")
		enc.AddTag("address", string(ip))
		enc.AddTag("name", cleanName)
		enc.AddField("player", lineprotocol.IntValue(int64(tStat.NumPlayers)))
		enc.AddField("maxplayer", lineprotocol.IntValue(int64(tStat.MaxPlayers)))
		enc.AddField("ping", lineprotocol.IntValue(int64(tStat.Ping)))

		// hash map to convert it to int
		h := fnv.New32a()
		_, _ = h.Write([]byte(tStat.Map))
		enc.AddField("map", lineprotocol.UintValue(uint64(h.Sum32())))

		enc.EndLine(time.Now())
		fmt.Print(string(enc.Bytes()))
	}
}

func (ip TF2Server) Stats() (*TF2Stat, error) {
	cmd := exec.Command("gamedig", "--type", "teamfortress2", string(ip)) //nolint:gosec
	bOut, err := cmd.CombinedOutput()
	if err != nil {
		return nil, err
	}

	nStat := new(TF2Stat)
	if err = json.Unmarshal(bOut, &nStat); err != nil {
		return nil, err
	}
	if nStat.Error != nil {
		return nil, fmt.Errorf("gamedig returned error: %v", *nStat.Error)
	}

	return nStat, nil
}