feat: aws segment

src/properties.go

@@ -23,6 +23,8 @@ const (
 	DisplayVersion Property = "display_version"
 	// AlwaysEnabled decides whether or not to always display the info
 	AlwaysEnabled Property = "always_enabled"
+	// SegmentTemplate is the template to use to render the information
+	SegmentTemplate Property = "template"
 )
 
 type properties struct {

src/segment.go

@@ -90,6 +90,8 @@ const (
 	ExecutionTime SegmentType = "executiontime"
 	// Ruby writes which ruby version is currently active
 	Ruby SegmentType = "ruby"
+	// Aws writes the active aws context
+	Aws SegmentType = "aws"
 )
 
 func (segment *Segment) string() string {
@@ -147,6 +149,7 @@ func (segment *Segment) mapSegmentWithWriter(env environmentInfo) error {
 		YTM:           &ytm{},
 		ExecutionTime: &executiontime{},
 		Ruby:          &ruby{},
+		Aws:           &aws{},
 	}
 	if writer, ok := functions[segment.Type]; ok {
 		props := &properties{

src/segment_aws.go

@@ -0,0 +1,81 @@
+package main
+
+import (
+	"fmt"
+	"strings"
+)
+
+type aws struct {
+	props   *properties
+	env     environmentInfo
+	Profile string
+	Region  string
+}
+
+const (
+	defaultUser = "default"
+)
+
+func (a *aws) init(props *properties, env environmentInfo) {
+	a.props = props
+	a.env = env
+}
+
+func (a *aws) enabled() bool {
+	getEnvFirstMatch := func(envs ...string) string {
+		for _, env := range envs {
+			value := a.env.getenv(env)
+			if value != "" {
+				return value
+			}
+		}
+		return ""
+	}
+	a.Profile = getEnvFirstMatch("AWS_VAULT", "AWS_PROFILE")
+	a.Region = getEnvFirstMatch("AWS_DEFAULT_REGION", "AWS_REGION")
+	if a.Profile != "" && a.Region != "" {
+		return true
+	}
+	if a.Profile == "" && a.Region != "" {
+		a.Profile = defaultUser
+		return true
+	}
+	a.getConfigFileInfo()
+	return a.Profile != ""
+}
+
+func (a *aws) getConfigFileInfo() {
+	configPath := a.env.getenv("AWS_CONFIG_FILE")
+	if configPath == "" {
+		configPath = fmt.Sprintf("%s/.aws/config", a.env.homeDir())
+	}
+	config := a.env.getFileContent(configPath)
+	configSection := "[default]"
+	if a.Profile != "" {
+		configSection = fmt.Sprintf("[profile %s]", a.Profile)
+	}
+	configLines := strings.Split(config, "\n")
+	var sectionActive bool
+	for _, line := range configLines {
+		if strings.HasPrefix(line, configSection) {
+			sectionActive = true
+			continue
+		}
+		if sectionActive && strings.HasPrefix(line, "region") {
+			a.Region = strings.TrimSpace(strings.Split(line, "=")[1])
+			break
+		}
+	}
+	if a.Profile == "" && a.Region != "" {
+		a.Profile = defaultUser
+	}
+}
+
+func (a *aws) string() string {
+	segmentTemplate := a.props.getString(SegmentTemplate, "{{.Profile}}{{if .Region}}@{{.Region}}{{end}}")
+	template := &textTemplate{
+		Template: segmentTemplate,
+		Context:  a,
+	}
+	return template.render()
+}

src/segment_aws_test.go

@@ -0,0 +1,65 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAWSSegment(t *testing.T) {
+	cases := []struct {
+		Case            string
+		ExpectedString  string
+		ExpectedEnabled bool
+		Profile         string
+		Vault           string
+		Region          string
+		DefaultRegion   string
+		ConfigFile      string
+		Template        string
+	}{
+		{Case: "disabled", ExpectedString: "", ExpectedEnabled: false},
+		{Case: "enabled no region", ExpectedString: "company", ExpectedEnabled: true, Profile: "company"},
+		{Case: "enabled with region", ExpectedString: "company@eu-west", ExpectedEnabled: true, Profile: "company", Region: "eu-west"},
+		{
+			Case:            "template: enabled no region",
+			ExpectedString:  "profile: company",
+			ExpectedEnabled: true,
+			Profile:         "company",
+			Template:        "profile: {{.Profile}}{{if .Region}} in {{.Region}}{{end}}",
+		},
+		{
+			Case:            "template: enabled with region",
+			ExpectedString:  "profile: company in eu-west",
+			ExpectedEnabled: true,
+			Profile:         "company",
+			Region:          "eu-west",
+			Template:        "profile: {{.Profile}}{{if .Region}} in {{.Region}}{{end}}",
+		},
+		{Case: "template: invalid", ExpectedString: invalidTemplate, ExpectedEnabled: true, Profile: "c", Template: "{{ .Burp"},
+	}
+
+	for _, tc := range cases {
+		env := new(MockedEnvironment)
+		env.On("getenv", "AWS_VAULT").Return(tc.Vault)
+		env.On("getenv", "AWS_PROFILE").Return(tc.Profile)
+		env.On("getenv", "AWS_REGION").Return(tc.Region)
+		env.On("getenv", "AWS_DEFAULT_REGION").Return(tc.DefaultRegion)
+		env.On("getenv", "AWS_CONFIG_FILE").Return(tc.ConfigFile)
+		env.On("getFileContent", "/usr/home/.aws/config").Return("")
+		env.On("homeDir", nil).Return("/usr/home")
+		props := &properties{
+			values: map[Property]interface{}{},
+		}
+		if tc.Template != "" {
+			props.values[SegmentTemplate] = tc.Template
+		}
+
+		aws := &aws{
+			env:   env,
+			props: props,
+		}
+		assert.Equal(t, tc.ExpectedEnabled, aws.enabled(), tc.Case)
+		assert.Equal(t, tc.ExpectedString, aws.string(), tc.Case)
+	}
+}