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) + } +}