From a1a2f86a2d6c970366a55e04c46529c657c3630b Mon Sep 17 00:00:00 2001
From: Thomas FOUBERT <perso@thomas-foubert.com>
Date: Fri, 24 May 2024 15:10:04 -0400
Subject: [PATCH] fix: this should do the trick and user inputs shouldn't have
 any \n left

---
 cmd/agent/main.go           |  8 ++++----
 config/agent_live_change.go |  8 +++++---
 go.mod                      |  1 +
 go.sum                      |  2 ++
 utils/strings.go            | 16 ++++++++++++++++
 5 files changed, 28 insertions(+), 7 deletions(-)
 create mode 100644 utils/strings.go

diff --git a/cmd/agent/main.go b/cmd/agent/main.go
index 19a2b2f..12dbab8 100644
--- a/cmd/agent/main.go
+++ b/cmd/agent/main.go
@@ -3,11 +3,11 @@ package main
 import (
 	"fmt"
 	"os"
-	"strings"
 
 	"github.com/equals215/deepsentinel/agent"
 	"github.com/equals215/deepsentinel/config"
 	"github.com/equals215/deepsentinel/daemonize"
+	"github.com/equals215/deepsentinel/utils"
 	"github.com/spf13/cobra"
 )
 
@@ -66,21 +66,21 @@ func installCmd(rootCmd *cobra.Command) {
 
 			fmt.Print("Input the server address (ex.: http://localhost:8080): ")
 			fmt.Scanln(&serverAddress)
-			err = config.AgentSetServerAddress(strings.ReplaceAll(serverAddress, "\n", ""))
+			err = config.AgentSetServerAddress(utils.CleanString(serverAddress))
 			if err != nil {
 				return err
 			}
 
 			fmt.Print("Input the auth token: ")
 			fmt.Scanln(&authToken)
-			err = config.AgentSetAuthToken(strings.ReplaceAll(authToken, "\n", ""))
+			err = config.AgentSetAuthToken(utils.CleanString(authToken))
 			if err != nil {
 				return err
 			}
 
 			fmt.Print("Input the machine name: ")
 			fmt.Scanln(&machineName)
-			err = config.AgentSetMachineName(strings.ReplaceAll(machineName, "\n", ""))
+			err = config.AgentSetMachineName(utils.CleanString(machineName))
 			if err != nil {
 				return err
 			}
diff --git a/config/agent_live_change.go b/config/agent_live_change.go
index 4b1fddd..5731558 100644
--- a/config/agent_live_change.go
+++ b/config/agent_live_change.go
@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"net/url"
 
+	"github.com/equals215/deepsentinel/utils"
+	"github.com/mrz1836/go-sanitize"
 	"github.com/spf13/viper"
 )
 
@@ -20,7 +22,7 @@ func AgentSetServerAddress(args ...any) error {
 		return fmt.Errorf("too many arguments")
 	}
 
-	address := args[0].(string)
+	address := sanitize.URL(args[0].(string))
 	url, err := url.Parse(address)
 	if err != nil {
 		return fmt.Errorf("Invalid URL: %s", err)
@@ -49,7 +51,7 @@ func AgentSetAuthToken(args ...any) error {
 		return fmt.Errorf("too many arguments")
 	}
 
-	token := args[0].(string)
+	token := utils.CleanString(args[0].(string))
 
 	Agent.Lock()
 	viper.Set("auth-token", token)
@@ -71,7 +73,7 @@ func AgentSetMachineName(args ...any) error {
 		return fmt.Errorf("too many arguments")
 	}
 
-	name := args[0].(string)
+	name := utils.CleanString(args[0].(string))
 
 	Agent.Lock()
 	viper.Set("machine-name", name)
diff --git a/go.mod b/go.mod
index 23789e3..bac10f0 100644
--- a/go.mod
+++ b/go.mod
@@ -10,6 +10,7 @@ require (
 	github.com/grongor/panicwatch v1.2.0
 	github.com/jxsl13/osfacts v0.4.0
 	github.com/kristinjeanna/redact v1.0.0
+	github.com/mrz1836/go-sanitize v1.3.2
 	github.com/sirupsen/logrus v1.9.3
 	github.com/spf13/cobra v1.8.0
 	github.com/spf13/pflag v1.0.5
diff --git a/go.sum b/go.sum
index 3ffb976..95d936b 100644
--- a/go.sum
+++ b/go.sum
@@ -70,6 +70,8 @@ github.com/mattn/go-runewidth v0.0.15 h1:UNAjwbU9l54TA3KzvqLGxwWjHmMgBUVhBiTjelZ
 github.com/mattn/go-runewidth v0.0.15/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
 github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=
 github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo=
+github.com/mrz1836/go-sanitize v1.3.2 h1:sGhusPxP4L+7NAUVAUl/WrrBazNSNECsvYh3yumuJXQ=
+github.com/mrz1836/go-sanitize v1.3.2/go.mod h1:wvRS2ALFDxOCK3ORQPwKUxl7HTIBUV8S3U34Hwn96r4=
 github.com/pelletier/go-toml/v2 v2.1.0 h1:FnwAJ4oYMvbT/34k9zzHuZNrhlz48GB3/s6at6/MHO4=
 github.com/pelletier/go-toml/v2 v2.1.0/go.mod h1:tJU2Z3ZkXwnxa4DPO899bsyIoywizdUvyaeZurnPPDc=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
diff --git a/utils/strings.go b/utils/strings.go
new file mode 100644
index 0000000..548ab4e
--- /dev/null
+++ b/utils/strings.go
@@ -0,0 +1,16 @@
+package utils
+
+import (
+	"strings"
+	"unicode"
+)
+
+// CleanString removes all non-graphic characters from a string
+func CleanString(dirty string) string {
+	return strings.Map(func(r rune) rune {
+		if unicode.IsGraphic(r) {
+			return r
+		}
+		return -1
+	}, dirty)
+}