nexpose_servicenow 0.4.24 → 0.5.1

data/lib/nexpose_servicenow/chunker.rb

@@ -1,3 +1,4 @@
+#TODO: Check if using site_id is OK. If the object is changed, the code on SN will likely need changed.
 
 module NexposeServiceNow
   class Chunker
@@ -12,28 +13,28 @@ module NexposeServiceNow
 
     def setup_logging
       @log = NexposeServiceNow::NxLogger.instance
-      @log.log_message("Chunker File Limit: #{@size_limit}MB");
-      @log.log_message("Chunk Row Limit: #{@row_limit}")
+      @log.log_message("Chunk File Limit:\t#{@size_limit}MB")
+      @log.log_message("Chunk Row Limit:\t#{@row_limit}")
     end
 
-    #Grab the header from the first file
+    # Grab the header from the first file
     def get_header
-      file = File.open(@report_details.first[:report_name], "r")
+      file = File.open(@report_details.first[:report_name], 'r')
       header = file.readline
       file.close
 
       header
     end
 
-    def preprocess(nexpose_ids=nil)
+    def preprocess
       all_chunks = []
       @report_details.each do |report|
-        @log.log_message("Breaking file #{report[:report_name]} down into chunks.")
+        @log.log_message("Dividing file #{report[:report_name]} into chunks.")
         chunks = process_file(report[:report_name], report[:id])
         all_chunks.concat chunks
       end 
 
-      @log.log_message("Files broken down into #{all_chunks.count} chunks")
+      @log.log_message("Files divided into #{all_chunks.count} chunks")
 
       puts all_chunks.to_json
     end
@@ -41,12 +42,12 @@ module NexposeServiceNow
     def process_file(file_path, site_id=nil)
       relative_size_limit = @size_limit - @header.bytesize
       chunk = { site_id: site_id,
-                start: @header.bytesize, 
+                start: @header.bytesize,
                 length: 0, 
                 row_count: 0 }
       
       chunks = []
-      csv_file = CSV.open(file_path, "r", headers: true)
+      csv_file = CSV.open(file_path, 'r', headers: true)
       while(true)
         position = csv_file.pos
         line = csv_file.shift
@@ -62,12 +63,13 @@ module NexposeServiceNow
         else
           chunks << chunk
 
+          # TODO: Make generic?
           #Initialise chunk with this row information
           chunk = { site_id: site_id,
                     start: position, 
                     length: row_length, 
                     row_count: 1 }
-        end  
+        end
       end
       csv_file.close
 
@@ -83,7 +85,7 @@ module NexposeServiceNow
     end
 
     def get_file(site_id=nil)
-      #-1 indicates a single query report
+      # -1 indicates a single query report
       return @report_details.first[:report_name] if site_id.to_i <= 0
 
       report = @report_details.find { |r| r[:id].to_s == site_id.to_s }
@@ -91,12 +93,15 @@ module NexposeServiceNow
     end
 
     def read_chunk(start, length, site_id=nil)
-      @log.log_message("Returning chunk. Start: #{start}, Length: #{length}, File: #{@file_path}")
+      file_path = get_file(site_id)
+      msg = "Returning chunk. Start: #{start}, " \
+            "Length: #{length}, File: #{file_path}"
+      @log.log_message(msg)
 
       #If the header isn't in the chunk, prepend it
-      header = start == 0 ? "" : @header
+      header = start == 0 ? '' : @header
 
-      file = File.open(get_file(site_id), "rb")
+      file = File.open(file_path, 'rb')
       file.seek(start)
       puts header + file.read(length)
       file.close

data/lib/nexpose_servicenow/historical_data.rb

@@ -3,37 +3,81 @@ require_relative './nx_logger'
 
 module NexposeServiceNow
   class HistoricalData
-    REPORT_FILE = "Nexpose-ServiceNow-latest_scans.csv"
-    STORED_FILE = "last_scan_data.csv"
-    TIMESTAMP_FILE = "last_vuln_run.csv"
-    NEW_TIMESTAMP_FILE = "new_vuln_timestamp.csv"
 
