From fef0e78ac3e9971f13675612cce333d4f415d5fa Mon Sep 17 00:00:00 2001 From: Jonathan Claudius Date: Mon, 22 Aug 2016 20:55:03 -0400 Subject: [PATCH 1/2] Improve exception handling logic --- lib/ssh_scan/client.rb | 3 +++ lib/ssh_scan/error.rb | 2 ++ lib/ssh_scan/error/connection_refused.rb | 12 ++++++++++++ lib/ssh_scan/error/disconnected.rb | 12 ++++++++++++ lib/ssh_scan/scan_engine.rb | 2 ++ 5 files changed, 31 insertions(+) create mode 100644 lib/ssh_scan/error/connection_refused.rb create mode 100644 lib/ssh_scan/error/disconnected.rb 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..efde83aa 100644 --- a/lib/ssh_scan/scan_engine.rb +++ b/lib/ssh_scan/scan_engine.rb @@ -26,6 +26,8 @@ def scan_target(target, opts) net_ssh_session.close rescue Net::SSH::ConnectionTimeout => e result[:error] = SSHScan::Error::ConnectTimeout.new(e.message) + rescue Net::SSH::Disconnect => 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)") From 790286d7cb55d4bac1c4de144eace658dda58ce2 Mon Sep 17 00:00:00 2001 From: Jonathan Claudius Date: Mon, 22 Aug 2016 21:29:03 -0400 Subject: [PATCH 2/2] Constrain threads to a finite set of workers --- bin/ssh_scan | 8 +++++++- lib/ssh_scan/scan_engine.rb | 30 ++++++++++++++++++++++++------ 2 files changed, 31 insertions(+), 7 deletions(-) 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/scan_engine.rb b/lib/ssh_scan/scan_engine.rb index efde83aa..324cf063 100644 --- a/lib/ssh_scan/scan_engine.rb +++ b/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