add availability and usage threhold option

Signed-off-by: Simon Schneider <[email protected]>

Author: sschne

main.go

@@ -5,6 +5,7 @@ import (
 	"flag"
 	"fmt"
 	"os"
+	"strconv"
 	"strings"
 
 	wildcard "git.iglou.eu/Imported/go-wildcard"
@@ -37,6 +38,9 @@ var (
 	themeOpt = flag.String("theme", defaultThemeName(), "color themes: dark, light")
 	styleOpt = flag.String("style", defaultStyleName(), "style: unicode, ascii")
 
+	availThreshold = flag.String("avail-threshold", "10G,1G", "specifies the coloring threshold (yellow, red) of the avail column, must be integer with optional SI prefixes")
+	usageThreshold = flag.String("usage-threshold", "0.5,0.9", "specifies the coloring threshold (yellow, red) of the usage bars as a floating point number from 0 to 1")
+
 	inodes     = flag.Bool("inodes", false, "list inode information instead of block usage")
 	jsonOutput = flag.Bool("json", false, "output all devices in JSON format")
 	warns      = flag.Bool("warnings", false, "output all warnings to STDERR")
@@ -243,6 +247,34 @@ func main() {
 		m = mounts
 	}
 
+	// validate availability thresholds
+	availbilityThresholds := strings.Split(*availThreshold, ",")
+	if len(availbilityThresholds) != 2 {
+		fmt.Fprintln(os.Stderr, fmt.Errorf("error parsing avail-threshold: invalid option '%s'", *availThreshold))
+		os.Exit(1)
+	}
+	for _, thresold := range availbilityThresholds {
+		_, err = stringToSize(thresold)
+		if err != nil {
+			fmt.Fprintln(os.Stderr, "error parsing avail-threshold:", err)
+			os.Exit(1)
+		}
+	}
+
+	// validate usage thresholds
+	usageThresholds := strings.Split(*usageThreshold, ",")
+	if len(usageThresholds) != 2 {
+		fmt.Fprintln(os.Stderr, fmt.Errorf("error parsing usage-threshold: invalid option '%s'", *usageThreshold))
+		os.Exit(1)
+	}
+	for _, thresold := range usageThresholds {
+		_, err = strconv.ParseFloat(thresold, 64)
+		if err != nil {
+			fmt.Fprintln(os.Stderr, "error parsing usage-threshold:", err)
+			os.Exit(1)
+		}
+	}
+
 	// print out warnings
 	if *warns {
 		for _, warning := range warnings {

table.go

@@ -3,6 +3,8 @@ package main
 import (
 	"fmt"
 	"os"
+	"regexp"
+	"strconv"
 	"strings"
 
 	"github.com/jedib0t/go-pretty/v6/table"
@@ -153,10 +155,12 @@ func spaceTransformer(val interface{}) string {
 	free := val.(uint64)
 
 	var s = termenv.String(sizeToString(free))
+	redAvail, _ := stringToSize(strings.Split(*availThreshold, ",")[1])
+	yellowAvail, _ := stringToSize(strings.Split(*availThreshold, ",")[0])
 	switch {
-	case free < 1<<30:
+	case free < redAvail:
 		s = s.Foreground(theme.colorRed)
-	case free < 10*1<<30:
+	case free < yellowAvail:
 		s = s.Foreground(theme.colorYellow)
 	default:
 		s = s.Foreground(theme.colorGreen)
@@ -183,10 +187,12 @@ func barTransformer(val interface{}) string {
 	}
 
 	// apply color to progress-bar
+	redUsage, _ := strconv.ParseFloat(strings.Split(*usageThreshold, ",")[1], 64)
+	yellowUsage, _ := strconv.ParseFloat(strings.Split(*usageThreshold, ",")[0], 64)
 	switch {
-	case usage >= 0.9:
+	case usage >= redUsage:
 		s = s.Foreground(theme.colorRed)
-	case usage >= 0.5:
+	case usage >= yellowUsage:
 		s = s.Foreground(theme.colorYellow)
 	default:
 		s = s.Foreground(theme.colorGreen)
@@ -260,6 +266,43 @@ func sizeToString(size uint64) (str string) {
 	return
 }
 
+// stringToSize transforms an SI size into a number.
+func stringToSize(s string) (size uint64, err error) {
+	regex := regexp.MustCompile(`^(\d+)([KMGTPE]?)$`)
+	matches := regex.FindStringSubmatch(s)
+	if len(matches) == 0 {
+		return 0, fmt.Errorf("'%s' is not valid, must have integer with optional SI prefix", s)
+	}
+
+	num, err := strconv.ParseUint(matches[1], 10, 64)
+	if err != nil {
+		return 0, err
+	}
+	if matches[2] != "" {
+		prefix := matches[2]
+		switch {
+		case prefix == "K":
+			size = num << 10
+		case prefix == "M":
+			size = num << 20
+		case prefix == "G":
+			size = num << 30
+		case prefix == "T":
+			size = num << 40
+		case prefix == "P":
+			size = num << 50
+		case prefix == "E":
+			size = num << 60
+		default:
+			err = fmt.Errorf("prefix '%s' not allowed, valid prefixes are K, M, G, T, P, E", prefix)
+			return
+		}
+	} else {
+		size = num
+	}
+	return
+}
+
 // stringToColumn converts a column name to its index.
 func stringToColumn(s string) (int, error) {
 	s = strings.ToLower(s)