lib/aws_su.rb in aws_su-0.1.6 vs lib/aws_su.rb in aws_su-0.1.7

@@ -64,11 +64,10 @@
     region = AWSConfig.profiles[@profile][:region] unless
         AWSConfig.profiles[@profile][:region].nil?
     @region = options[:region].nil? ? region : options[:region]
     raise('Unable to determine region') if @region.nil?
 
-    export_aws_sudo_file
     assume_role
   end
 
   # Configure the ec2 client
   def ec2_client
@@ -108,23 +107,23 @@
   # Assume a role
   # @param duration A string integer representing the role session duration
   def assume_role(duration = DURATION)
     if session_valid?
       # Recover persisted session and use that to update AWS.config
-      Aws.config.update(
+      config = Aws.config.update(
         credentials: Aws::Credentials.new(
           parse_access_key,
           parse_secret_access_key,
           parse_session_token
         )
       )
+      export_config_to_environment(config)
     else
       # Session has expired so auth again
       assume_role_mfa(duration)
+      export_aws_sudo_file
     end
-    # For the benefit of anything downstream we are running
-    export_aws_sudo_file
   end
 
   # Assume a role using an MFA Token
   def assume_role_mfa(duration, mfa_code = nil)
     mfa_code = prompt_for_mfa_code if mfa_code.nil?
@@ -173,13 +172,13 @@
     end
   end
 
   # Export the AWS values to the ENV
   def export_config_to_environment(config)
-    ENV['AWS_ACCESS_KEY_ID'] = config.credentials.access_key_id
-    ENV['AWS_SECRET_ACCESS_KEY'] = config.credentials.secret_access_key
-    ENV['AWS_SESSION_TOKEN'] = config.credentials.session_token
-    ENV['AWS_SECURITY_TOKEN'] = config.credentials.session_token
+    ENV['AWS_ACCESS_KEY_ID'] = config[:credentials].access_key_id
+    ENV['AWS_SECRET_ACCESS_KEY'] = config[:credentials].secret_access_key
+    ENV['AWS_SESSION_TOKEN'] = config[:credentials].session_token
+    ENV['AWS_SECURITY_TOKEN'] = config[:credentials].session_token
     ENV['AWS_TOKEN_TTL'] = @token_ttl
     ENV['AWS_DEFAULT_REGION'] = @region
   end
 
   # Load the user's AWS Secrets