Feature: add role assumption support

README.md

@@ -56,6 +56,10 @@ set :ec2_region, %w{} # REQUIRED
 set :ec2_contact_point, nil
 
 set :ec2_filter_by_status_ok?, nil
+
+set :ec2_assume_role, true # enable role assumption
+set :ec2_role_assumption, 'role assumption arn'
+set :ec2_role_session_name, 'role assumption session name'
 ```
 
 #### Order of inheritance

lib/cap-ec2/ec2-handler.rb

@@ -13,14 +13,38 @@ def initialize
       end
     end
 
-    def ec2_connect(region=nil)
+    def ec2_regular_connect(region)
+      Aws::EC2::Client.new(
+          access_key_id: fetch(:ec2_access_key_id),
+          secret_access_key: fetch(:ec2_secret_access_key),
+          region: region
+      )
+    end
+
+    def ec2_role_assumption_connect(region)
+      role_credentials = Aws::AssumeRoleCredentials.new(
+          client: Aws::STS::Client.new(
+              access_key_id: fetch(:ec2_access_key_id),
+              secret_access_key: fetch(:ec2_secret_access_key),
+              region: region
+          ),
+          role_arn: fetch(:ec2_role_assumption),
+          role_session_name: fetch(:ec2_role_session_name),
+          )
       Aws::EC2::Client.new(
-        access_key_id: fetch(:ec2_access_key_id),
-        secret_access_key: fetch(:ec2_secret_access_key),
-        region: region
+          credentials: role_credentials,
+          region: region
       )
     end
 
+    def ec2_connect(region=nil)
+      if fetch(:ec2_assume_role)
+        ec2_role_assumption_connect(region)
+      else
+        ec2_regular_connect(region)
+      end
+    end
+
     def status_table
       CapEC2::StatusTable.new(
         defined_roles.map {|r| get_servers_for_role(r)}.flatten.uniq {|i| i.instance_id}

lib/cap-ec2/tasks/ec2.rake

@@ -32,5 +32,9 @@ namespace :load do
     set :ec2_secret_access_key, nil
     set :ec2_region, %w{}
 
+    set :ec2_assume_role, false
+    set :ec2_role_assumption, nil
+    set :ec2_role_session_name, nil
+
   end
 end

lib/cap-ec2/utils.rb

@@ -68,6 +68,10 @@ def load_config
           set :ec2_region, config['regions'] if config['regions']
 
           set :ec2_filter_by_status_ok?, config['filter_by_status_ok?'] if config['filter_by_status_ok?']
+
+          set :ec2_assume_role, !!config['ec2_assume_role'] if config['ec2_assume_role'].to_s == 'true'
+          set :ec2_role_assumption, config['ec2_role_assumption'] if config['ec2_role_assumption']
+          set :ec2_role_session_name, config['ec2_role_session_name'] if config['ec2_role_session_name']
         end
       end
     end