Merge c000016 into 108b4bc

bin/ssh_scan

@@ -149,16 +149,15 @@ scan") do |file|
     options["unit_test"] = true
   end
 
-  opts.on("-V", "--verbosity [STD_LOGGING_LEVEL]", "File to write JSON output to") do |verb|
+  opts.on("-V", "--verbosity [STD_LOGGING_LEVEL]",
+          "File to write JSON output to") do |verb|
     case verb
     when "INFO"
       options["logger"].level ==  Logger::INFO
     when "WARN"
       options["logger"].level ==  Logger::WARN
     when "DEBUG"
       options["logger"].level ==  Logger::DEBUG
-    when "WARN"
-      options["logger"].level ==  Logger::WARN
     when "ERROR"
       options["logger"].level ==  Logger::ERROR
     when "FATAL"

lib/ssh_scan/api.rb

@@ -18,6 +18,10 @@ class API < Sinatra::Base
       configure do
         set :job_queue, JobQueue.new()
         set :authentication, false
+        config_file = File.join(Dir.pwd, "./config/api/config.yml")
+        opts = YAML.load_file(config_file)
+        opts["config_file"] = config_file
+        set :db, SSHScan::Database.from_hash(opts)
       end
     end
 
@@ -50,6 +54,12 @@ class API < Sinatra::Base
       headers "Pragma" => "no-cache"
     end
 
+    helpers do
+      def cache_valid?(start_time)
+        (Time.now - Time.parse(start_time.to_s)) / (60 * 60 * 24) < 1
+      end
+    end
+
     # Custom 404 handling
     not_found do
       content_type "text/plain"
@@ -103,6 +113,18 @@ class API < Sinatra::Base
         options[:sockets] <<
           "#{params[:target]}:#{params[:port] ? params[:port] : "22"}"
         options[:policy_file] = options[:policy]
