Merge pull request #125 from common-fate/Add-support-for--exec-flag

Add support for --exec flag in linux and mac

.vscode/launch.json

@@ -73,6 +73,18 @@
         "FORCE_NO_ALIAS": "true",
         "AWS_PROFILE": "demo-sandbox1"
       }
+    },
+    {
+      "name": "Run command exec",
+      "type": "go",
+      "request": "launch",
+      "mode": "auto",
+      "program": "cmd/assume/main.go",
+      "args": ["cf-testing", "--exec", "aws s3 ls"],
+      "env": {
+        "FORCE_NO_ALIAS": "true",
+        "AWS_PROFILE": "demo-sandbox1"
+      }
     }
   ]
 }

pkg/assume/assume.go

@@ -2,6 +2,9 @@ package assume
 
 import (
 	"fmt"
+	"os/exec"
+	"runtime"
+	"strings"
 	"sync"
 	"time"
 
@@ -26,6 +29,10 @@ func AssumeCommand(c *cli.Context) error {
 	if err != nil {
 		return err
 	}
+
+	if assumeFlags.String("exec") != "" && runtime.GOOS == "windows" {
+		return fmt.Errorf("--exec flag is not currently supported on Windows. Let us know if you'd like support for this: https://github.com/common-fate/granted/issues/new")
+	}
 	var wg sync.WaitGroup
 	activeRoleProfile := assumeFlags.String("granted-active-aws-role-profile")
 	activeRoleFlag := assumeFlags.Bool("active-role")
@@ -137,7 +144,8 @@ func AssumeCommand(c *cli.Context) error {
 					return err
 				}
 			}
-			fmt.Printf("GrantedOutput %s", url)
+			// return the url via stdout through the cli wrapper script
+			fmt.Print(MakeGrantedOutput(url))
 		} else {
 			browsers.PromoteUseFlags(labels)
 			fmt.Fprintf(color.Error, "\nOpening a console for %s in your browser...\n", profile.Name)
@@ -157,13 +165,15 @@ func AssumeCommand(c *cli.Context) error {
 		} else {
 			green.Fprintf(color.Error, "\n[%s](%s) session credentials ready\n", profile.Name, region)
 		}
+		if assumeFlags.String("exec") != "" {
+			return RunExecCommandWithCreds(assumeFlags.String("exec"), creds, region)
+		}
 		// DO NOT REMOVE, this interacts with the shell script that wraps the assume command, the shell script is what configures your shell environment vars
 		// to export more environment variables, add then in the assume and assume.fish scripts then append them to this output preparation function
 		// the shell script treats "None" as an emprty string and will not set a value for that positional output
 		output := PrepareStringsForShellScript([]string{creds.AccessKeyID, creds.SecretAccessKey, creds.SessionToken, profile.Name, region, sessionExpiration})
 		fmt.Printf("GrantedAssume %s %s %s %s %s %s", output...)
 	}
-
 	return nil
 }
 
@@ -181,3 +191,37 @@ func PrepareStringsForShellScript(in []string) []interface{} {
 	}
 	return out
 }
+
+// RunExecCommandWithCreds takes in a command, which may be a program and arguments sperated by spaces
+// it splits these then runs the command with teh credentials as the environment.
+// The output of this is returned via the assume script to stdout so it may be processed further by piping
+func RunExecCommandWithCreds(cmd string, creds aws.Credentials, region string) error {
+	fmt.Print(MakeGrantedOutput(""))
+	args := strings.Split(cmd, " ")
+	c := exec.Command(args[0], args[1:]...)
+	c.Stdout = os.Stdout
+	c.Stderr = color.Error
+	c.Env = append(c.Env, EnvKeys(creds, region)...)
+	return c.Run()
+}
+
+// EnvKeys is used to set the env for the "exec" flag
+func EnvKeys(creds aws.Credentials, region string) []string {
+	return []string{"AWS_ACCESS_KEY_ID=" + creds.AccessKeyID,
+		"AWS_SECRET_ACCESS_KEY=" + creds.SecretAccessKey,
+		"AWS_SESSION_TOKEN=" + creds.SessionToken,
+		"AWS_REGION=" + region}
+}
+
+// MakeGrantedOutput formats a string to match the requirements of granted output in the shell script
+// Currently in windows, the grantedoutput is handled differently, as linux and mac support the exec cli flag whereas windows does not yet have support
+// this method may be changed in future if we implement support for "--exec" in windows
+func MakeGrantedOutput(s string) string {
+	out := "GrantedOutput"
+	if runtime.GOOS != "windows" {
+		out += "\n"
+	} else {
+		out += " "
+	}
+	return out + s
+}

pkg/assume/entrypoint.go

@@ -30,7 +30,9 @@ func GlobalFlags() []cli.Flag {
 		&cli.BoolFlag{Name: "verbose", Usage: "Log debug messages"},
 		&cli.StringFlag{Name: "update-checker-api-url", Value: build.UpdateCheckerApiUrl, EnvVars: []string{"UPDATE_CHECKER_API_URL"}, Hidden: true},
 		&cli.StringFlag{Name: "granted-active-aws-role-profile", EnvVars: []string{"AWS_PROFILE"}, Hidden: true},
-		&cli.BoolFlag{Name: "auto-configure-shell", Usage: "Configure shell alias without prompts"}}
+		&cli.BoolFlag{Name: "auto-configure-shell", Usage: "Configure shell alias without prompts"},
+		&cli.StringFlag{Name: "exec", Usage: "assume a profile then execute this command"},
+	}
 }
 
 func GetCliApp() *cli.App {

pkg/updates/updates.go

@@ -19,7 +19,7 @@ import (
 // the last day checked is stored in the local config cache
 // this function will fail silently
 func Check(c *cli.Context) (string, bool) {
-	if os.Getenv("GRANTED_DISABLE_UPDATE_CHECK") == "true" {
+	if os.Getenv("GRANTED_DISABLE_UPDATE_CHECK") == "true" || build.Version == "dev" {
 		return "", false
 	}
 	updateCheckerApiUrl := c.String("update-checker-api-url")

scripts/assume

@@ -85,8 +85,12 @@ if [ "$GRANTED_FLAG" = "GrantedAssume" ]; then
     fi
   done
 fi
+
+# The GrantedOutput flag should be followed by a newline, then the output. 
+# This way, the shell script can omit the first line containing the flag and return the unaltered output to the stdout
+# This is great as it works well with the -exec flag
 if [ "$GRANTED_FLAG" = "GrantedOutput" ];then
-  echo ${GRANTED_1}
+  echo ${GRANTED_OUTPUT} | sed -n '1!p'
 fi
 
 if [ "$GRANTED_RETURN_STATUS" = "true" ]; then

scripts/assume.fish

@@ -73,7 +73,11 @@ else if test "$GRANTED_FLAG" = "GrantedAssume"
     end
   end
 else if test "$GRANTED_FLAG" = "GrantedOutput"
-  echo $GRANTED_1
+  for line in $GRANTED_OUTPUT
+      if test "$line" != "GrantedOutput"
+        echo $line
+      end
+  end
 end
 
 exit $GRANTED_STATUS