Adding support for CMD, PowerShell and Fish

.gitignore

@@ -14,4 +14,5 @@
 # Dependency directories (remove the comment below to include it)
 # vendor/
 *.json
-build/*
+build/*
+*.DS_Store

README.md

@@ -26,6 +26,16 @@ export AWS_SECRET_ACCESS_KEY="BBBBBBBBBBBBBBBBBB"
 export AWS_SESSION_TOKEN="CCCCCCCCCCCCCCCCCCCCCCCC"
 ```
 
+As an alternative for people using [`fish`](https://fishshell.com/), [`cmd`](https://en.wikipedia.org/wiki/Cmd.exe) or [`PowerShell`](https://en.wikipedia.org/wiki/PowerShell) (defaults to `bash`), it is possible to set which output is preferred. As example for `PowerShell`:
+
+```powershell
+PS> aws sts assume-role --role-arn ${ROLE_ARN} --role-session-name ${SESSION_NAME} --external-id ${EXTERNAL_ID} | atg -powershell
+
+$Env:AWS_ACCESS_KEY_ID="AAAAAAAAAAAAAAA"
+$Env:AWS_SECRET_ACCESS_KEY="BBBBBBBBBBBBBBBBBB"
+$Env:AWS_SESSION_TOKEN="CCCCCCCCCCCCCCCCCCCCCCCC"
+```
+
 Read from `stdin` and import 
 
 ```sh
@@ -45,18 +55,18 @@ $> aws sts get-caller-identity
 }
 ```
 
-Read from `role.json` file:
+Read from `role.json` file for `fish`:
 
 ```sh
-$> aws sts get-caller-identity
+~> aws sts get-caller-identity
 {
     "UserId": "AIDA11111111111111111",
     "Account": "111111111111",
     "Arn": "arn:aws:iam::111111111111:user/myuser"
 }
 
-$> eval $(atg -json role.json)
-$> aws sts get-caller-identity
+~> eval $(atg -json role.json -fish)
+~> aws sts get-caller-identity
 {
     "UserId": "AROA22222222222222222:${SESSION_NAME}",
     "Account": "222222222222",

main.go

@@ -1,28 +1,95 @@
+/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements.  See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership.  The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License.  You may obtain a copy of the License at
+   http://www.apache.org/licenses/LICENSE-2.0
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied.  See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/
+
 package main
 
 import (
 	"flag"
+	"fmt"
+	"os"
 
 	"github.com/gocaio/atg/modules/parser"
 )
 
-var jsonfile string
+var (
+	jsonfile   string
+	powershell bool
+	cmd        bool
+	fish       bool
+	bash       bool
+)
 
 func init() {
 	flag.StringVar(&jsonfile, "json", "", "JSON file to use")
+	flag.BoolVar(&bash, "bash", false, "Export bash/zsh environmental variables [default]")
+	flag.BoolVar(&cmd, "cmd", false, "Export Windows CMD environmental variables")
+	flag.BoolVar(&fish, "fish", false, "Export fish environmental variables")
+	flag.BoolVar(&powershell, "powershell", false, "Export Powershell environmental variables")
+
 	flag.Parse()
 }
 
 func main() {
+	boolElements := [4]bool{bash, cmd, powershell, fish}
+	if Count(boolElements) > 1 {
+		fmt.Println("Choose one [bash, cmd, fish, powershell]")
+		fmt.Println("Use atg -h")
+		//flag.PrintDefaults()
+		os.Exit(0)
+	}
 
 	if jsonfile == "" {
 		onStdin := parser.CheckStdin()
 		if !onStdin {
 			flag.PrintDefaults()
 		} else {
-			parser.ParseStdin("AA")
+			access, secret, token := parser.ParseStdin()
+			if cmd == true {
+				parser.KeyPrinter(access, secret, token, "cmd")
+			} else if fish == true {
+				parser.KeyPrinter(access, secret, token, "fish")
+			} else if powershell == true {
+				parser.KeyPrinter(access, secret, token, "powershell")
+			} else {
+				parser.KeyPrinter(access, secret, token, "bash")
+			}
 		}
 	} else {
-		parser.ParseJSON(jsonfile)
+		access, secret, token := parser.ParseJSON(jsonfile)
+		if cmd == true {
+			parser.KeyPrinter(access, secret, token, "cmd")
+		} else if fish == true {
+			parser.KeyPrinter(access, secret, token, "fish")
+		} else if powershell == true {
+			parser.KeyPrinter(access, secret, token, "powershell")
+		} else {
+			parser.KeyPrinter(access, secret, token, "bash")
+		}
+	}
+}
+
+// Count will count how many
+// bool flags are there in total
+func Count(boolElements [4]bool) int {
+	var True = 0
+	for i := range boolElements {
+		if boolElements[i] == true {
+			True = True + 1
+		}
 	}
+	return True
 }

modules/parser/parser.go

@@ -1,3 +1,20 @@
+/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements.  See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership.  The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License.  You may obtain a copy of the License at
+   http://www.apache.org/licenses/LICENSE-2.0
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied.  See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/
+
 package parser
 
 import (
@@ -13,14 +30,19 @@ import (
 	"github.com/gocaio/atg/modules/structs"
 )
 
+// CheckStdin will check if there is
+// something in the stdin buffer
 func CheckStdin() bool {
 	if termutil.Isatty(os.Stdin.Fd()) {
 		return false
 	}
 	return true
 }
 
-func ParseStdin(stdin string) {
+// ParseStdin will parse the output from
+// the aws sts assume-role command from
+// stdin piping ATG
+func ParseStdin() (string, string, string) {
 	var buf bytes.Buffer
 	reader := bufio.NewReader(os.Stdin)
 
@@ -43,11 +65,12 @@ func ParseStdin(stdin string) {
 	err := json.Unmarshal(buf.Bytes(), &data)
 	CheckErr(err)
 
-	KeyPrinter(data.Credentials.AccessKeyID, data.Credentials.SecretAccessKey, data.Credentials.SessionToken)
-
+	return data.Credentials.AccessKeyID, data.Credentials.SecretAccessKey, data.Credentials.SessionToken
 }
 
-func ParseJSON(jsonfile string) {
+// ParseJSON will spect the output json from the
+// aws sts assume-role command
+func ParseJSON(jsonfile string) (string, string, string) {
 
 	jsonFile, err := os.Open(jsonfile)
 	CheckErr(err)
@@ -56,14 +79,31 @@ func ParseJSON(jsonfile string) {
 	byteValue, _ := ioutil.ReadAll(jsonFile)
 	var data structs.AssumeRole
 	json.Unmarshal(byteValue, &data)
-	//pp.Print(data)
-	KeyPrinter(data.Credentials.AccessKeyID, data.Credentials.SecretAccessKey, data.Credentials.SessionToken)
+	return data.Credentials.AccessKeyID, data.Credentials.SecretAccessKey, data.Credentials.SessionToken
 }
 
-func KeyPrinter(access, secret, token string) {
-	fmt.Printf("export AWS_ACCESS_KEY_ID=\"%v\"\n", access)
-	fmt.Printf("export AWS_SECRET_ACCESS_KEY=\"%v\"\n", secret)
-	fmt.Printf("export AWS_SESSION_TOKEN=\"%v\"\n", token)
+// KeyPrinter will print the environment
+// variables depending of the chosen shell
+func KeyPrinter(access, secret, token, shell string) {
+	switch shell {
+	case "cmd":
+		fmt.Printf("set AWS_ACCESS_KEY_ID=\"%v\"\n", access)
+		fmt.Printf("set AWS_SECRET_ACCESS_KEY=\"%v\"\n", secret)
+		fmt.Printf("set AWS_SESSION_TOKEN=\"%v\"\n", token)
+	case "fish":
+		fmt.Printf("set -x AWS_ACCESS_KEY_ID=\"%v\"\n", access)
+		fmt.Printf("set -x AWS_SECRET_ACCESS_KEY=\"%v\"\n", secret)
+		fmt.Printf("set -x AWS_SESSION_TOKEN=\"%v\"\n", token)
+	case "powershell":
+		fmt.Printf("$Env:AWS_ACCESS_KEY_ID=\"%v\"\n", access)
+		fmt.Printf("$Env:AWS_SECRET_ACCESS_KEY=\"%v\"\n", secret)
+		fmt.Printf("$Env:AWS_SESSION_TOKEN=\"%v\"\n", token)
+	default:
+		fmt.Printf("export AWS_ACCESS_KEY_ID=\"%v\"\n", access)
+		fmt.Printf("export AWS_SECRET_ACCESS_KEY=\"%v\"\n", secret)
+		fmt.Printf("export AWS_SESSION_TOKEN=\"%v\"\n", token)
+	}
+
 }
 
 // CheckErr will handle errors

modules/structs/structs.go

@@ -1,3 +1,20 @@
+/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements.  See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership.  The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License.  You may obtain a copy of the License at
+   http://www.apache.org/licenses/LICENSE-2.0
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied.  See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/
+
 package structs
 
 import "time"