batch_secgen: Automatically back up the projects/ directory of successful projects, added prompt on excluding ranges currently in table from randomised ip generation.

.gitignore

@@ -5,3 +5,4 @@ unusedcode
 mount
 log
 .directory
+batch/

lib/batch/batch_secgen.rb

@@ -1,3 +1,4 @@
+require 'fileutils'
 require 'getoptlong'
 require 'open3'
 require 'pg'
@@ -8,6 +9,7 @@ require_relative '../helpers/constants.rb'
 # Globals
 @db_conn = nil
 @secgen_args = ''
+@ranges_in_table = nil
 
 # Displays secgen_batch usage data
 def usage
@@ -76,9 +78,10 @@ end
 
 def get_delete_opts
   delete_options = misc_opts + [['--id', GetoptLong::REQUIRED_ARGUMENT],
-                                ['--all', GetoptLong::OPTIONAL_ARGUMENT]]
+                                ['--all', GetoptLong::OPTIONAL_ARGUMENT],
+                                ['--failed', GetoptLong::OPTIONAL_ARGUMENT]]
   options = parse_opts(GetoptLong.new(*delete_options))
-  if options[:id] == '' and options[:all] == false
+  if options[:id] == '' and options[:all] == false and options[:failed] == false
     Print.err 'Error: The delete command requires an argument.'
     usage
   else
@@ -100,6 +103,8 @@ def parse_opts(opts)
         options[:random_ips] = arg.to_i
       when '--all'
         options[:all] = true
+      when '--failed'
+        options[:failed] = true
       else
         Print.err 'Invalid argument'
         exit(false)
@@ -135,6 +140,11 @@ def start(options)
   # Start in SecGen's ROOT_DIR
   Dir.chdir ROOT_DIR
 
+  # Create directories
+  Dir.mkdir 'log' unless Dir.exists? 'log'
+  FileUtils.mkdir_p 'batch/successful' unless Dir.exists? 'batch/successful'
+  FileUtils.mkdir_p 'batch/failed' unless Dir.exists? 'batch/failed'
+
   # Start the service and call secgen.rb
   current_threads = []
   while true
@@ -143,7 +153,6 @@ def start(options)
         current_job = get_jobs[0]
         job_id = current_job['id']
         update_status(job_id, :running)
-
         secgen_args = current_job['secgen_args']
 
         # execute secgen
@@ -151,26 +160,42 @@ def start(options)
         stdout, stderr, status = Open3.capture3("ruby secgen.rb #{secgen_args}")
         puts "Job #{job_id} Complete"
 
+        # Update job status and back-up paths
         if status.exitstatus == 0
           update_status(job_id, :success)
           log_prefix = ''
+          backup_path = 'batch/successful/'
         else
           update_status(job_id, :error)
           log_prefix = 'ERROR_'
+          backup_path = 'batch/failed/'
         end
 
-        # Log output
-        Dir.mkdir 'log' unless Dir.exists? 'log'
-        project_path = stdout.split('Creating project: ')[1].split('...')[0]
-        project_id = project_path.split('projects/')[1]
-        log = File.new("log/#{log_prefix}#{project_id}", 'w')
-        log.write("SecGen project path::: #{project_path}\n\n\n")
-        log.write("SecGen arguments::: #{secgen_args}\n\n\n")
-        log.write("SecGen output::: \n\n\n")
-        log.write(stdout)
-        log.write("\n\n\nGenerator local output::: \n\n\n")
-        log.write(stderr)
-        log.close
+        # Get project data from stderr
+        stderr_project_split = stdout.split('Creating project: ')
+        if stderr_project_split.size > 1
+          project_path = stderr_project_split[1].split('...')[0]
+          project_id = project_path.split('projects/')[1]
+
+          # Log output
+          log_name = "#{log_prefix}#{project_id}"
+          log_path = "log/#{log_name}"
+          log = File.new(log_path, 'w')
+          log.write("SecGen project path::: #{project_path}\n\n\n")
+          log.write("SecGen arguments::: #{secgen_args}\n\n\n")
+          log.write("SecGen output::: \n\n\n")
+          log.write(stdout)
+          log.write("\n\n\nGenerator local output::: \n\n\n")
+          log.write(stderr)
+          log.close
+
+          # Back up project and log file
+          FileUtils.cp_r(project_path, backup_path)
+          FileUtils.cp(log_path, (backup_path + project_id + '/' + log_name))
+        else
+          Print.err("Fatal error on job #{job_id}: SecGen crashed before project creation.")
+          Print.err('Check your scenario file.')
+        end
       }
       sleep(1)
     else
@@ -196,6 +221,8 @@ end
 def delete(options)
   if options[:id] != ''
     delete_id(options[:id])
+  elsif options[:failed]
+    delete_failed
   elsif options[:all]
     delete_all
   end
@@ -237,6 +264,17 @@ def update_status(job_id, status)
   @db_conn.exec_prepared(statement,[status_enum[status], job_id])
 end
 
+def delete_failed
+  Print.info 'Are you sure you want to DELETE failed jobs from the queue table? [y/N]'
+  input = STDIN.gets.chomp
+  if input == 'Y' or input == 'y'
+    Print.info "'Deleting all jobs with status == 'error' from Queue table"
+    @db_conn.exec_params("DELETE FROM queue WHERE status = 'error';")
+  else
+    exit
+  end
+end
+
 def delete_all
   Print.info 'Are you sure you want to DELETE all jobs from the queue table? [y/N]'
   input = STDIN.gets.chomp
@@ -255,20 +293,52 @@ def delete_id(id)
   @db_conn.exec_prepared(statement, [id])
 end
 
+def secgen_arg_network_ranges(secgen_args)
+  ranges_in_arg = []
+  split_args = secgen_args.split(' ')
+  network_ranges_index = split_args.find_index('--network-ranges')
+  if network_ranges_index != nil
+    range = split_args[network_ranges_index + 1]
+    if range.include?(',')
+      range.split(',').each { |split_range| ranges_in_arg << split_range }
+    else
+      ranges_in_arg << range
+    end
+  end
+  ranges_in_arg
+end
+
 def generate_range_arg(options)
   range_arg = ''
   if options.has_key? :random_ips
-    network_ranges = []
+
+    # Check if there are jobs in the DB containing ips. Assign once so that repeated calls don't get added to the list.
+    if @ranges_in_table == nil
+      @ranges_in_table = []
+
+      # Promt to see if we're excluding ranges in the table
+      Print.info 'Do you want to exclude ranges in the database from your random IP generation? [y/N]'
+      input = STDIN.gets.chomp
+      if input == 'Y' or input == 'y'
+        table_entries = select_all
+        table_entries.each { |job|
+          @ranges_in_table += secgen_arg_network_ranges(job['secgen_args'])
+        }
+      end
+    end
+
+    generated_network_ranges = []
     scenario_networks_qty = options[:random_ips]
     scenario_networks_qty.times {
       range = generate_range
       # Check for uniqueness
-      while network_ranges.include?(range)
+      while @ranges_in_table.include?(range)
         range = generate_range
       end
-      network_ranges << range
+      @ranges_in_table << range
+      generated_network_ranges << range
     }
-    random_ip_string = network_ranges.join(',')
+    random_ip_string = generated_network_ranges.join(',')
     range_arg = "--network-ranges #{random_ip_string} "
   end
   range_arg