+        options[:force] = params[:force] ? params[:force] : false
+
+        unless options[:force] == 'true'
+          available_result = settings.db.fetch_cached_result(params)
+          unless available_result.nil?
+            if cache_valid?(available_result[:start_time])
+              return {
+                uuid: available_result[:uuid]
+              }.to_json
+            end
+          end
+        end
         options[:uuid] = SecureRandom.uuid
         settings.job_queue.add(options)
         {
@@ -112,13 +134,9 @@ class API < Sinatra::Base
 
       get '/scan/results' do
         uuid = params[:uuid]
-
         return {"scan" => "not found"}.to_json if uuid.nil? || uuid.empty?
-
         result = settings.db.find_scan_result(uuid)
-
         return {"scan" => "not found"}.to_json if result.nil?
-
         return result.to_json
       end
 
@@ -158,12 +176,16 @@ class API < Sinatra::Base
       post '/work/results/:worker_id/:uuid' do
         worker_id = params['worker_id']
         uuid = params['uuid']
+        result = JSON.parse(request.body.first).first
+        socket = {}
+        socket[:target] = result['ip']
+        socket[:port] = result['port']
 
         if worker_id.empty? || uuid.empty?
           return {"accepted" => "false"}.to_json
         end
 
-        settings.db.add_scan(worker_id, uuid, JSON.parse(request.body.first).first)
+        settings.db.add_scan(worker_id, uuid, result, socket)
       end
 
       get '/__lbheartbeat__' do

lib/ssh_scan/database.rb

@@ -32,8 +32,8 @@ def self.from_hash(opts)
     # @param [String] uuid
     # @param [Hash] result
     # @return [Nil]
-    def add_scan(worker_id, uuid, result)
-      @database.add_scan(worker_id, uuid, result)
+    def add_scan(worker_id, uuid, result, socket)
+      @database.add_scan(worker_id, uuid, result, socket)
       return nil
     end
 
@@ -52,5 +52,10 @@ def delete_all
     def find_scan_result(uuid)
       @database.find_scan_result(uuid)
     end
+
+    # @return [Hash] result
+    def fetch_cached_result(socket)
+      @database.fetch_cached_result(socket)
+    end
   end
 end

lib/ssh_scan/database/mongo.rb

@@ -26,9 +26,12 @@ def self.from_hash(opts)
       # @param [String] worker_id
       # @param [String] uuid
       # @param [Hash] result
-      def add_scan(worker_id, uuid, result)
-        @collection.insert_one("uuid" => uuid, "scan" => result,
-                          "worker_id" => worker_id, "scanned_on" => Time.now)
+       def add_scan(worker_id, uuid, result, socket)
+        @collection.insert_one("uuid" => uuid,
+                          "target" => socket[:target],
+                          "port" => socket[:port],
+                          "scan" => result,
+                          "worker_id" => worker_id)
       end
 
       def delete_scan(uuid)
@@ -47,6 +50,18 @@ def find_scan_result(uuid)
 
         return nil
       end
+
+      def fetch_cached_result(socket)
+        results = @collection.find(:target => socket[:target], :port => socket[:port])
+        results = results.skip(results.count() - 1)
+        return nil if results.count.zero?
+        result = {}
+        results.each do |result|
+          result[:uuid] = result[:uuid]
+          result[:start_time] = result[:scan][:start_time]
+          return result
+        end
+      end
     end
   end
 end

lib/ssh_scan/database/sqlite.rb

@@ -31,9 +31,10 @@ def self.from_hash(opts)
           db.execute <<-SQL
             create table ssh_scan (
               uuid varchar(100),
+              target varchar(100),
+              port varchar(100),
               result json,
-              worker_id varchar(100),
-              scanned_on datetime
+              worker_id varchar(100)
             );
           SQL
         end
@@ -42,13 +43,14 @@ def self.from_hash(opts)
       end
 
       def size
-        count = @database.execute("select count() from api_schema")
+        count = @database.execute("select count() from ssh_scan")
         return count
       end
 
-      def add_scan(worker_id, uuid, result)
-        @database.execute "insert into ssh_scan values ( ? , ? , ? , ? )",
-                    [uuid, result.to_json, worker_id, Time.now.to_s]
+      def add_scan(worker_id, uuid, result, socket)
+        @database.execute "insert into ssh_scan values ( ? , ? , ? , ? , ? )",
+                    [uuid, socket[:target], socket[:port],
+                     result.to_json, worker_id]
       end
 
       def delete_scan(uuid)
@@ -67,10 +69,23 @@ def find_scan_result(uuid)
           "select * from ssh_scan where uuid like ( ? )",
           uuid
         ) do |row|
-          return JSON.parse(row[1])
+          return JSON.parse(row[3])
         end
         return nil
       end
+
+      def fetch_cached_result(socket)
+        result = {}
+        results = @database.execute(
+          "select uuid, result from ssh_scan
+          where target like ( ? ) and port like ( ? )",
+          [socket[:target], socket[:port]]
+        )
+        return nil if results == []
+        result[:uuid] = results[result.length()-1][0]
+        result[:start_time] = JSON.parse(results[result.length()-1][1])["start_time"]
+        return result
+      end
     end
   end
 end

spec/ssh_scan/api_spec.rb

@@ -64,10 +64,10 @@ def app
     }.to_json)
   end
 
-  it "should say ConnectTimeout for bad IP, and return valid JSON" do
-    bad_ip = "192.168.255.255"
-    port = "999"
-    post "/api/v#{SSHScan::API_VERSION}/scan", {:target => bad_ip, :port => port}
+  it "should return string uuid" do
+    ip = "192.168.1.1"
+    port = "22"
+    post "/api/v#{SSHScan::API_VERSION}/scan", {:target => ip, :port => port}
     expect(last_response.status).to eql(200)
     expect(last_response.body).to be_kind_of(::String)
     expect(last_response["Content-Type"]).to eql("application/json") 

