skolo/go/censustaker/censustaker.go

package main

import (
	"context"
	"flag"
	"fmt"
	"os"

	"go.skia.org/infra/go/sklog"
)

var (
	scriptDir     = flag.String("script_dir", "/usr/local/share/trooper_tools/censustaker/", "Path in which the ansible scripts and configurations are stored. This will be the working dir when executing the ansible scripts to enumerate all bots on the network.")
	ansibleOutput = flag.String("ansible_out", "/tmp/census_output", "File in which the ansible script should dump its intermediary output.")

	switchAddress = flag.String("switch_address", "", "The IP address of the switch to pull the port numbers from.")
)

func enumerateBots(ctx context.Context, names BotNameGetter, ports BotPortGetter) ([]Bot, error) {
	nameList, err := names.GetBotNamesAddresses(ctx)
	if err != nil {
		return nil, fmt.Errorf("Could not fetch bot names and mac addresses: %s", err)
	}

	portList, err := ports.GetBotPortsAddresses()
	if err != nil {
		return nil, fmt.Errorf("Could not fetch bot ports and mac addresses: %s", err)
	}

	SENTINAL_PORT := -1
	botMap := map[string]Bot{}
	for _, b := range nameList {
		if b.MACAddress != "" {
			b.Port = SENTINAL_PORT
			botMap[b.MACAddress] = b
		}
	}
	for _, b := range portList {
		if _, ok := botMap[b.MACAddress]; ok && b.MACAddress != "" {
			a := botMap[b.MACAddress]
			a.Port = b.Port
			botMap[b.MACAddress] = a
		}
	}

	botList := []Bot{}
	for _, b := range botMap {
		if b.Port != SENTINAL_PORT {
			botList = append(botList, b)
		}
	}

	return botList, nil
}

func main() {
	flag.Parse()

	if *scriptDir == "" || *switchAddress == "" {
		sklog.Fatal("--script_dir and --switch_address cannot be empty")
	}

	if _, err := os.Stat(*scriptDir); os.IsNotExist(err) {
		sklog.Fatalf("--script_dir %s points to a non-existent directory", *scriptDir)
	}

	ctx := context.Background()
	ansible := NewAnsibleBotNameGetter(*scriptDir, *ansibleOutput)
	edgeswitch := NewEdgeSwitchBotPortGetter(*switchAddress)
	bots, err := enumerateBots(ctx, ansible, edgeswitch)
	sklog.Infof("Found bots %v (With err %v)", bots, err)
}