-    def initialize(options)
-      local_dir = File.expand_path(options[:output_dir])
-      @ids = options[:nexpose_ids]
+    REPORT_FILE = 'Nexpose-ServiceNow-latest_scans.csv'
+    STORED_FILE = 'last_scan_data.csv'
+    TIMESTAMP_FILE = 'last_vuln_run.csv'
+    NEW_TIMESTAMP_FILE = 'new_vuln_timestamp.csv'
+
+    DAG_TIMESTAMP_FILE = 'last_scan_data_dag.csv'
+    NEW_DAG_TIMESTAMP_FILE = 'new_dag_timestamp.csv'
+
+    SITE_IDENTIFIER = 'site_id'
+    SITE_DELTA_VALUE = 'last_scan_id'
+    SITE_BASE_VALUE = 0
+
+    DAG_IDENTIFIER = 'asset_group_id'
+    DAG_DELTA_VALUE = 'last_import'
+    DAG_BASE_VALUE = '1985-01-01 00:00:00'
+
+    def initialize(output_dir, nexpose_ids, id_type, start_time)
+      local_dir = File.expand_path(output_dir)
+      @ids = nexpose_ids
+      @id_type = id_type
+
+      if @id_type == :site
+        current_file = STORED_FILE
+        new_file = REPORT_FILE
+      elsif @id_type == :asset_group
+        current_file = DAG_TIMESTAMP_FILE
+        new_file = NEW_DAG_TIMESTAMP_FILE
+      end
+
+      @start_time = start_time
 
-      @local_file = File.join(local_dir, STORED_FILE)
-      @remote_file = File.join(local_dir, REPORT_FILE)
+      @local_file = File.join(local_dir, current_file)
+      @remote_file = File.join(local_dir, new_file)
 
       # File containing the timestamp used in vulnerability queries
       @timestamp_file = File.join(local_dir, TIMESTAMP_FILE)
       @prev_timestamp_file = File.join(local_dir, NEW_TIMESTAMP_FILE)
 
       @log = NexposeServiceNow::NxLogger.instance
-      @log.log_message "Retrieving environment variables."
     end
 
-    #Filters the saved report down to the sites being queried
-    #This can then be used as a basis to update last_scan_data
+    # TODO: Remove site references here? Will there be remote CSV here if we're using DAGs? No scan IDs.
+    # Filters the saved report down to the sites being queried
+    # This can then be used as a basis to update last_scan_data
     def filter_report
-      #Create a full last_scan_data if it doesn't already exist
-      create_last_scan_data unless File.exist? @local_file
+      # Create a full last_scan_data if it doesn't already exist
+      create_base_delta_file unless File.exist? @local_file
 
-      @log.log_message 'Filtering report down sites which will be queried'
+      @log.log_message 'Filtering report down to sites which will be queried'
 
       remote_csv = load_scan_id_report
-      site_ids = @ids.map(&:to_s)
-      filtered_csv = remote_csv.delete_if { |r| !site_ids.include?(r['site_id']) }
+      nexpose_ids = @ids.map(&:to_s)
+      identifier = if @id_type == :site
+                     SITE_IDENTIFIER
+                   elsif @id_type ==:asset_group
+                     DAG_IDENTIFIER
+                   end
+
+      if @id_type == :asset_group
+        header = [DAG_IDENTIFIER, DAG_DELTA_VALUE]
+        rows = []
+
+        @ids.each do |i|
+          rows << CSV::Row.new(header, [i, @start_time])
+        end
+
+        remote_csv = CSV::Table.new(rows)
+      end
+
+      # TODO: Why is this done? Aren't these already filtered?
+      filtered_csv = remote_csv.delete_if do |r|
+        !nexpose_ids.include?(r[identifier])
+      end
+
       File.open(@remote_file, 'w') do |f|
         f.write(remote_csv.to_csv)
       end
@@ -41,45 +85,58 @@ module NexposeServiceNow
       puts filtered_csv
     end
 
-    #Reads the downloaded report containing LATEST scan IDs
+    # Reads the downloaded report containing LATEST scan IDs
     def load_scan_id_report
-      @log.log_message "Loading scan data report"
+      @log.log_message 'Loading scan data report'
       unless File.exists? @remote_file
-        @log.log_message "No existing report file found."
+        @log.log_message 'No existing report file found.'
         return nil
       end
       CSV.read(@remote_file, headers: true)
     end
 
-    #Loads the last scan data file as CSV.
-    #It may be necessary to create one first.
+    # Loads the last scan data file as CSV.
+    # It may be necessary to create one first.
     def load_last_scan_data
-      @log.log_message "Loading last scan data"
+      @log.log_message 'Loading last scan data.'
 
-      create_last_scan_data unless File.exist? @local_file
+      create_base_delta_file unless File.exist? @local_file
       CSV.read(@local_file, headers: true)
     end
 
-    def last_scan_ids(sites)
+    def stored_delta_values(nexpose_ids)
       return [] if !File.exist? @local_file
 
