diff --git a/bin/ssh_scan b/bin/ssh_scan index 00eb9fc6..cbead01a 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.db") + "fingerprint_database" => File.join(File.dirname(__FILE__),"../data/fingerprints.yml") } # Reorder arguments before parsing @@ -246,6 +246,11 @@ ssh_scan' to get the latest" 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 + options["policy_file"] = SSHScan::Policy.from_file(options["policy"]) # Perform scan and get results diff --git a/lib/ssh_scan/fingerprint_database.rb b/lib/ssh_scan/fingerprint_database.rb index f1241a97..89d94506 100644 --- a/lib/ssh_scan/fingerprint_database.rb +++ b/lib/ssh_scan/fingerprint_database.rb @@ -1,39 +1,38 @@ -require 'sqlite3' +require 'yaml/store' module SSHScan class FingerprintDatabase def initialize(database_name) - if File.exists?(database_name) - @db = ::SQLite3::Database.open(database_name) - else - @db = ::SQLite3::Database.new(database_name) - self.create_schema - end - end - - def create_schema - @db.execute <<-SQL - create table fingerprints ( - fingerprint varchar(100), - ip varchar(100) - ); - SQL + @store = YAML::Store.new(database_name) end def clear_fingerprints(ip) - @db.execute "delete from fingerprints where ip like ( ? )", [ip] + @store.transaction do + @store[ip] = [] + end end def add_fingerprint(fingerprint, ip) - @db.execute "insert into fingerprints values ( ?, ? )", [fingerprint, ip] + @store.transaction do + @store[ip] = [] if @store[ip].nil? + @store[ip] << fingerprint + end end def find_fingerprints(fingerprint) - ips = [] - @db.execute( "select * from fingerprints where fingerprint like ( ? )", [fingerprint] ) do |row| - ips << row[1] + ip_matches = [] + + @store.transaction(true) do + @store.roots.each do |ip| + @store[ip].each do |other_fingerprint| + if fingerprint == other_fingerprint + ip_matches << ip + end + end + end end - return ips + + return ip_matches.uniq end end end diff --git a/spec/ssh_scan/fingerprint_database_spec.rb b/spec/ssh_scan/fingerprint_database_spec.rb index 7a0d3eb1..abc3a9b0 100644 --- a/spec/ssh_scan/fingerprint_database_spec.rb +++ b/spec/ssh_scan/fingerprint_database_spec.rb @@ -9,11 +9,20 @@ #start with a known good state File.unlink(file_name) if File.exists?(file_name) - expect(File.exist?(file_name)).to eql(false) + database = SSHScan::FingerprintDatabase.new(file_name) expect(database).to be_kind_of(SSHScan::FingerprintDatabase) + + # The file isn't created until something is written + expect(File.exist?(file_name)).to eql(false) + + # Write something, to trigger file creation + database.add_fingerprint("hello_world", "192.168.1.1") + + # Verify the file exists now expect(File.exist?(file_name)).to eql(true) + File.unlink(file_name) #clean up after ourselves end @@ -24,20 +33,14 @@ File.unlink(file_name) if File.exists?(file_name) # Create a pre-existing DB - db = ::SQLite3::Database.new(file_name) - db.execute <<-SQL - create table fingerprints ( - fingerprint varchar(100), - ip varchar(100) - ); - SQL - db.execute "insert into fingerprints values ( ?, ? )", ["fake_fingerprint", "127.0.0.1"] + database = SSHScan::FingerprintDatabase.new(file_name) + database.add_fingerprint("hello_world", "192.168.1.1") expect(File.exist?(file_name)).to eql(true) - database = SSHScan::FingerprintDatabase.new(file_name) - expect(database).to be_kind_of(SSHScan::FingerprintDatabase) + database2 = SSHScan::FingerprintDatabase.new(file_name) + expect(database2).to be_kind_of(SSHScan::FingerprintDatabase) - expect(database.find_fingerprints("fake_fingerprint")).to eql(["127.0.0.1"]) + expect(database2.find_fingerprints("hello_world")).to eql(["192.168.1.1"]) File.unlink(file_name) #clean up after ourselves end end diff --git a/ssh_scan.gemspec b/ssh_scan.gemspec index cfc11eee..48ebb8fe 100644 --- a/ssh_scan.gemspec +++ b/ssh_scan.gemspec @@ -32,7 +32,6 @@ Gem::Specification.new do |s| s.add_dependency('bindata', '~> 2.0') s.add_dependency('netaddr') s.add_dependency('net-ssh') - s.add_dependency('sqlite3') s.add_dependency('sshkey') s.add_development_dependency('pry') s.add_development_dependency('rspec', '~> 3.0')