From 853eaa38904863cd7e720dede4784fa3a0fc4b9f Mon Sep 17 00:00:00 2001
From: Jan De Dobbeleer <jan.de.dobbeleer@gmail.com>
Date: Mon, 15 Feb 2021 22:01:09 +0100
Subject: [PATCH] feat(session): support text template

relates to #376
---
 docs/docs/segment-session.md |  3 ++
 src/segment_session.go       | 14 +++++-
 src/segment_session_test.go  | 84 ++++++++++++++++++++++++++++++++++++
 3 files changed, 99 insertions(+), 2 deletions(-)

diff --git a/docs/docs/segment-session.md b/docs/docs/segment-session.md
index 2ee32bd9..6867b32a 100644
--- a/docs/docs/segment-session.md
+++ b/docs/docs/segment-session.md
@@ -33,11 +33,14 @@ to `\uF817 `
 - default_user_name: `string` - name of the default user - defaults to empty
 - display_default_user: `boolean` - display the segment or not when the user matches `default_user_name` - defaults
 to `true`
+- template: `string` - A go [text/template][go-text-template] template extended with [sprig][sprig] utilizing the
+properties below. Only used when a value is set, making the above properties obsolete.
 
 ## Template Properties
 
 - `.UserName`: `string` - the current user's name
 - `.ComputerName`: `string` - the current computer's name
 - `.SSHSession`: `boolean` - active SSH session or not
+- `.Root`: `boolean` - are you a root/admin user or not
 
 [colors]: /docs/configure#colors
diff --git a/src/segment_session.go b/src/segment_session.go
index 7ca9bb90..885707eb 100644
--- a/src/segment_session.go
+++ b/src/segment_session.go
@@ -11,6 +11,7 @@ type session struct {
 	UserName     string
 	ComputerName string
 	SSHSession   bool
+	Root         bool
 }
 
 const (
@@ -53,12 +54,22 @@ func (s *session) init(props *properties, env environmentInfo) {
 
 func (s *session) getFormattedText() string {
 	s.ComputerName = s.getComputerName()
+	s.SSHSession = s.activeSSHSession()
+	segmentTemplate := s.props.getString(SegmentTemplate, "")
+	if segmentTemplate != "" {
+		s.Root = s.env.isRunningAsRoot()
+		template := &textTemplate{
+			Template: segmentTemplate,
+			Context:  s,
+		}
+		return template.render()
+	}
 	separator := ""
 	if s.props.getBool(DisplayHost, true) && s.props.getBool(DisplayUser, true) {
 		separator = s.props.getString(UserInfoSeparator, "@")
 	}
 	var sshIcon string
-	if s.activeSSHSession() {
+	if s.SSHSession {
 		sshIcon = s.props.getString(SSHIcon, "\uF817 ")
 	}
 	userColor := s.props.getColor(UserColor, s.props.foreground)
@@ -97,7 +108,6 @@ func (s *session) activeSSHSession() bool {
 	for _, key := range keys {
 		content := s.env.getenv(key)
 		if content != "" {
-			s.SSHSession = true
 			return true
 		}
 	}
diff --git a/src/segment_session_test.go b/src/segment_session_test.go
index d1781ad5..45addd63 100644
--- a/src/segment_session_test.go
+++ b/src/segment_session_test.go
@@ -151,3 +151,87 @@ func TestActiveSSHSessionActiveBoth(t *testing.T) {
 	}
 	assert.True(t, s.activeSSHSession())
 }
+
+func TestSessionSegmentTemplate(t *testing.T) {
+	cases := []struct {
+		Case           string
+		ExpectedString string
+		UserName       string
+		ComputerName   string
+		SSHSession     bool
+		Root           bool
+		Template       string
+	}{
+		{
+			Case:           "user and computer",
+			ExpectedString: "john@company-laptop",
+			ComputerName:   "company-laptop",
+			UserName:       "john",
+			Template:       "{{.UserName}}{{if .ComputerName}}@{{.ComputerName}}{{end}}",
+		},
+		{
+			Case:           "user only",
+			ExpectedString: "john",
+			UserName:       "john",
+			Template:       "{{.UserName}}{{if .ComputerName}}@{{.ComputerName}}{{end}}",
+		},
+		{
+			Case:           "user with ssh",
+			ExpectedString: "john on remote",
+			UserName:       "john",
+			SSHSession:     true,
+			ComputerName:   "remote",
+			Template:       "{{.UserName}}{{if .SSHSession}} on {{.ComputerName}}{{end}}",
+		},
+		{
+			Case:           "user without ssh",
+			ExpectedString: "john",
+			UserName:       "john",
+			SSHSession:     false,
+			ComputerName:   "remote",
+			Template:       "{{.UserName}}{{if .SSHSession}} on {{.ComputerName}}{{end}}",
+		},
+		{
+			Case:           "user with root and ssh",
+			ExpectedString: "super john on remote",
+			UserName:       "john",
+			SSHSession:     true,
+			ComputerName:   "remote",
+			Root:           true,
+			Template:       "{{if .Root}}super {{end}}{{.UserName}}{{if .SSHSession}} on {{.ComputerName}}{{end}}",
+		},
+		{
+			Case:           "no template",
+			ExpectedString: "\uf817 <>john</>@<>remote</>",
+			UserName:       "john",
+			SSHSession:     true,
+			ComputerName:   "remote",
+			Root:           true,
+		},
+	}
+
+	for _, tc := range cases {
+		env := new(MockedEnvironment)
+		env.On("getCurrentUser", nil).Return(tc.UserName)
+		env.On("getRuntimeGOOS", nil).Return("burp")
+		env.On("getHostName", nil).Return(tc.ComputerName, nil)
+		var SSHSession string
+		if tc.SSHSession {
+			SSHSession = "zezzion"
+		}
+		env.On("getenv", "SSH_CONNECTION").Return(SSHSession)
+		env.On("getenv", "SSH_CLIENT").Return(SSHSession)
+		env.On("isRunningAsRoot", nil).Return(tc.Root)
+		props := &properties{
+			values: map[Property]interface{}{
+				SegmentTemplate: tc.Template,
+			},
+		}
+		session := &session{
+			env:   env,
+			props: props,
+		}
+		_ = session.enabled()
+		assert.Equal(t, tc.ExpectedString, session.string(), tc.Case)
+	}
+}