Simplify color map system

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2024 CapCreeperGR
+Copyright (c) 2024 EnumDev
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

config/ascii/bedrock

@@ -0,0 +1,12 @@
+#/7;8;8;7
+${C1}   \\\\\\\\\\\\
+    \\\      \\\
+     \\\      \\\
+      \\\      \\\\\\\\\\\\\\\\\
+       \\\                    \\\
+        \\\                    \\\
+         \\\        ${C2}______      \\\
+          \\\                   ///
+           \\\                 ///
+            \\\               ///
+             \\\////////////////

config/config.yaml

@@ -1,7 +1,5 @@
 distro_ascii: auto
 fetch_script: auto
-ansii_colors: []
-force_config_ansii: false
 show_fs_type: true
 hidden_partitions: []
 # Hiding squashfs prevents snaps from showing up

src/color.go

@@ -0,0 +1,23 @@
+package main
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+)
+
+func setupColorMap(asciiArt string) map[string]string {
+	colorMap := make(map[string]string)
+
+	// Read colors from ascii art
+	if strings.HasPrefix(asciiArt, "#/") {
+		firstLine := strings.Split(asciiArt, "\n")[0]
+		ansiColors := strings.Split(strings.TrimPrefix(firstLine, "#/"), ";")
+
+		for i, ansiColor := range ansiColors {
+			colorMap["C"+strconv.Itoa(i+1)] = fmt.Sprintf("\033[38;5;%sm", ansiColor)
+		}
+	}
+
+	return colorMap
+}

src/main.go

