From efae7c4b93c4ccbb3fbc9a7c0bc478b40c94891c Mon Sep 17 00:00:00 2001
From: CapCreeperGR <nickvavour@gmail.com>
Date: Sun, 9 Jun 2024 10:49:11 +0300
Subject: [PATCH] Split main.go code into more functions,improved extra newline
 removal and added --ascii flag

---
 main.go  | 117 +++++++++++++++++++++++++++++++------------------------
 utils.go |   2 +-
 2 files changed, 67 insertions(+), 52 deletions(-)

diff --git a/main.go b/main.go
index dee803f..a95a8ef 100644
--- a/main.go
+++ b/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"flag"
 	"fmt"
 	yaml "gopkg.in/yaml.v3"
 	"log"
@@ -33,6 +34,9 @@ type StormfetchConfig struct {
 
 func main() {
 	readConfig()
+	readFlags()
+	checkDependencies()
+	runStormfetch()
 }
 
 func readConfig() {
@@ -72,6 +76,14 @@ func readConfig() {
 	} else {
 		log.Fatalf("Fetch script file not found: %s", err.Error())
 	}
+}
+
+func readFlags() {
+	flag.StringVar(&config.Ascii, "ascii", config.Ascii, "help message for flagname")
+	flag.Parse()
+}
+
+func checkDependencies() {
 	// Show Dependency warning if enabled
 	if config.DependencyWarning {
 		var dependencies []string
@@ -89,6 +101,58 @@ func readConfig() {
 			fmt.Println("You can disable this warning through your stormfetch config")
 		}
 	}
+}
+
+func SetupFetchEnv() []string {
+	var env = make(map[string]string)
+	env["DISTRO_LONG_NAME"] = getDistroInfo().LongName
+	env["DISTRO_SHORT_NAME"] = getDistroInfo().ShortName
+	env["CPU_MODEL"] = getCPUName()
+	env["CPU_THREADS"] = strconv.Itoa(getCPUThreads())
+	memory := GetMemoryInfo()
+	if memory != nil {
+		env["MEM_TOTAL"] = strconv.Itoa(memory.MemTotal)
+		env["MEM_USED"] = strconv.Itoa(memory.MemTotal - memory.MemAvailable)
+		env["MEM_FREE"] = strconv.Itoa(memory.MemAvailable)
+	}
+	partitions := getMountedPartitions()
+	if len(partitions) != 0 {
+		env["MOUNTED_PARTITIONS"] = strconv.Itoa(len(partitions))
+		for i, part := range partitions {
+			env["PARTITION"+strconv.Itoa(i+1)+"_DEVICE"] = part.Device
+			env["PARTITION"+strconv.Itoa(i+1)+"_MOUNTPOINT"] = part.MountPoint
+			if part.Label != "" {
+				env["PARTITION"+strconv.Itoa(i+1)+"_LABEL"] = part.Label
+			}
+			env["PARTITION"+strconv.Itoa(i+1)+"_TOTAL_SIZE"] = FormatBytes(part.TotalSize)
+			env["PARTITION"+strconv.Itoa(i+1)+"_USED_SIZE"] = FormatBytes(part.UsedSize)
+			env["PARTITION"+strconv.Itoa(i+1)+"_FREE_SIZE"] = FormatBytes(part.FreeSize)
+		}
+	}
+	env["DE_WM"] = GetDEWM()
+	env["USER_SHELL"] = GetShell()
+	env["DISPLAY_PROTOCOL"] = GetDisplayProtocol()
+	monitors := getMonitorResolution()
+	if len(monitors) != 0 {
+		env["CONNECTED_MONITORS"] = strconv.Itoa(len(monitors))
+		for i, monitor := range monitors {
+			env["MONITOR"+strconv.Itoa(i+1)] = monitor
+		}
+	}
+	if getGPUName() != "" {
+		env["GPU_MODEL"] = getGPUName()
+	}
+
+	var ret = make([]string, len(env))
+	i := 0
+	for key, value := range env {
+		ret[i] = fmt.Sprintf("%s=%s", key, value)
+		i++
+	}
+	return ret
+}
+
+func runStormfetch() {
 	// Fetch ascii art and apply colors
 	colorMap := make(map[string]string)
 	colorMap["C0"] = "\033[0m"
@@ -183,55 +247,6 @@ func readConfig() {
 		}
 		final += lastAsciiColor + line + "\n"
 	}
-	final = strings.TrimSpace(final)
-	fmt.Print(final + "\033[0m")
-}
-
-func SetupFetchEnv() []string {
-	var env = make(map[string]string)
-	env["DISTRO_LONG_NAME"] = getDistroInfo().LongName
-	env["DISTRO_SHORT_NAME"] = getDistroInfo().ShortName
-	env["CPU_MODEL"] = getCPUName()
-	env["CPU_THREADS"] = strconv.Itoa(getCPUThreads())
-	memory := GetMemoryInfo()
-	if memory != nil {
-		env["MEM_TOTAL"] = strconv.Itoa(memory.MemTotal)
-		env["MEM_USED"] = strconv.Itoa(memory.MemTotal - memory.MemAvailable)
-		env["MEM_FREE"] = strconv.Itoa(memory.MemAvailable)
-	}
-	partitions := getMountedPartitions()
-	if len(partitions) != 0 {
-		env["MOUNTED_PARTITIONS"] = strconv.Itoa(len(partitions))
-		for i, part := range partitions {
-			env["PARTITION"+strconv.Itoa(i+1)+"_DEVICE"] = part.Device
-			env["PARTITION"+strconv.Itoa(i+1)+"_MOUNTPOINT"] = part.MountPoint
-			if part.Label != "" {
-				env["PARTITION"+strconv.Itoa(i+1)+"_LABEL"] = part.Label
-			}
-			env["PARTITION"+strconv.Itoa(i+1)+"_TOTAL_SIZE"] = FormatBytes(part.TotalSize)
-			env["PARTITION"+strconv.Itoa(i+1)+"_USED_SIZE"] = FormatBytes(part.UsedSize)
-			env["PARTITION"+strconv.Itoa(i+1)+"_FREE_SIZE"] = FormatBytes(part.FreeSize)
-		}
-	}
-	env["DE_WM"] = GetDEWM()
-	env["USER_SHELL"] = GetShell()
-	env["DISPLAY_PROTOCOL"] = GetDisplayProtocol()
-	monitors := getMonitorResolution()
-	if len(monitors) != 0 {
-		env["CONNECTED_MONITORS"] = strconv.Itoa(len(monitors))
-		for i, monitor := range monitors {
-			env["MONITOR"+strconv.Itoa(i+1)] = monitor
-		}
-	}
-	if getGPUName() != "" {
-		env["GPU_MODEL"] = getGPUName()
-	}
-
-	var ret = make([]string, len(env))
-	i := 0
-	for key, value := range env {
-		ret[i] = fmt.Sprintf("%s=%s", key, value)
-		i++
-	}
-	return ret
+	final = strings.TrimRight(final, "\n\t ")
+	fmt.Println(final + "\033[0m")
 }
diff --git a/utils.go b/utils.go
index 5fb4746..e5d8e44 100644
--- a/utils.go
+++ b/utils.go
@@ -103,7 +103,7 @@ func getDistroAsciiArt() string {
 		if err != nil {
 			return defaultAscii
 		}
-		return string(bytes)
+		return strings.TrimRight(string(bytes), "\n\t ")
 	} else {
 		return defaultAscii
 	}