From bf44a446be9ff40055924607741e3e220edf47fa Mon Sep 17 00:00:00 2001 From: Jan De Dobbeleer Date: Sun, 7 Feb 2021 10:55:09 +0100 Subject: [PATCH] feat: aws segment --- src/properties.go | 2 + src/segment.go | 3 ++ src/segment_aws.go | 81 +++++++++++++++++++++++++++++++++++++++++ src/segment_aws_test.go | 65 +++++++++++++++++++++++++++++++++ 4 files changed, 151 insertions(+) create mode 100644 src/segment_aws.go create mode 100644 src/segment_aws_test.go diff --git a/src/properties.go b/src/properties.go index 49cb3a19..c45f4799 100644 --- a/src/properties.go +++ b/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 { diff --git a/src/segment.go b/src/segment.go index 607b4f98..881593d0 100644 --- a/src/segment.go +++ b/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{ diff --git a/src/segment_aws.go b/src/segment_aws.go new file mode 100644 index 00000000..158b24ec --- /dev/null +++ b/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() +} diff --git a/src/segment_aws_test.go b/src/segment_aws_test.go new file mode 100644 index 00000000..db4b7a1d --- /dev/null +++ b/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) + } +}