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 'xmlsimple'

class VulnerabilityProcessor
  #Initialise the Vulnerability processor object by creating a new VulnerabilityHelper object
  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)
  # @param scenario_vulns [String] Attributes specified in scenario.xml
  # @return [Hash] Vulnerability values
  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 vulnerability_query.difficulty.length > 0
        puts "Searching for vulnerability matching difficulty: " + vulnerability_query.difficulty
        search_list.delete_if{|x| x.difficulty != vulnerability_query.difficulty}
      end

      if vulnerability_query.cvss_rating.length > 0
        puts "Searching for vulnerability matching cvss rating: " + vulnerability_query.cvss_rating
        remove_by_cvss(vulnerability_query, search_list)
      end

      if vulnerability_query.vector_string.length > 0
        puts "Searching for vulnerability based on vector string: " + vulnerability_query.vector_string
        remove_by_vector(vulnerability_query, search_list)
      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

  # Get array of vulnerability objects
  # @return [Array] Array containing vulnerability objects
  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

  # Remove vulnerability queries from search list
  # @param vulnerability_query [String] Vulnerability query
  # @param search_list [Array] List containing all remaining vulnerabilities
  def remove_by_cvss (vulnerability_query, search_list)
    puts case vulnerability_query.cvss_rating
           when 'none' # 0.0
             search_list.delete_if{|x| x.cvss_score.to_f > 0 }
           when 'low' # 0.1 - 3.9
             search_list.delete_if{|x| x.cvss_score.to_f == 0 or x.cvss_score.to_f > 4 }
           when 'medium' # 4.0 - 6.9
             search_list.delete_if{|x| x.cvss_score.to_f < 4 or x.cvss_score.to_f > 7 }
           when 'high' # 7.0 - 8.9
             search_list.delete_if{|x| x.cvss_score.to_f < 7 and x.cvss_score.to_f <= 9 }
           when 'critical' # 9.0 - 10
             search_list.delete_if{|x| x.cvss_score.to_f < 9 }
         end
  end

  # method which removes vulnerabilities from the search_list based on vector string provided
  # in the vulnerability_query (i.e. a user specified <vulnerability> in scenario.xml)
  # @param query_vulnerability [String] Vector string for desired vulnerabilities
  # @param search_list [Array] List containing all remaining vulnerabilities
  def remove_by_vector(query_vulnerability, search_list)

    query_vector_hash = query_vulnerability.get_vector_hash

    for query_vector_pair in query_vector_hash
      search_list.delete_if{ |vulnerability|
        search_vector_hash = vulnerability.get_vector_hash
        search_vector_pair = search_vector_hash.assoc(query_vector_pair[0])
        if search_vector_pair != nil
          query_vector_pair[1] != search_vector_pair[1]
        else
          true
        end
      }
    end
  end

end