spec/ssh_scan/database/mongodb_spec.rb

@@ -23,10 +23,11 @@
     worker_id = SecureRandom.uuid
     uuid = SecureRandom.uuid
     result = {"ip" => "127.0.0.1", "port" => 1337, "foo" => "bar", "biz" => "baz"}
+    socket = {"target" => "127.0.0.1", "port" => 1337}
 
     temp_file = Tempfile.new('sqlite_database_file')
 
-    @mongodb.add_scan(worker_id, uuid, result)
+    @mongodb.add_scan(worker_id, uuid, result, socket)
 
     # Emulate the retrieval process
     doc = @mongodb.collection.find(:uuid => uuid).first
@@ -35,19 +36,17 @@
     expect(doc["uuid"]).to eql(uuid)
     expect(doc["scan"]).to eql(result)
     expect(doc["worker_id"]).to eql(worker_id)
-
-    #Example: "2017-01-05 14:08:08 -0500"
-    expect(doc["scanned_on"].to_s).to match(/\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2} UTC/)
   end
 
   it "should #delete_scan only the scan we ask it to" do
     worker_id = SecureRandom.uuid
     uuid1 = SecureRandom.uuid
     uuid2 = SecureRandom.uuid
     result = {"ip" => "127.0.0.1", "port" => 1337, "foo" => "bar", "biz" => "baz"}
+    socket = {"target" => "127.0.0.1", "port" => 1337}
 
-    @mongodb.add_scan(worker_id, uuid1, result)
-    @mongodb.add_scan(worker_id, uuid2, result)
+    @mongodb.add_scan(worker_id, uuid1, result, socket)
+    @mongodb.add_scan(worker_id, uuid2, result, socket)
 
     # Verify that we now have two entries in the DB
     first_docs = @mongodb.collection.find(:worker_id => worker_id)
@@ -67,9 +66,10 @@
     uuid1 = SecureRandom.uuid
     uuid2 = SecureRandom.uuid
     result = {"ip" => "127.0.0.1", "port" => 1337, "foo" => "bar", "biz" => "baz"}
+    socket = {"target" => "127.0.0.1", "port" => 1337}
 
-    @mongodb.add_scan(worker_id, uuid1, result)
-    @mongodb.add_scan(worker_id, uuid2, result)
+    @mongodb.add_scan(worker_id, uuid1, result, socket)
+    @mongodb.add_scan(worker_id, uuid2, result, socket)
 
     # Let's delete all of them via the collection
     @mongodb.delete_all
@@ -85,9 +85,10 @@
     uuid2 = SecureRandom.uuid
     result1 = {"ip" => "127.0.0.1", "port" => 1337, "foo" => "bar", "biz" => "baz"}
     result2 = {"ip" => "127.0.0.1", "port" => 1337, "foo" => "bar2", "biz" => "baz2"}
+    socket = {"target" => "127.0.0.1", "port" => 1337}
 
-    @mongodb.add_scan(worker_id, uuid1, result1)
-    @mongodb.add_scan(worker_id, uuid2, result2)
+    @mongodb.add_scan(worker_id, uuid1, result1, socket)
+    @mongodb.add_scan(worker_id, uuid2, result2, socket)
 
     # It should find the first scan
     response1 = @mongodb.find_scan_result(uuid1)
@@ -105,8 +106,9 @@
     uuid1 = SecureRandom.uuid
     bogus_uuid = SecureRandom.uuid
     result1 = {"ip" => "127.0.0.1", "port" => 1337, "foo" => "bar", "biz" => "baz"}
+    socket = {"target" => "127.0.0.1", "port" => 1337}
 
-    @mongodb.add_scan(worker_id, uuid1, result1)
+    @mongodb.add_scan(worker_id, uuid1, result1, socket)
 
     # It should return nil for a non-existant uuid
     response = @mongodb.find_scan_result(bogus_uuid)