updates to the kdi importer

classes/nvd_downloader.rb

@@ -1,12 +1,14 @@
 require 'rest-client'
 require 'zlib'
 require 'json'
+require 'active_support/all'
 #require '../modules/nvd_tools'
 
 module NvdTools
   class NvdDownloader
     attr_accessor :version, :base_url, :base_filename, :years, :filenames_json, :filenames_meta, :client
    
+    CURRENT_TIME = Time.now 
     MIN_YEAR = '2002'
     MAX_YEAR = '2019'
     AVAILABLE_YEARS = (MIN_YEAR..MAX_YEAR).to_a

kenna_kdi_importer/kenna_kdi_importer.rb

@@ -1,7 +1,42 @@
+require 'optparse'
 require './asset_generator'
 require './cve_info'
 require './kdi_jsonify'
 require 'pry'
 
 asset_generator = KennaKdi::AssetGenerator.new('./data/cve')
-Pry.start
+
+ARGV << '-h' if ARGV.empty?
+
+params = {}
+
+OptionParser.new do |opts|
+  opts.banner = "Usage: kenna_kdi_importer.rb [OPTIONS]"
+
+  opts.on('-a NUM', '--assets', Integer, "Number of assets to create") { |a| params[:assets] = a }
+  opts.on('-v NUM', '--max_vulns', Integer, "Max number of vulns that may be created per asset") {|v| params[:max_vulns] = v }
+
+  opts.on_tail("-h", "--help", "Show this message") do
+    puts opts
+    exit
+  end
+end.parse!(into: params)
+
+raise OptionParser::MissingArgument if params[:assets].nil?
+raise OptionParser::MissingArgument if params[:max_vulns].nil?
+
+num_of_assets = params[:assets]
+max_vulns = params[:max_vulns]
+
+puts "Now creating #{num_of_assets} assets. Each asset with a random set of vulnerabilities between 1 and #{max_vulns}"
+puts "\n"
+
+while true do
+  hash = asset_generator.create_assets(num_of_assets, max_vulns)
+  if hash
+    break
+  end
+end
+
+json = hash.to_json
+puts JSON.pretty_generate(JSON.parse(json))