SecGen/lib/helpers/vulnerability_processor.rb

require_relative '../../lib/constants' #CW - Could this only be included in secgen.rb and have the constants initialized at runtime?
require_relative '../../lib/objects/vulnerability'
require_relative 'vulnerability_helper'
require 'nokogiri'
require 'xmlsimple'

class VulnerabilityProcessor
    def initialize()
      @vulnerability_helper = VulnerabilityHelper.new
    end
    # returns a hash of compatible vulnerabilities based on what is provided in scenario.xml (scenario_vulns)
    # based on the attributes optionally specified in scenario.xml (scenario_vulns)
    def process(scenario_vulns)

      return_vulns = {}

      all_vulnerabilities = get_vulnerabilities_array

      scenario_vulns.each do |vulnerability_query|
        # select based on selected type, access, cve...
        search_list = all_vulnerabilities.clone
        # shuffle order of available vulnerabilities
        search_list.shuffle!
        # remove all the vulns that don't match the current selection (type, etc)
        if vulnerability_query.type.length > 0
          puts "Searching for vulnerability matching type: " + vulnerability_query.type
          search_list.delete_if{|x| x.type != vulnerability_query.type}
        end
        if vulnerability_query.access.length > 0
          puts "Searching for vulnerability matching access: " + vulnerability_query.access
          search_list.delete_if{|x| x.access != vulnerability_query.access}
        end
        if vulnerability_query.cve.length > 0
          puts "Searching for vulnerability matching CVE: " + vulnerability_query.cve
          search_list.delete_if{|x| x.cve != vulnerability_query.cve}
        end

        if search_list.length == 0
          puts VULN_NOT_FOUND
          puts "(note: you can only have one of each type of vulnerability per system)"
          exit
        else
          # use from the top of the top of the randomised list
          return_vulns[vulnerability_query.id] = search_list[0]
          if search_list[0].type.length > 0
            puts "Selected vulnerability : " + search_list[0].type
          end

          # enforce only one of any vulnerability type (remove from available)
          search_list.delete_if{|x| x.type == vulnerability_query.type}
        end
      end

      return return_vulns.values


    end

    def get_vulnerabilities_array
      vulnerabilities = []
      Dir.glob("#{ROOT_DIR}/modules/vulnerabilities/**/**/secgen_metadata.xml").each do |file|
        vulnerability_hash = XmlSimple.xml_in(file, {})
        vulnerability = @vulnerability_helper.getVulnerabilityObject(vulnerability_hash)
        vulnerabilities.push(vulnerability)
      end

      return vulnerabilities
    end
end