unifi-shutd/main.go

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"time"

	"github.com/rwinkhart/go-boilerplate/front"
	"github.com/rwinkhart/go-boilerplate/other"
	nut "github.com/rwinkhart/uni-nut"
)

const configPath = "/etc/unifi-shutd.json"

type configT struct {
	IP                  string  `json:"ip"`
	Port                string  `json:"port"`
	Username            *string `json:"username"`
	Password            *string `json:"password"`
	ScanIntervalSeconds int     `json:"scan_interval_seconds"`
	BadScanLimit        int     `json:"bad_scan_limit"`
}

func main() {
	var config configT
	if len(os.Args) < 2 {
		printUsageInfo()
	}
	switch os.Args[1] {
	case "start":
		configBytes, err := os.ReadFile(configPath)
		if err != nil {
			other.PrintError("Failed to read ("+configPath+"): "+err.Error(), 1)
		}
		err = json.Unmarshal(configBytes, &config)
		if err != nil {
			other.PrintError("Failed to unmarshal config file: "+err.Error(), 2)
		}
	case "config":
		if os.Geteuid() != 0 {
			fmt.Printf("\"%s config\" must be run as root\n", os.Args[0])
			os.Exit(0)
		}
		var username, password *string
		if front.InputBinary("Does your NUT server use authentication (username/password)?") {
			usernameStr := front.Input("NUT Username:")
			passwordStr := string(front.InputHidden("NUT Password:"))
			username = &usernameStr
			password = &passwordStr
		}
		config = configT{
			IP:                  front.Input("NUT IP:"),
			Port:                front.Input("NUT Port:"),
			Username:            username,
			Password:            password,
			ScanIntervalSeconds: front.InputInt("Scan Interval (seconds):", 1, -1),
			BadScanLimit:        front.InputInt("Bad Scan Limit:", 1, -1),
		}
		configBytes, err := json.Marshal(config)
		if err != nil {
			other.PrintError("Failed to marshal config file: "+err.Error(), 3)
		}
		err = os.WriteFile(configPath, configBytes, 0600)
		if err != nil {
			other.PrintError("Failed to write ("+configPath+"): "+err.Error(), 4)
		}
		os.Exit(0)
	default:
		printUsageInfo()
	}

	intervalDur := time.Duration(config.ScanIntervalSeconds) * time.Second
	client, err := nut.Dial(config.IP + ":" + config.Port)
	if err != nil {
		other.PrintError("Failed to dial NUT server ("+os.Args[1]+"): "+err.Error(), 7)
	}

	if config.Username != nil {
		err = client.Authenticate(*config.Username, *config.Password)
		if err != nil {
			other.PrintError("Failed to authenticate with NUT server: "+err.Error(), 8)
		}
	}

	err = client.AutomaticallySetID()
	if err != nil {
		other.PrintError("Failed to detect UPS ID: "+err.Error(), 9)
	}

	log.Println("Starting NUT monitoring loop...")
	cmd := exec.Command("/sbin/poweroff")
	var value string
	var badCycles int
	for {
		time.Sleep(intervalDur)
		value, err = client.GetVar("ups.status")
		if err != nil {
			other.PrintError("Failed to get status from \""+os.Args[2]+"\": "+err.Error(), 10)
		}

		switch value {
		case "OB DISCHRG":
			badCycles++
			log.Printf("BAD cycle: %d", badCycles)
		default:
			badCycles = 0
		}
		if badCycles >= config.BadScanLimit {
			log.Println("Bad scan limit reached; initiating shutdown...")
			err := cmd.Run()
			if err != nil {
				other.PrintError("Failed to execute shutdown command: "+err.Error(), 11)
			}
		}
	}
}

func printUsageInfo() {
	fmt.Printf("Usage: %s [start | config]\n", os.Args[0])
	os.Exit(0)
}