+      if @id_type == :site
+        identifier = SITE_IDENTIFIER
+        delta_column = SITE_DELTA_VALUE
+        base_value = SITE_BASE_VALUE
+      elsif @id_type == :asset_group
+        identifier = DAG_IDENTIFIER
+        delta_column = DAG_DELTA_VALUE
+        base_value = DAG_BASE_VALUE
+      end
+
       csv = load_last_scan_data
-      last_scan_ids = {}
-      sites.each do |id|
-        row = csv.find { |r| r['site_id'] == id.to_s } || { 'last_scan_id' => '0' }
-        last_scan_ids[id.to_s] = row['last_scan_id']
+      delta_values = {}
+      nexpose_ids.each do |id|
+        row = csv.find { |r| r[identifier] == id.to_s }
+        row ||= { delta_column => base_value }
+        delta_values[id.to_s] = row[delta_column]
       end
 
-      last_scan_ids
+      delta_values
     end
 
-    #Compares stored scan IDs versus remote scan IDs.
-    #This determines which scans are included as filters.
-    def sites_to_scan(imported_sites_only=false)
+    # Compares stored scan IDs versus remote scan IDs.
+    # This determines which scans are included as filters.
+    def collections_to_import(previously_imported_only=false)
       return @ids unless File.exist? @remote_file
+      @log.log_message "Filtering for #{@id_type}s with new scans"
+      self.send("#{@id_type}s_to_import", previously_imported_only)
+    end
 
-      @log.log_message 'Filtering for sites with new scans'
-
+    def sites_to_import(previously_imported_only=false)
       remote_csv = CSV.read(@remote_file, headers: true)
       local_csv = load_last_scan_data
 
@@ -93,7 +150,7 @@ module NexposeServiceNow
         local_scan_id = local_scan_id['last_scan_id'] || 0
 
         # Check if only allowing sites which were previously imported
-        next if local_scan_id.to_s == "0" && imported_sites_only
+        next if local_scan_id.to_s == '0' && previously_imported_only
 
         filtered_sites << id if local_scan_id.to_i < remote_scan_id.to_i
       end
@@ -101,22 +158,50 @@ module NexposeServiceNow
       @ids = filtered_sites
     end
 
-    #Creates a base last scan data file from a downloaded report
-    def create_last_scan_data
-      @log.log_message 'Creating base last scan data file'
+    def asset_groups_to_import(previously_imported_only=false)
+      filtered_asset_groups = []
+      local_csv = load_last_scan_data
+
+      @ids.each do |id|
+        local_id = local_csv.find { |r| r[DAG_IDENTIFIER] == id.to_s } || {}
+        local_id = local_id[DAG_DELTA_VALUE] || DAG_BASE_VALUE
+
+        next if local_id == DAG_BASE_VALUE && previously_imported_only
+
+        filtered_asset_groups << id
+      end
 
+      @ids = filtered_asset_groups
+    end
+
+    # Creates a base last scan data file from a downloaded report
+    def create_base_delta_file
+      @log.log_message 'Creating base delta file'
+      self.send("create_#{@id_type}_base_file")
+    end
+
+    def create_site_base_file
       csv = load_scan_id_report
       csv.delete('finished')
-      csv.each { |l| l['last_scan_id'] = 0 }
+      csv.each { |l| l['last_scan_id'] = SITE_BASE_VALUE }
 
       save_last_scan_data(csv)
     end
 
-    #Updates only the rows that were affected by this scan
-    def update_last_scan_data
-      @log.log_message "Updating last scan data"
+    def create_asset_group_base_file
+      CSV.open(@local_file, 'w') do |csv|
+        csv << %w(asset_group_id last_import)
+        @ids.each do |n|
+          csv << [n, DAG_BASE_VALUE]
+        end
+      end
+    end
+
+    # Updates only the rows that were affected by this scan
+    def update_delta_file
+      @log.log_message 'Updating last scan data'
 
-      if !(File.exist? @local_file) && !(File.exist? @remote_file) 
+      if !(File.exist? @local_file) && !(File.exist? @remote_file)
         @log.log_message 'Last scan data does not exist yet.'
         return
       end
@@ -124,23 +209,51 @@ module NexposeServiceNow
       updated_csv = load_last_scan_data
       remote_csv = load_scan_id_report
 
+      method = "update_#{@id_type}_delta_file"
+      updated_csv = self.send(method, updated_csv, remote_csv)
+
+      save_last_scan_data(updated_csv)
+    end
+
+    def update_site_delta_file(updated_csv, remote_csv)
       #merge changes in from remote_csv
       remote_csv.each do |row|
         updated_row = updated_csv.find { |r| r['site_id'] == row['site_id'] }
         if updated_row.nil?
           row.delete 'finished'