@@ -3,7 +3,6 @@ package main
 import (
 	"flag"
 	"fmt"
-	"gopkg.in/yaml.v3"
 	"log"
 	"os"
 	"os/exec"
@@ -12,6 +11,8 @@ import (
 	"strconv"
 	"strings"
 	"time"
+
+	"gopkg.in/yaml.v3"
 )
 
 var systemConfigDir = "/etc/"
@@ -24,8 +25,6 @@ var TimeTaken = false
 var config = StormfetchConfig{
 	Ascii:            "auto",
 	FetchScript:      "auto",
-	AnsiiColors:      make([]int, 0),
-	ForceConfigAnsii: false,
 	ShowFSType:       false,
 	HiddenPartitions: make([]string, 0),
 	HiddenGPUS:       make([]int, 0),
@@ -102,7 +101,7 @@ func SetupFetchEnv(showTimeTaken bool) []string {
 		env[key] = setter()
 		end := time.Now().UnixMilli()
 		if showTimeTaken {
-			fmt.Println(fmt.Sprintf("Setting '%s' took %d milliseconds", key, end-start))
+			fmt.Printf("Setting '%s' took %d milliseconds\n", key, end-start)
 		}
 	}
 	setVariable("PACKAGES", func() string { return GetInstalledPackages() })
@@ -115,7 +114,7 @@ func SetupFetchEnv(showTimeTaken bool) []string {
 	memory := GetMemoryInfo()
 	end := time.Now().UnixMilli()
 	if showTimeTaken {
-		fmt.Println(fmt.Sprintf("Setting '%s' took %d milliseconds", "MEM_*", end-start))
+		fmt.Printf("Setting '%s' took %d milliseconds\n", "MEM_*", end-start)
 	}
 	if memory != nil {
 		env["MEM_TOTAL"] = strconv.Itoa(memory.MemTotal)
@@ -126,7 +125,7 @@ func SetupFetchEnv(showTimeTaken bool) []string {
 	partitions := GetMountedPartitions(config.HiddenPartitions, config.HiddenFilesystems)
 	end = time.Now().UnixMilli()
 	if showTimeTaken {
-		fmt.Println(fmt.Sprintf("Setting '%s' took %d milliseconds", "PARTITION_*", end-start))
+		fmt.Printf("Setting '%s' took %d milliseconds\n", "PARTITION_*", end-start)
 	}
 	if len(partitions) != 0 {
 		env["MOUNTED_PARTITIONS"] = strconv.Itoa(len(partitions))
@@ -154,7 +153,7 @@ func SetupFetchEnv(showTimeTaken bool) []string {
 	monitors := GetMonitorResolution()
 	end = time.Now().UnixMilli()
 	if showTimeTaken {
-		fmt.Println(fmt.Sprintf("Setting '%s' took %d milliseconds", "MONITOR_*", end-start))
+		fmt.Printf("Setting '%s' took %d milliseconds\n", "MONITOR_*", end-start)
 	}
 	if len(monitors) != 0 {
 		env["CONNECTED_MONITORS"] = strconv.Itoa(len(monitors))
@@ -166,7 +165,7 @@ func SetupFetchEnv(showTimeTaken bool) []string {
 	gpus := GetGPUModels()
 	end = time.Now().UnixMilli()
 	if showTimeTaken {
-		fmt.Println(fmt.Sprintf("Setting '%s' took %d milliseconds", "GPU_*", end-start))
+		fmt.Printf("Setting '%s' took %d milliseconds\n", "GPU_*", end-start)
 	}
 	if len(gpus) != 0 {
 		env["CONNECTED_GPUS"] = strconv.Itoa(len(gpus))
@@ -188,49 +187,23 @@ func SetupFetchEnv(showTimeTaken bool) []string {
 }
 
 func runStormfetch() {
-	// Fetch ascii art and apply colors
-	colorMap := make(map[string]string)
-	colorMap["C0"] = "\033[0m"
-	setColorMap := func() {
-		for i := 0; i < 6; i++ {
-			if i > len(config.AnsiiColors)-1 {
-				colorMap["C"+strconv.Itoa(i+1)] = "\033[0m"
-				continue
-			}
-			colorMap["C"+strconv.Itoa(i+1)] = fmt.Sprintf("\033[1m\033[38;5;%dm", config.AnsiiColors[i])
-		}
-	}
-	setColorMap()
-	ascii := GetDistroAsciiArt()
-	if strings.HasPrefix(ascii, "#/") {
-		firstLine := strings.Split(ascii, "\n")[0]
-		if !config.ForceConfigAnsii {
-			ansiiColors := strings.Split(strings.TrimPrefix(firstLine, "#/"), ";")
-			for i, color := range ansiiColors {
-				atoi, err := strconv.Atoi(color)
-				if err != nil {
-					log.Fatal(err)
-				}
-				if i < len(config.AnsiiColors) {
-					config.AnsiiColors[i] = atoi
-				} else {
-					config.AnsiiColors = append(config.AnsiiColors, atoi)
-				}
-			}
-			setColorMap()
-		}
-		ascii = os.Expand(ascii, func(s string) string {
-			return colorMap[s]
-		})
-		ascii = strings.TrimPrefix(ascii, firstLine+"\n")
-	} else {
-		ascii = os.Expand(ascii, func(s string) string {
+	// Fetch ascii art
+	asciiArt := GetDistroAsciiArt()
+
+	// Setup color map
+	colorMap := setupColorMap(asciiArt)
+	if len(colorMap) > 0 {
+		asciiArt = os.Expand(asciiArt, func(s string) string {
 			return colorMap[s]
 		})
+
+		asciiArt = strings.SplitN(asciiArt, "\n", 2)[1]
 	}
-	asciiSplit := strings.Split(ascii, "\n")
-	asciiNoColor := StripAnsii(ascii)
-	//Execute fetch script
+
+	asciiSplit := strings.Split(asciiArt, "\n")
+	asciiNoColor := StripAnsii(asciiArt)
+
+	// Execute fetch script
 	cmd := exec.Command("/bin/bash", fetchScriptPath)
 	cmd.Dir = path.Dir(fetchScriptPath)
 	cmd.Env = os.Environ()

src/system.go

@@ -1,11 +1,12 @@
 package main
 
 import (
-	"github.com/mitchellh/go-ps"
 	"os"
 	"os/exec"
 	"path"
 	"strings"
+
+	"github.com/mitchellh/go-ps"
 )
 
 type DistroInfo struct {
@@ -24,13 +25,24 @@ func GetDistroInfo() DistroInfo {
 		info.LongName = strings.TrimSpace(config.DistroName)
 		info.ShortName = strings.TrimSpace(config.DistroName)
 	}
-	var releaseMap = make(map[string]string)
-	if _, err := os.Stat("/etc/os-release"); err == nil {
-		releaseMap, err = ReadKeyValueFile("/etc/os-release")
+
+	// Detect release file location
+	var releaseFile string
+	if _, err := os.Stat("/bedrock/etc/os-release"); os.Getenv("BEDROCK_RESTRICT") == "" && err == nil {
+		// Using Bedrock Linux
+		releaseFile = "/bedrock/etc/os-release"
+	} else if _, err := os.Stat("/etc/os-release"); err == nil {
+		// Using a regular linux distribution
+		releaseFile = "/etc/os-release"
+	} else {
+		return info
+	}
+
+	releaseMap, err := ReadKeyValueFile(releaseFile)
 	if err != nil {
 		return info
 	}
-	}
+
 	if id, ok := releaseMap["ID"]; ok {
 		info.ID = id
 	}