From 9638c0ccd63848d19d72a73ff915aa13b829e8c9 Mon Sep 17 00:00:00 2001
From: Jonathan Claudius <jclaudius@mozilla.com>
Date: Tue, 15 Jan 2019 11:32:35 -0500
Subject: [PATCH] Fix up fingerprint storage capabilities

---
 bin/ssh_scan                |  8 ++++----
 lib/ssh_scan/scan_engine.rb | 18 ++++++------------
 lib/ssh_scan/subprocess.rb  | 26 ++++++++++++++++++++++++++
 3 files changed, 36 insertions(+), 16 deletions(-)
 create mode 100644 lib/ssh_scan/subprocess.rb

diff --git a/bin/ssh_scan b/bin/ssh_scan
index ffca12a1..cc9a2162 100755
--- a/bin/ssh_scan
+++ b/bin/ssh_scan
@@ -18,7 +18,7 @@ options = {
   "threads" => 5,
   "verbosity" => nil,
   "logger" => Logger.new(STDERR),
-  "fingerprint_database" => File.join(File.dirname(__FILE__),"../data/fingerprints.yml")
+  "fingerprint_database" => ENV['HOME']+'/.ssh_scan_fingerprints.yml'
 }
 
 # Reorder arguments before parsing
@@ -247,9 +247,9 @@ end
 #end
 
 # Limit scope of fingerprints DB to (per scan)
-if options["fingerprint_database"] && File.exists?(options["fingerprint_database"])
-  File.unlink(options["fingerprint_database"])
-end
+# if options["fingerprint_database"] && File.exists?(options["fingerprint_database"])
+#   File.unlink(options["fingerprint_database"])
+# end
 
 options["policy_file"] = SSHScan::Policy.from_file(options["policy"])
 
diff --git a/lib/ssh_scan/scan_engine.rb b/lib/ssh_scan/scan_engine.rb
index 999e1ba6..631be1b2 100644
--- a/lib/ssh_scan/scan_engine.rb
+++ b/lib/ssh_scan/scan_engine.rb
@@ -1,7 +1,8 @@
 require 'socket'
 require 'ssh_scan/client'
 require 'ssh_scan/crypto'
-#require 'ssh_scan/fingerprint_database'
+require 'ssh_scan/fingerprint_database'
+require 'ssh_scan/subprocess'
 require 'net/ssh'
 require 'logger'
 require 'open3'
@@ -122,17 +123,10 @@ def scan_target(socket, opts)
 
       output = ""
 
-      begin
-        Timeout::timeout(timeout) {
-          stdin, stdout, stderr, wait_thr = Open3.popen3('ssh-keyscan', '-t', 'rsa,dsa', '-p', port.to_s, target)
-          output = stdout.gets(nil) if port.nil?
-          stdout.close
-          output = stderr.gets(nil) if !port.nil?
-          stderr.close
-          exit_code = wait_thr.value
-        }
-      rescue Timeout::Error
-        #nop
+      cmd = ['ssh-keyscan', '-t', 'rsa,dsa', '-p', port.to_s, target].join(" ")
+
+      Utils::Subprocess.new(cmd) do |stdout, stderr, thread|
+        output += stdout
       end
 
       host_keys = output.split
diff --git a/lib/ssh_scan/subprocess.rb b/lib/ssh_scan/subprocess.rb
new file mode 100644
index 00000000..660d8d2a
--- /dev/null
+++ b/lib/ssh_scan/subprocess.rb
@@ -0,0 +1,26 @@
+require 'open3'
+
+module Utils
+  class Subprocess
+    def initialize(cmd, &block)
+      # see: http://stackoverflow.com/a/1162850/83386
+      Open3.popen3(cmd) do |stdin, stdout, stderr, thread|
+        # read each stream from a new thread
+        { :out => stdout, :err => stderr }.each do |key, stream|
+          Thread.new do
+            until (line = stream.gets).nil? do
+              # yield the block depending on the stream
+              if key == :out
+                yield line, nil, thread if block_given?
+              else
+                yield nil, line, thread if block_given?
+              end
+            end
+          end
+        end
+
+        thread.join # don't exit until the external process is done
+      end
+    end
+  end
+end