diff --git a/bin/ssh_scan b/bin/ssh_scan index b58460a6..6215b8a5 100755 --- a/bin/ssh_scan +++ b/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 diff --git a/lib/ssh_scan/client.rb b/lib/ssh_scan/client.rb index f53fc27c..3d8bca1c 100644 --- a/lib/ssh_scan/client.rb +++ b/lib/ssh_scan/client.rb @@ -28,6 +28,9 @@ def connect() rescue Errno::ETIMEDOUT => e @error = SSHScan::Error::ConnectTimeout.new(e.message) @sock = nil + rescue Errno::ECONNREFUSED => e + @error = SSHScan::Error::ConnectionRefused.new(e.message) + @sock = nil else @raw_server_banner = @sock.gets.chomp @server_banner = SSHScan::Banner.read(@raw_server_banner) diff --git a/lib/ssh_scan/error.rb b/lib/ssh_scan/error.rb index b59b01a3..da65d218 100644 --- a/lib/ssh_scan/error.rb +++ b/lib/ssh_scan/error.rb @@ -1,2 +1,4 @@ require 'ssh_scan/error/connect_timeout' require 'ssh_scan/error/closed_connection' +require 'ssh_scan/error/connection_refused' +require 'ssh_scan/error/disconnected' diff --git a/lib/ssh_scan/error/connection_refused.rb b/lib/ssh_scan/error/connection_refused.rb new file mode 100644 index 00000000..2c54448a --- /dev/null +++ b/lib/ssh_scan/error/connection_refused.rb @@ -0,0 +1,12 @@ +module SSHScan + module Error + class ConnectionRefused < Exception + def initialize(message) + @message = message + end + def to_s + "#{self.class.to_s.split('::')[-1]}: #{@message}" + end + end + end +end diff --git a/lib/ssh_scan/error/disconnected.rb b/lib/ssh_scan/error/disconnected.rb new file mode 100644 index 00000000..b5269eb6 --- /dev/null +++ b/lib/ssh_scan/error/disconnected.rb @@ -0,0 +1,12 @@ +module SSHScan + module Error + class Disconnected < Exception + def initialize(message) + @message = message + end + def to_s + "#{self.class.to_s.split('::')[-1]}: #{@message}" + end + end + end +end diff --git a/lib/ssh_scan/scan_engine.rb b/lib/ssh_scan/scan_engine.rb index e5c1c21b..324cf063 100644 --- a/lib/ssh_scan/scan_engine.rb +++ b/lib/ssh_scan/scan_engine.rb @@ -25,7 +25,17 @@ 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/) warn("WARNING: net-ssh could not find a mutually acceptable encryption algorithm (fingerprints and auth_methods will not be available)") @@ -72,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