-          updated_csv << row    
+          updated_csv << row
         else
-          updated_row['last_scan_id'] = row['last_scan_id'] 
+          updated_row['last_scan_id'] = row['last_scan_id']
         end
       end
 
-      save_last_scan_data(updated_csv) 
+      updated_csv
+    end
+
+    def update_asset_group_delta_file(updated_csv, remote_csv)
+      #merge changes in from remote_csv
+      remote_csv.each do |row|
+        updated_row = updated_csv.find do |r|
+          r['asset_group_id'] == row['asset_group_id']
+        end
 
-      #puts updated_csv
+        if updated_row.nil?
+          updated_csv << row
+        else
+          updated_row['last_import'] = row['last_import']
+        end
+      end
+
+      # Insert any NEW IDs with baseline time
+      @ids.each do |i|
+        row = updated_csv.find { |r| r[DAG_IDENTIFIER] == i }
+        updated_csv << [i, DAG_BASE_VALUE] if row.nil?
+      end
+
+      updated_csv
     end
 
-    #Overwrite the last scan data file with new csv
+    # Overwrite the last scan data file with new csv
     def save_last_scan_data(csv)
       @log.log_message 'Saving last scan data'
       File.open(@local_file, 'w') do |f|
@@ -148,9 +261,7 @@ module NexposeServiceNow
       end
     end
 
-
-    #insert sites?
-    def save_vuln_timestamp(sites=[])
+    def save_vuln_timestamp(nexpose_ids=[])
       start_time = Time.new
 
       #Read timestamp from new timestamp file (substitute base time)
@@ -167,32 +278,32 @@ module NexposeServiceNow
 
       last_run ||= Time.new(1985)
       last_sites ||= []
-      last_run = last_run.strftime("%Y-%m-%d") if last_run.class.to_s == 'Time'
+      last_run = last_run.strftime('%Y-%m-%d') if last_run.class.to_s == 'Time'
       create_last_vuln_data(last_run, last_sites)
 
       file = File.expand_path(@prev_timestamp_file)
       CSV.open(file, 'w') do |csv|
         csv << ['Last Scan Time', 'Sites']
-        csv << [start_time.strftime("%Y-%m-%d"), sites.join(',')]
+        csv << [start_time.strftime('%Y-%m-%d'), nexpose_ids.join(',')]
       end
     end
 
-    def create_last_vuln_data(time=nil, sites=[])
+    def create_last_vuln_data(time=nil, nexpose_ids=[])
       @log.log_message 'Creating last vulnerability scan time file.'
 
       time ||= Time.new(1985)
-      time = time.strftime("%Y-%m-%d") if time.class.to_s == 'Time'
-      sites = sites.join(',') if sites.class.to_s == 'Array'
+      time = time.strftime('%Y-%m-%d') if time.class.to_s == 'Time'
+      nexpose_ids = nexpose_ids.join(',') if nexpose_ids.class.to_s == 'Array'
 
       file = File.expand_path(@timestamp_file)
 
       CSV.open(file, 'w') do |csv|
         csv << ['Last Scan Time', 'Sites']
-        csv << [time, sites]
+        csv << [time, nexpose_ids]
       end
     end
 
-    #Current IDs are inserted into the updated CSV file.
+    # Current IDs are inserted into the updated CSV file.
     def last_vuln_run
       @log.log_message 'Retrieving the last vulnerability timestamp'
 
@@ -206,8 +317,11 @@ module NexposeServiceNow
       last_run
     end
 
-    #Experimental
-    #These should probably return strings that can be mlog'd
+    #########################################################
+    #                    Experimental                       #
+    #########################################################
+
+    # These should probably return strings that can be mlog'd
     def log_and_print(message)
       puts message
       @log.log_message message unless @log.nil?
@@ -227,15 +341,15 @@ module NexposeServiceNow
 
       csv = load_last_scan_data
       row = csv.find { |r| r['site_id'] == nexpose_id }
-      
+
       if row.nil?
         csv << [nexpose_id, scan_id]
       else
         row['last_scan_id'] = scan_id
       end
-        
+
       save_last_scan_data csv
-      
+
       log_and_print 'Last scan data updated.'
     end
 
@@ -253,7 +367,7 @@ module NexposeServiceNow
 
       new_name = "#{filename}.#{Time.new.strftime('%Y-%m-%d.%H:%M:%S')}"
       begin
-        #Delete existing file with same name
+        # Delete existing file with same name
         File.delete new_name if File.exist? new_name
         File.rename(filename, new_name)
       rescue Exception => e