supports credential source AWS_PROFILE export

trek10inc · Mar 1, 2020 · e61b248 · e61b248
1 parent cae07c7
commit e61b248
Show file tree

Hide file tree

Showing 9 changed files with 87 additions and 28 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1,5 +1,6 @@
 # misc
 *.local
+*.local.*
 node_modules
 
 # Byte-compiled / optimized / DLL files

diff --git a/awsume/__data__.py b/awsume/__data__.py
@@ -1,4 +1,4 @@
-version = '4.2.6'
+version = '4.2.7a1'
 
 name = 'awsume'
 author = 'Trek10, Inc'

diff --git a/awsume/awsumepy/app.py b/awsume/awsumepy/app.py
@@ -258,6 +258,7 @@ def run(self, system_arguments: list):
                     str(credentials.get('SessionToken')),
                     str(credentials.get('Region')),
                     str(args.target_profile_name),
+                    str(credentials.get('AwsProfile')),
                 ])
         except exceptions.EarlyExit:
             logger.debug('', exc_info=True)

diff --git a/awsume/awsumepy/default_plugins.py b/awsume/awsumepy/default_plugins.py
@@ -437,6 +437,13 @@ def get_credentials_no_mfa(config: dict, arguments: argparse.Namespace, profiles
     return return_session
 
 
+def get_credentials_from_credential_source(config: dict, arguments: argparse.Namespace, profiles: dict, target_profile: dict):
+    region = profile_lib.get_region(profiles, arguments, config)
+    return_session = {'AwsProfile': arguments.target_profile_name}
+    return_session['Region'] = region
+    return return_session
+
+
 def get_session_token_credentials(config: dict, arguments: argparse.Namespace, profiles: dict, target_profile: dict):
     region = profile_lib.get_region(profiles, arguments, config)
     mfa_serial = profile_lib.get_mfa_serial(profiles, arguments.target_profile_name)
@@ -479,15 +486,15 @@ def get_credentials(config: dict, arguments: argparse.Namespace, profiles: dict)
         else:
             if mfa_serial:
                 user_session = get_session_token_credentials(config, arguments, profiles, target_profile)
+            elif 'credential_source' in target_profile:
+                user_session = get_credentials_from_credential_source(config, arguments, profiles, target_profile)
             else:
                 user_session = get_credentials_no_mfa(config, arguments, profiles, target_profile)
 
     if config.get('is_interactive'):
-        if user_session:
-            if user_session.get('Expiration'):
-                safe_print('Session token will expire at {}'.format(profile_lib.parse_time(user_session['Expiration'])), colorama.Fore.GREEN)
-        if role_session:
-            if role_session.get('Expiration'):
-                safe_print('Role credentials will expire {}'.format(profile_lib.parse_time(role_session['Expiration'])), colorama.Fore.GREEN)
+        if user_session and user_session.get('Expiration'):
+            safe_print('Session token will expire at {}'.format(profile_lib.parse_time(user_session['Expiration'])), colorama.Fore.GREEN)
+        if role_session and role_session.get('Expiration'):
+            safe_print('Role credentials will expire {}'.format(profile_lib.parse_time(role_session['Expiration'])), colorama.Fore.GREEN)
 
     return role_session or user_session
diff --git a/awsume/awsumepy/lib/exceptions.py b/awsume/awsumepy/lib/exceptions.py
@@ -24,7 +24,7 @@ def __init__(self, profile_name, message=''):
         self.profile_name = profile_name
         self.message = message
     def __str__(self):
-        return 'Invalid profile {} {}'.format(self.profile_name, self.message)
+        return 'Invalid profile [{}]: {}'.format(self.profile_name, self.message)
 
 
 class ValidationException(AwsumeException):

diff --git a/awsume/awsumepy/lib/profile.py b/awsume/awsumepy/lib/profile.py
@@ -77,12 +77,16 @@ def validate_profile(config: dict, arguments: argparse.Namespace, profiles: dict
     # validate user profile
     if user_profile:
         missing_keys = []
-        if 'aws_access_key_id' not in user_profile:
-            missing_keys.append('aws_access_key_id')
-        if 'aws_secret_access_key' not in user_profile:
-            missing_keys.append('aws_secret_access_key')
+        if 'credential_source' in profile:
+            if profile.get('credential_source') not in VALID_CREDENTIAL_SOURCES:
+                raise exceptions.InvalidProfileError(user_profile_name, message='unsupported awsume credential_source profile option: {}'.format(profile.get('credential_source')))
+        else:
+            if 'aws_access_key_id' not in user_profile:
+                missing_keys.append('aws_access_key_id')
+            if 'aws_secret_access_key' not in user_profile:
+                missing_keys.append('aws_secret_access_key')
         if missing_keys:
-            raise exceptions.InvalidProfileError(user_profile_name, message='Missing keys {}'.format(', '.join(missing_keys)))
+            raise exceptions.InvalidProfileError(user_profile_name, message='Missing keys {}, or credential_source'.format(', '.join(missing_keys)))
 
     # validate arguments with profile
     if 'role_arn' not in profile and arguments.auto_refresh:

diff --git a/shell_scripts/awsume b/shell_scripts/awsume
@@ -99,8 +99,12 @@ elif [ "$AWSUME_FLAG" = "Awsume" ]; then
   unset AWS_PROFILE
   unset AWS_DEFAULT_PROFILE
   unset AWSUME_PROFILE
-  export AWS_ACCESS_KEY_ID=${AWSUME_1}
-  export AWS_SECRET_ACCESS_KEY=${AWSUME_2}
+  if [ ! "${AWSUME_1}" = "None" ]; then
+    export AWS_ACCESS_KEY_ID=${AWSUME_1}
+  fi
+  if [ ! "${AWSUME_2}" = "None" ]; then
+    export AWS_SECRET_ACCESS_KEY=${AWSUME_2}
+  fi
   if [ ! "${AWSUME_3}" = "None" ]; then
     export AWS_SESSION_TOKEN=${AWSUME_3}
     export AWS_SECURITY_TOKEN=${AWSUME_3}
@@ -112,11 +116,19 @@ elif [ "$AWSUME_FLAG" = "Awsume" ]; then
   if [ ! "${AWSUME_5}" = "None" ]; then
     export AWSUME_PROFILE=${AWSUME_5}
   fi
+  if [ ! "${AWSUME_6}" = "None" ]; then
+    export AWS_PROFILE=${AWSUME_6}
+    export AWS_DEFAULT_PROFILE=${AWSUME_6}
+  fi
   for AWSUME_var in "$@"
   do
     if [[ "$AWSUME_var" == "-s"* ]]; then
-      echo export AWS_ACCESS_KEY_ID=${AWSUME_1}
-      echo export AWS_SECRET_ACCESS_KEY=${AWSUME_2}
+      if [ ! "${AWSUME_1}" = "None" ]; then
+        echo export AWS_ACCESS_KEY_ID=${AWSUME_1}
+      fi
+      if [ ! "${AWSUME_2}" = "None" ]; then
+        echo export AWS_SECRET_ACCESS_KEY=${AWSUME_2}
+      fi
       if [ ! "${AWSUME_3}" = "None" ]; then
         echo export AWS_SESSION_TOKEN=${AWSUME_3}
         echo export AWS_SECURITY_TOKEN=${AWSUME_3}
@@ -128,6 +140,10 @@ elif [ "$AWSUME_FLAG" = "Awsume" ]; then
       if [ ! "${AWSUME_5}" = "None" ]; then
         echo export AWSUME_PROFILE=${AWSUME_5}
       fi
+      if [ ! "${AWSUME_6}" = "None" ]; then
+        echo export AWS_PROFILE=${AWSUME_6}
+        echo export AWS_DEFAULT_PROFILE=${AWSUME_6}
+      fi
     fi
   done
 fi
diff --git a/shell_scripts/awsume.bat b/shell_scripts/awsume.bat
@@ -9,7 +9,7 @@ FOR %%A IN (%*) DO (
     IF "%%A"=="-s" (set "SHOW=y")
 )
 
-for /f "tokens=1,2,3,4,5,6 delims= " %%a in ("%AWSUME_TEXT%") do (
+for /f "tokens=1,2,3,4,5,6,7 delims= " %%a in ("%AWSUME_TEXT%") do (
     if "%%a" == "Auto" (
         set AWS_SECRET_ACCESS_KEY=
         set AWS_SESSION_TOKEN=
@@ -93,25 +93,33 @@ for /f "tokens=1,2,3,4,5,6 delims= " %%a in ("%AWSUME_TEXT%") do (
         set AWS_DEFAULT_PROFILE=
         set AWSUME_PROFILE=
 
-        set AWS_ACCESS_KEY_ID=%%b
-        set AWS_SECRET_ACCESS_KEY=%%c
+        if "%%b" NEQ "None" (
+            set AWS_ACCESS_KEY_ID=%%b)
+
+        if "%%c" NEQ "None" (
+            set AWS_SECRET_ACCESS_KEY=%%c)
 
         if "%%d" NEQ "None" (
             set AWS_SESSION_TOKEN=%%d
             set AWS_SECURITY_TOKEN=%%d)
 
-
         if "%%e" NEQ "None" (
             set AWS_REGION=%%e
             set AWS_DEFAULT_REGION=%%e)
 
         if "%%f" NEQ "None" (
             set AWSUME_PROFILE=%%f)
 
+        if "%%g" NEQ "None" (
+            set AWS_PROFILE=%%g
+            set AWS_DEFAULT_PROFILE=%%g)
+
         IF defined SHOW (
             for /f "tokens=1,2,3,4,5 delims= " %%a in ("%AWSUME_TEXT%") do (
-                echo set AWS_ACCESS_KEY_ID=%%b
-                echo set AWS_SECRET_ACCESS_KEY=%%c
+                if "%%b" NEQ "None" (
+                    echo set AWS_ACCESS_KEY_ID=%%b)
+                if "%%c" NEQ "None" (
+                    echo set AWS_SECRET_ACCESS_KEY=%%c)
 
                 if "%%d" NEQ "None" (
                     echo set AWS_SESSION_TOKEN=%%d
@@ -123,6 +131,10 @@ for /f "tokens=1,2,3,4,5,6 delims= " %%a in ("%AWSUME_TEXT%") do (
 
                 if "%%f" NEQ "None" (
                     echo set AWSUME_PROFILE=%%f)
+
+                if "%%g" NEQ "None" (
+                    echo set AWS_PROFILE=%%g
+                    echo set AWS_DEFAULT_PROFILE=%%g)
             )
         )
     )

diff --git a/shell_scripts/awsume.ps1 b/shell_scripts/awsume.ps1
@@ -2,7 +2,7 @@
 
 #AWSUME_FLAG - what awsumepy told the shell to do
 #AWSUME_n - the data from awsumepy
-$AWSUME_FLAG, $AWSUME_1, $AWSUME_2, $AWSUME_3, $AWSUME_4, $AWSUME_5 = `
+$AWSUME_FLAG, $AWSUME_1, $AWSUME_2, $AWSUME_3, $AWSUME_4, $AWSUME_5, $AWSUME_6 = `
 $(awsumepy $args) -split '\s+'
 
 #if incorrect flag/help
@@ -108,8 +108,12 @@ elseif ( $AWSUME_FLAG -eq "Awsume") {
     $env:AWS_DEFAULT_PROFILE = ""
     $env:AWSUME_PROFILE = ""
 
-    $env:AWS_ACCESS_KEY_ID = $AWSUME_1
-    $env:AWS_SECRET_ACCESS_KEY = $AWSUME_2
+    if ( $AWSUME_1 -ne "None" ) {
+        $env:AWS_ACCESS_KEY_ID = $AWSUME_1
+    }
+    if ( $AWSUME_2 -ne "None" ) {
+        $env:AWS_SECRET_ACCESS_KEY = $AWSUME_2
+    }
 
     if ( $AWSUME_3 -ne "None" ) {
         $env:AWS_SESSION_TOKEN = $AWSUME_3
@@ -125,10 +129,19 @@ elseif ( $AWSUME_FLAG -eq "Awsume") {
         $env:AWSUME_PROFILE = $AWSUME_5
     }
 
+    if ( $AWSUME_6 -ne "None" ) {
+        $env:AWS_PROFILE = $AWSUME_6
+        $env:AWS_DEFAULT_PROFILE = $AWSUME_6
+    }
+
     #if enabled, show the exact commands to use in order to assume the role we just assumed
     if ($args -like "-s") {
-        Write-Host "`$env:AWS_ACCESS_KEY_ID =" $env:AWS_ACCESS_KEY_ID
-        Write-Host "`$env:AWS_SECRET_ACCESS_KEY =" $env:AWS_SECRET_ACCESS_KEY
+        if ( $AWSUME_1 -ne "None" ) {
+            Write-Host "`$env:AWS_ACCESS_KEY_ID =" $env:AWS_ACCESS_KEY_ID
+        }
+        if ( $AWSUME_2 -ne "None" ) {
+            Write-Host "`$env:AWS_SECRET_ACCESS_KEY =" $env:AWS_SECRET_ACCESS_KEY
+        }
 
         if ( $AWSUME_3 -ne "None" ) {
             Write-Host "`$env:AWS_SESSION_TOKEN =" $env:AWS_SESSION_TOKEN
@@ -143,5 +156,10 @@ elseif ( $AWSUME_FLAG -eq "Awsume") {
         if ( $AWSUME_5 -ne "None" ) {
             Write-Host "`$env:AWSUME_PROFILE =" $env:AWSUME_PROFILE
         }
+
+        if ( $AWSUME_6 -ne "None" ) {
+            Write-Host "`$env:AWS_PROFILE =" $AWSUME_6
+            Write-Host "`$env:AWS_DEFAULT_PROFILE =" $AWSUME_6
+        }
     }
 }