Support authentication and auto-UPS-ID-detection

main.go

@@ -1,43 +1,104 @@
 package main
 
 import (
+	"encoding/json"
 	"fmt"
 	"log"
 	"os"
 	"os/exec"
-	"strconv"
 	"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() {
-	if len(os.Args) != 5 {
+	var config configT
+	if len(os.Args) < 2 {
 		printUsageInfo()
 	}
-	intervalInt, err := strconv.Atoi(os.Args[3])
-	if err != nil {
-		other.PrintError("Failed to parse scan interval as integer: "+err.Error(), 2)
+	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(intervalInt) * time.Second
-	shutdownMultiplier, err := strconv.Atoi(os.Args[4])
+
+	intervalDur := time.Duration(config.ScanIntervalSeconds) * time.Second
+	client, err := nut.Dial(config.IP + ":" + config.Port)
 	if err != nil {
-		other.PrintError("Failed to parse bad scan limit as integer: "+err.Error(), 3)
+		other.PrintError("Failed to dial NUT server ("+os.Args[1]+"): "+err.Error(), 7)
 	}
-	client, err := nut.Dial(os.Args[1])
+
+	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 dial NUT server ("+os.Args[1]+"): "+err.Error(), 4)
+		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(os.Args[2], "ups.status")
+		value, err = client.GetVar("ups.status")
 		if err != nil {
-			other.PrintError("Failed to get status from \""+os.Args[2]+"\": "+err.Error(), 5)
+			other.PrintError("Failed to get status from \""+os.Args[2]+"\": "+err.Error(), 10)
 		}
 
 		switch value {
@@ -47,17 +108,17 @@ func main() {
 		default:
 			badCycles = 0
 		}
-		if badCycles >= shutdownMultiplier {
+		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(), 6)
+				other.PrintError("Failed to execute shutdown command: "+err.Error(), 11)
 			}
 		}
 	}
 }
 
 func printUsageInfo() {
-	fmt.Print("\nUsage: unifi-shutd <host (ip:port)> <ups id> <scan interval (seconds)> <bad scan limit (int)>\n\n")
+	fmt.Printf("Usage: %s [start | config]\n", os.Args[0])
 	os.Exit(0)
 }