Merge pull request #4 from tecracer-chef/fix/better-ec2-enumeration

Fix EC2 enumeration

Author: thheinen

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+## Version 0.2.1
+
+- Fix EC2 instance enumeration to use paging
+
 ## Version 0.2.0
 
 - Add more parameter and status checks

README.md

@@ -25,6 +25,7 @@ You can build and install this gem on your local system as well via a Rake task:
 | `execution_timeout`  | Maximum time until timeout                        | 60               |
 | `recheck_invocation` | Interval of rechecking AWS command invocation     | 1.0              |
 | `recheck_execution`  | Interval of rechecking completion of command      | 1.0              |
+| `instance_pagesize`  | Paging size for EC2 instance retrieval            | 100              |
 
 ## Limitations
 

lib/train-awsssm/connection.rb

@@ -7,7 +7,7 @@ module TrainPlugins
   module AWSSSM
     class Connection < Train::Plugins::Transport::BaseConnection
       attr_reader :instance_id, :options
-      attr_writer :ssm, :ec2
+      attr_writer :ssm, :ec2, :instances
 
       def initialize(options)
         super(options)
@@ -85,8 +85,6 @@ def check_options
       def resolve_instance_id(address)
         logger.debug format("[AWS-SSM] Trying to resolve address %s", address)
 
-        instances = ec2.describe_instances.reservations.collect { |r| r.instances.first }
-
         # Resolve, if DNS name and not Amazon default
         if dns_name?(address) && !amazon_dns?(address)
           address = Resolv.getaddress(address)
@@ -112,6 +110,28 @@ def resolve_instance_id(address)
         raise ArgumentError, format("Error looking up Instance ID for %s: %s", address, e.message)
       end
 
+      # List up EC2 instances in the account.
+      #
+      # @param [Boolean] cache Cache results
+      # @return [Array] List of instances
+      # @todo Implement paging
+      def instances(caching: true)
+        return @instances unless @instances.nil? || !caching
+
+        results = []
+
+        ec2_instances = ec2.describe_instances(max_results: options[:instance_pagesize])
+        loop do
+          results.concat ec2_instances.reservations.map(&:instances).flatten
+
+          break unless ec2_instances.next_token
+
+          ec2_instances = ec2.describe_instances(max_results: options[:instance_pagesize], next_token: ec2_instances.next_token)
+        end
+
+        @instances = results
+      end
+
       # Check if this is an IP address
       #
       # @param [String] address Host, IP address or other input

lib/train-awsssm/transport.rb

@@ -11,6 +11,7 @@ class Transport < Train.plugin(1)
       option :execution_timeout,  default: 60.0
       option :recheck_invocation, default: 1.0
       option :recheck_execution,  default: 1.0
+      option :instance_pagesize,  default: 100
 
       def connection(_instance_opts = nil)
         @connection ||= TrainPlugins::AWSSSM::Connection.new(@options)

lib/train-awsssm/version.rb

@@ -1,5 +1,5 @@
 module TrainPlugins
   module AWSSSM
-    VERSION = "0.2.0".freeze
+    VERSION = "0.2.1".freeze
   end
 end