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

    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

    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)
    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