From 08aad783dffe453c1feba678487f9b3d8bcd4047 Mon Sep 17 00:00:00 2001 From: thomashaw Date: Fri, 20 Oct 2017 14:26:16 +0100 Subject: [PATCH] Minor cleanup --- lib/batch/README.md | 12 +++++++++++- lib/batch/batch_secgen.rb | 5 ++--- lib/objects/system.rb | 29 ++++++++--------------------- lib/output/project_files_creator.rb | 14 ++++++++------ secgen.rb | 4 ++-- 5 files changed, 31 insertions(+), 33 deletions(-) diff --git a/lib/batch/README.md b/lib/batch/README.md index 70fb40fbe..9da29e4d8 100644 --- a/lib/batch/README.md +++ b/lib/batch/README.md @@ -1,6 +1,16 @@ # Batch Processing with SecGen -TODO: Write a brief introduction +Generating multiple VMs in a batch is now possible through the use of batch_secgen, a ruby script which uses postgresql +as a job queue to mass-create VMs with SecGen. + +There are helper commands available to add jobs, list jobs in the table, remove jobs, and reset the status of jobs from 'running' or 'error' to 'todo'. + +When adding multiple jobs to the queue, it is possible to prefix the VM names with unique strings. +The example below demonstrates adding 3 copies of the flawed_fortress scenario, which results in the VM names being prefixed with 'tom_', 'cliffe_', and 'aimee_'. + +``` +ruby batch_secgen.rb add --instances tom,cliffe,aimee --- -s scenarios/ctf/flawed_fortress_1.xml r +``` ## Initialise the Database diff --git a/lib/batch/batch_secgen.rb b/lib/batch/batch_secgen.rb index d10e5121a..a7e112d2a 100644 --- a/lib/batch/batch_secgen.rb +++ b/lib/batch/batch_secgen.rb @@ -250,7 +250,6 @@ def reset(options) end end - def delete(options) if options[:id] != '' delete_id(options[:id]) @@ -365,7 +364,7 @@ def generate_range_arg(options) if @ranges_in_table == nil @ranges_in_table = [] - # Promt to see if we're excluding ranges in the table + # Prompt 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 == '' or input == 'Y' or input == 'y' @@ -430,4 +429,4 @@ case ARGV[0] delete(get_delete_opts) else usage -end +end \ No newline at end of file diff --git a/lib/objects/system.rb b/lib/objects/system.rb index 17874416f..222dbd65c 100644 --- a/lib/objects/system.rb +++ b/lib/objects/system.rb @@ -8,7 +8,6 @@ class System attr_accessor :module_selectors # (filters) attr_accessor :module_selections # (after resolution) attr_accessor :num_actioned_module_conflicts - attr_accessor :system_networks attr_accessor :options #(command line options hash) # Initalizes System object @@ -21,22 +20,19 @@ class System self.module_selectors = module_selectors self.module_selections = [] self.num_actioned_module_conflicts = 0 - self.system_networks = [] - self.options = {} end # selects from the available modules, based on the selection filters that have been specified # @param [Object] available_modules all available modules (vulnerabilities, services, bases) - # @param [Object] opts command line options hash + # @param [Object] options command line options hash # @return [Object] the list of selected modules - def resolve_module_selection(available_modules, opts) - @options = opts + def resolve_module_selection(available_modules, options) retry_count = 0 # Replace $IP_addresses with options ip_ranges if required begin - if @options[:ip_ranges] and $datastore['IP_addresses'] and !$datastore['replaced_ranges'] - unused_opts_ranges = @options[:ip_ranges].clone + if options[:ip_ranges] and $datastore['IP_addresses'] and !$datastore['replaced_ranges'] + unused_opts_ranges = options[:ip_ranges].clone option_range_map = {} # k = ds_range, v = opts_range new_ip_addresses = [] @@ -142,25 +138,26 @@ class System # filter to those that satisfy the attribute filters # select based on selected type, access, cve... - search_list.delete_if { |module_for_possible_exclusion| + search_list.delete_if{|module_for_possible_exclusion| !module_for_possible_exclusion.matches_attributes_requirement(required_attributes) } Print.verbose "Filtered to modules matching: #{required_attributes.inspect} ~= (n=#{search_list.size})" # remove non-options due to conflicts - search_list.delete_if { |module_for_possible_exclusion| + search_list.delete_if{|module_for_possible_exclusion| check_conflicts_with_list(module_for_possible_exclusion, previously_selected_modules) } # check if modules need to be unique # write_module_path_to_datastore if write_module_path_to_datastore != nil && $datastore[write_module_path_to_datastore] != nil - search_list.delete_if { |module_for_possible_exclusion| + search_list.delete_if{|module_for_possible_exclusion| ($datastore[write_module_path_to_datastore] ||=[]).include? module_for_possible_exclusion.module_path } Print.verbose "Filtering to remove non-unique #{$datastore[write_module_path_to_datastore]} ~= (n=#{search_list.size})" end + if search_list.length == 0 raise 'failed' Print.err 'Could not find a matching module. Please check the scenario specification' @@ -435,14 +432,4 @@ class System modules_to_add end - def get_networks - if (self.system_networks = []) # assign the networks - self.module_selections.each do |mod| - if mod.module_type == 'network' - self.system_networks << mod - end - end - end - self.system_networks - end end \ No newline at end of file diff --git a/lib/output/project_files_creator.rb b/lib/output/project_files_creator.rb index 95cfe5736..38a6b60c6 100644 --- a/lib/output/project_files_creator.rb +++ b/lib/output/project_files_creator.rb @@ -39,13 +39,13 @@ class ProjectFilesCreator if File.exists? "#{@out_dir}/Vagrantfile" or File.exists? "#{@out_dir}/puppet" dest_dir = "#{@out_dir}/MOVED_#{Time.new.strftime("%Y%m%d_%H%M")}" Print.warn "Project already built to this directory -- moving last build to: #{dest_dir}" - Dir.glob("#{@out_dir}/**/*").select { |f| File.file?(f) }.each do |f| + Dir.glob( "#{@out_dir}/**/*" ).select { |f| File.file?( f ) }.each do |f| dest = "#{dest_dir}/#{f}" - FileUtils.mkdir_p(File.dirname(dest)) + FileUtils.mkdir_p( File.dirname( dest ) ) if f =~ /\.vagrant/ - FileUtils.cp(f, dest) + FileUtils.cp( f, dest ) else - FileUtils.mv(f, dest) + FileUtils.mv( f, dest ) end end end @@ -80,7 +80,7 @@ class ProjectFilesCreator if File.file? packerfile_path Print.info "Would you like to use the packerfile to create the packerfile from the given url (y/n)" - (Print.info "Exiting as vagrant needs the basebox to continue"; exit) unless ['y', 'yes'].include?(STDIN.gets.chomp.downcase) + (Print.info "Exiting as vagrant needs the basebox to continue"; exit) unless ['y','yes'].include?(STDIN.gets.chomp.downcase) Print.std "Packerfile #{packerfile_path.split('/').last} found, building basebox #{url.split('/').last} via packer" template_based_file_write(packerfile_path, packerfile_path.split(/.erb$/).first) @@ -156,11 +156,13 @@ class ProjectFilesCreator end # Resolves the network based on the scenario and ip_range. +# In the case that both command-line --network-ranges and datastores are provided, we have already handled the replacement of the ranges in the datastore. +# Because of this we prioritise datastore['IP_address'], then command line options (i.e. when no datastore is used, but the --network-ranges are passed), then the default network module's IP range. def resolve_network(network_module) current_network = network_module scenario_ip_range = network_module.attributes['range'].first - # Use datastore IP_address if we have them + # Prioritise datastore IP_address if current_network.received_inputs.include? 'IP_address' ip_address = current_network.received_inputs['IP_address'].first elsif @options.has_key? :ip_ranges diff --git a/secgen.rb b/secgen.rb index 5ced0ca26..25c4a9b61 100644 --- a/secgen.rb +++ b/secgen.rb @@ -188,7 +188,7 @@ end # # @author Jason Keighley # @return [Void] -def create_ewf_image(drive_path, image_output_location) +def create_ewf_image(drive_path ,image_output_location) ## Make E01 image Print.info "Creating E01 image with path #{image_output_location}.E01" Print.info 'This may take a while:' @@ -469,4 +469,4 @@ case ARGV[0] Print.err "Command not valid: #{ARGV[0]}" usage exit -end +end \ No newline at end of file