Constrain threads to a finite set of workers

bin/ssh_scan

@@ -13,7 +13,8 @@ options = {
   :port => 22,
   :policy => File.expand_path("../../policies/mozilla_modern.yml", __FILE__),
   :unit_test => false,
-  :timeout => 3
+  :timeout => 2,
+  :threads => 5,
 }
 
 target_parser = SSHScan::TargetParser.new()
@@ -62,6 +63,11 @@ opt_parser = OptionParser.new do |opts|
     options[:policy] = policy
   end
 
+  opts.on("--threads [NUMBER]",
+          "Number of worker threads (Default: 5)") do |threads|
+    options[:threads] = threads.to_i
+  end
+
   opts.on("-u", "--unit-test [FILE]",
           "Throw appropriate exit codes based on compliance status") do
     options[:unit_test] = true

lib/ssh_scan/scan_engine.rb

@@ -25,8 +25,16 @@ def scan_target(target, opts)
         host_key = net_ssh_session.host_keys.first
         net_ssh_session.close
       rescue Net::SSH::ConnectionTimeout => e
-          result[:error] = SSHScan::Error::ConnectTimeout.new(e.message)
+        warn("WARNING: net-ssh timed out attempting to connect to service (fingerprints and auth_methods will not be available)")
+        result['auth_methods'] = []
+        result['fingerprints'] = {}
+        result[:error] = e
+        result[:error] = SSHScan::Error::ConnectTimeout.new(e.message)
       rescue Net::SSH::Disconnect => e
+        warn("WARNING: net-ssh disconnected unexpectedly (fingerprints and auth_methods will not be available)")
+        result['auth_methods'] = []
+        result['fingerprints'] = {}
+        result[:error] = e
         result[:error] = SSHScan::Error::Disconnected.new(e.message)
       rescue Net::SSH::Exception => e
         if e.to_s.match(/could not settle on encryption_client algorithm/)
@@ -74,15 +82,25 @@ def scan_target(target, opts)
 
     def scan(opts)
       targets = opts[:targets]
+      threads = opts[:threads] || 5
 
       results = []
-      threads = []
-      targets.each_with_index do |target, index|
-        threads << Thread.new do
-          results << scan_target(target, opts)
+
+      work_queue = Queue.new
+      targets.each {|x| work_queue.push x }
+      workers = (0...threads).map do |worker_num|
+        Thread.new do
+          begin
+            while target = work_queue.pop(true)
+              results << scan_target(target, opts)
+            end
+          rescue ThreadError => e
+            raise e unless e.to_s.match(/queue empty/)
+          end
         end
       end
-      threads.map(&:join)
+      workers.map(&:join)
+
       return results
     end
   end