made this thing super fast! woohoO!

kenna_kdi_importer/asset_generator.rb

@@ -22,22 +22,22 @@ module KennaKdi
       }
     end
 
-    def create_assets(number_of_assets, max_number_vulns)
-      assets = assets_array(number_of_assets, max_number_vulns)
+    def create_assets(number_of_assets, number_of_vulns)
+      assets = assets_array(number_of_assets, number_of_vulns)
       vuln_hashes = assets.map {|asset| asset[:vulns]}.flatten
 
       header_and_assets = skip_autoclose.merge(assets_merge(assets))
       header_and_assets.merge(vuln_def_merge(vuln_hashes))
     end
 
-    def assets_array(number_of_assets, max_number_vulns)
+    def assets_array(number_of_assets, number_of_vulns)
       # this should be the primary logic that will generate the "assets": [asset1, asset2] data for the json
       number_of_assets.times.map do
-        random_asset_hash(rand(1..max_number_vulns))
+        multiple_vulns_asset_hash(number_of_vulns)
       end
     end
 
-    def random_asset_hash(number_of_vulns)
+    def multiple_vulns_asset_hash(number_of_vulns)
       # generate number of vulns specified vuln and vulndef pairs to work with
       v_and_vds = multiple_vulns(number_of_vulns)
       {

kenna_kdi_importer/generate_cve_ids.rb

@@ -0,0 +1,28 @@
+require './cve_report'
+
+module KennaKdi
+  class GenerateCveIdsJson
+    attr_accessor :cve_data_path, :cve_files
+
+    def initialize(cve_data_path)
+      @cve_data_path = cve_data_path
+      @cve_files = Dir.glob(File.join(cve_data_path, '**', '*')).select{|file| File.file?(file)}
+    end
+
+    def all_reports
+      cve_files.map do |cve_file|
+        CveReport.new(cve_file)
+      end
+    end
+
+    def perform
+      cve_reports = all_reports
+
+      cve_ids = cve_reports.map do |cve_report|
+        cve_report.cve_ids
+      end.flatten
+
+      cve_ids.to_json
+    end
+  end
+end

kenna_kdi_importer/kdi_format.rb

@@ -1,9 +1,5 @@
-require 'json'
-
 module KennaKdi
-	class KdiJsonify
-		def initialize
-		end
+	class KdiFormat
 	
 		def default_hash
 			# from https://help.kennasecurity.com/hc/en-us/articles/360026413111-Kenna-Data-Importer-JSON-Connector-

kenna_kdi_importer/kenna_kdi_importer.rb

@@ -1,10 +1,7 @@
 require 'optparse'
 require './asset_generator'
-require './cve_info'
-require './kdi_jsonify'
-require 'pry'
-
-asset_generator = KennaKdi::AssetGenerator.new('./data/cve')
+require './cve_report'
+require './kdi_format'
 
 ARGV << '-h' if ARGV.empty?
 
@@ -14,7 +11,7 @@ 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('-v NUM', '--vulns', Integer, "Number of vulns that may be created per asset") {|v| params[:vulns] = v }
 
   opts.on_tail("-h", "--help", "Show this message") do
     puts opts
@@ -23,16 +20,15 @@ OptionParser.new do |opts|
 end.parse!(into: params)
 
 raise OptionParser::MissingArgument if params[:assets].nil?
-raise OptionParser::MissingArgument if params[:max_vulns].nil?
+raise OptionParser::MissingArgument if params[:vulns].nil?
 
 num_of_assets = params[:assets]
-max_vulns = params[:max_vulns]
+num_of_vulns = params[:vulns]
 
-puts "Now creating #{num_of_assets} assets. Each asset with a random set of vulnerabilities between 1 and #{max_vulns}"
-puts "\n"
+asset_generator = KennaKdi::AssetGenerator.new('./data/cve_ids.json')
 
 while true do
-  hash = asset_generator.create_assets(num_of_assets, max_vulns)
+  hash = asset_generator.create_assets(num_of_assets, num_of_vulns)
   if hash
     break
   end

kenna_kdi_importer/vuln_generator.rb

@@ -1,10 +1,10 @@
 module KennaKdi
   class VulnGenerator
-    attr_accessor :cve_data_path
+    attr_accessor :path_to_cve_json, :cve_ids
 
-    def initialize(cve_data_path)
-      # path to a directory of json.gz nvd files for CveReport class
-      @cve_data_path = cve_data_path
+    def initialize(path_to_cve_json)
+      @path_to_cve_json = path_to_cve_json
+      @cve_ids = JSON.parse(File.read(path_to_cve_json)) 
     end
 
     def vulns(vulns_and_vuln_defs)
@@ -23,13 +23,24 @@ module KennaKdi
       num_of_vulns.times.map { random_vuln_and_vuln_def } 
     end
 
+    def sample_cve_ids(num_of_cve)
+      num_of_cve.times.map { cve_ids.sample }
+    end
+
     private
 
     def random_cve_report
-      cve_files = Dir.glob(File.join(cve_data_path, '**', '*')).select{|file| File.file?(file)}
       CveReport.new(cve_files.sample)
     end
 
+    def all_cve_reports
+      cve_files.map { |file| CveReport.new(file) }
+    end
+
+    def cve_report(file_path)
+      CveReport.new(file_path)
+    end
+
     def vuln_hash
       scanner_id = Faker::Code.nric
       t = Time.new
@@ -45,17 +56,12 @@ module KennaKdi
     end
 
     def vuln_def_hash(vuln_hash)
-      cve_data = random_cve_report
-      id = cve_data.cve_ids.sample
-      cve = cve_data.cve(id)
-      description = cve_data.description(id) 
-
+      id = cve_ids.sample
       {
         "scanner_identifier": vuln_hash[:scanner_identifier],
         "scanner_type": vuln_hash[:scanner_type],
         "cve_identifiers": id,
-        "name": "#{vuln_hash[:scanner_identifier]} - #{id}",
-        "description": description
+        "name": "#{vuln_hash[:scanner_identifier]} - #{id}"
       }
     end