this now creates multiple assets and randomizes the amount of vulns per asset

kenna_kdi_importer/asset_generator.rb

@@ -1,35 +1,49 @@
 require 'faker'
+require './vuln_generator'
 
 module KennaKdi
-  class AssetGenerator
-    attr_accessor :cve_data_path, :vuln_generator
+  class AssetGenerator < VulnGenerator
 
-    def initialize(cve_data_path)
-      @cve_data_path = cve_data_path
-      @vuln_generator = KennaKdi::VulnGenerator.new(cve_data_path)
+    def skip_autoclose(value = true)
+      { "skip_autoclose": value }
     end
 
-    def random_asset_hash
-      v_and_vd = random_vuln_and_vuln_def
+    def assets_merge(assets_array)
+      { "assets": assets_array }
+    end
+
+    def vuln_def_merge(vuln_hashes)
+      vuln_def_array = vuln_hashes.map do |vuln_hash|
+        vuln_def_hash(vuln_hash)
+      end
 
       {
-        "skip_autoclose": true,
-        "assets":[
-          {
-            "ip_address": Faker::Internet.ip_v4_address,
-            "vulns":[
-              v_and_vd[:vuln]
-            ]
-          }
-        ],
-        "vuln_defs":[
-          v_and_vd[:vuln_def]
-        ]
+        "vuln_defs": vuln_def_array
       }
     end
 
-    def random_vuln_and_vuln_def
-      vuln_generator.random_vuln_and_vuln_def
+    def create_assets(number_of_assets, max_number_vulns)
+      assets = assets_array(number_of_assets, max_number_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)
+      # 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))
+      end
+    end
+
+    def random_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)
+      {
+        "ip_address": Faker::Internet.ip_v4_address,
+        "vulns": vulns(v_and_vds)
+      }
     end
   end
-end
+end

kenna_kdi_importer/kenna_kdi_importer.rb

@@ -1,4 +1,7 @@
 require './asset_generator'
-require './vuln_generator'
 require './cve_info'
 require './kdi_jsonify'
+require 'pry'
+
+asset_generator = KennaKdi::AssetGenerator.new('./data/cve')
+Pry.start

kenna_kdi_importer/vuln_generator.rb

@@ -1,40 +1,26 @@
 module KennaKdi
   class VulnGenerator
-    attr_accessor :cve_data_path, :cve_data
+    attr_accessor :cve_data_path
 
     def initialize(cve_data_path)
       # path to a directory of json.gz nvd files for CveReport class
       @cve_data_path = cve_data_path
-      @cve_data = random_cve_report 
     end
 
-    def random_vuln_and_vuln_def
-      # spit out a pair of vuln/vuln_def hashes
-      scanner_id = Faker::Code.nric
-      t = Time.new
-      timestamp = t.strftime("%Y-%m-%d %H:%M:%S")
-      
+    def vulns(vulns_and_vuln_defs)
+      vulns_and_vuln_defs.flat_map do |vdata|
+        vdata[:vuln]
+      end
+    end
 
-      id = cve_data.cve_ids.sample
-      cve = cve_data.cve(id)
-      description = cve_data.description(id) 
+    def vuln_defs(vulns_and_vuln_defs)
+      vulns_and_vuln_defs.flat_map do |vdata|
+        vdata[:vuln_def]
+      end
+    end
 
-      {
-        "vuln": {
-          "scanner_identifier": scanner_id,
-          "scanner_type": "KDI Faker Data",
-          "created_at": timestamp,
-          "last_seen_at": timestamp,
-          "status": "open"
-        },
-        "vuln_def": {
-          "scanner_identifier": scanner_id,
-          "scanner_type": "KDI Faker Data",
-          "cve_identifiers": id,
-          "name": "#{scanner_id} - #{id}",
-          "description": description
-        }
-      }
+    def multiple_vulns(num_of_vulns)
+      num_of_vulns.times.map { random_vuln_and_vuln_def } 
     end
 
     private
@@ -43,5 +29,43 @@ module KennaKdi
       cve_files = Dir.glob(File.join(cve_data_path, '**', '*')).select{|file| File.file?(file)}
       CveReport.new(cve_files.sample)
     end
+
+    def vuln_hash
+      scanner_id = Faker::Code.nric
+      t = Time.new
+      timestamp = t.strftime("%Y-%m-%d %H:%M:%S")
+
+      {
+        "scanner_identifier": scanner_id,
+        "scanner_type": "KDI Faker Data",
+        "created_at": timestamp,
+        "last_seen_at": timestamp,
+        "status": "open"
+      }
+    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) 
+
+      {
+        "scanner_identifier": vuln_hash[:scanner_identifier],
+        "scanner_type": vuln_hash[:scanner_type],
+        "cve_identifiers": id,
+        "name": "#{vuln_hash[:scanner_identifier]} - #{id}",
+        "description": description
+      }
+    end
+
+    def random_vuln_and_vuln_def
+      # spit out a pair of vuln/vuln_def hashes
+      vuln = vuln_hash
+      {
+        "vuln": vuln,
+        "vuln_def": vuln_def_hash(vuln)
+      }
+    end
   end
 end