From 42adbcb853c9808d5dd70a7b42abbc01ebdfd41d Mon Sep 17 00:00:00 2001 From: Jjk422 Date: Fri, 19 Aug 2016 19:35:47 +0100 Subject: [PATCH] Added extra command line options to modify generated vms, need to be separated into a ruby class to ensure max values are not set (max ram larger then system ram) and to make sure that large ruby blocks are not in Vagrantfile.erb. Options added are: --memory-per-vm --total-memory --max-cpu-cores --max-cpu-usage --- lib/output/project_files_creator.rb | 4 +-- lib/templates/Vagrantfile.erb | 19 +++++++------- secgen.rb | 39 ++++++++++++++++++++++------- 3 files changed, 41 insertions(+), 21 deletions(-) diff --git a/lib/output/project_files_creator.rb b/lib/output/project_files_creator.rb index 273deb0cb..696af4880 100644 --- a/lib/output/project_files_creator.rb +++ b/lib/output/project_files_creator.rb @@ -13,7 +13,7 @@ class ProjectFilesCreator # @param [Object] systems list of systems that have been defined and randomised # @param [Object] out_dir the directory that the project output should be stored into # @param [Object] scenario the file path used to as a basis - def initialize(systems, out_dir, scenario, gui_output) + def initialize(systems, out_dir, scenario, options) @systems = systems @out_dir = out_dir @@ -24,7 +24,7 @@ class ProjectFilesCreator end @scenario = scenario @time = Time.new.to_s - @gui_output = gui_output + @options = options end diff --git a/lib/templates/Vagrantfile.erb b/lib/templates/Vagrantfile.erb index 38bc0abe6..88d20f37d 100644 --- a/lib/templates/Vagrantfile.erb +++ b/lib/templates/Vagrantfile.erb @@ -12,16 +12,15 @@ Vagrant.configure(VAGRANTFILE_API_VERSION) do |config| config.vm.define "<%= system.name %>" do |<%= system.name %>| # REMOVE: < %= system.name %>.vm.synced_folder "< %= MOUNT_DIR %>", "/mount" config.vm.provider :virtualbox do |vb| - <%= - if @gui_output == true - "vb.gui = true" - else - "vb.gui = false" - "vb.customize ['modifyvm', :id, '--pae', 'on']" - "vb.customize ['modifyvm', :id, '--hwvirtex', 'off']" - "vb.customize ['modifyvm', :id, '--vtxvpid', 'off']" - end - %> + <%= (@options.has_key? :gui_output && @options[:gui_output])?"vb.gui = true":" + vb.gui = false + vb.customize ['modifyvm', :id, '--pae', 'on'] + vb.customize ['modifyvm', :id, '--hwvirtex', 'off'] + vb.customize ['modifyvm', :id, '--vtxvpid', 'off']" %> + <%= (@options.has_key? :memory_per_vm)?"vb.memory = #{@options[:memory_per_vm]}":"vb.memory = 512" %> + <%= (@options.has_key? :total_memory)?"vb.memory = #{@options[:total_memory]}/#{@systems.length}":"vb.memory = 512" %> + <%= (@options.has_key? :max_cpu_cores)?"vb.cpus = #{@options[:max_cpu_cores]}":"vb.cpus = 4" %> + <%= (@options.has_key? :max_cpu_usage)?"vb.customize ['modifyvm', :id, '--cpuexecutioncap', '#{@options[:max_cpu_usage]}']":"vb.customize ['modifyvm', :id, '--cpuexecutioncap', '100']" %> end # SecGen modules <% system.module_selections.each do |selected_module| -%> diff --git a/secgen.rb b/secgen.rb index 5cb2809ed..a8aca262b 100644 --- a/secgen.rb +++ b/secgen.rb @@ -31,7 +31,7 @@ end # Builds the vagrant configuration file based on a scenario file # @return build_number [Integer] Current project's build number -def build_config(scenario, out_dir, gui_output) +def build_config(scenario, out_dir, options) Print.info 'Reading configuration file for virtual machines you want to create...' # read the scenario file describing the systems, which contain vulnerabilities, services, etc # this returns an array/hashes structure @@ -69,7 +69,7 @@ def build_config(scenario, out_dir, gui_output) Print.info "Creating project: #{out_dir}..." # create's vagrant file / report a starts the vagrant installation' - creator = ProjectFilesCreator.new(systems, out_dir, scenario, gui_output) + creator = ProjectFilesCreator.new(systems, out_dir, scenario, options) creator.write_files Print.info 'Project files created.' @@ -84,8 +84,8 @@ def build_vms(project_dir) end # Runs methods to run and configure a new vm from the configuration file -def run(scenario, project_dir, gui_output) - build_config(scenario, project_dir, gui_output) +def run(scenario, project_dir, options) + build_config(scenario, project_dir, options) build_vms(project_dir) end @@ -106,12 +106,16 @@ opts = GetoptLong.new( [ '--help', '-h', GetoptLong::NO_ARGUMENT ], [ '--project', '-p', GetoptLong::REQUIRED_ARGUMENT ], [ '--scenario', '-s', GetoptLong::REQUIRED_ARGUMENT ], - [ '--gui-output', '-g', GetoptLong::NO_ARGUMENT] + [ '--gui-output', '-g', GetoptLong::NO_ARGUMENT], + [ '--memory-per-vm', GetoptLong::REQUIRED_ARGUMENT], + [ '--total-memory', GetoptLong::REQUIRED_ARGUMENT], + [ '--max-cpu-cores', GetoptLong::REQUIRED_ARGUMENT], + [ '--max-cpu-usage', GetoptLong::REQUIRED_ARGUMENT], ) scenario = SCENARIO_XML project_dir = nil -gui_output = false +options = {} # process option arguments opts.each do |opt, arg| @@ -124,7 +128,24 @@ opts.each do |opt, arg| project_dir = arg; when '--gui-output' Print.info "Gui output set (virtual machines will be spawned)" - gui_output = true; + + options[:gui_output] = true; + when '--memory-per-vm' + Print.info "Memory per vm set to #{arg}" + options[:memory_per_vm] = arg + + when '--total-memory' + Print.info "Total memory to be used set to #{arg}" + options[:total_memory] = arg + + when '--max-cpu-cores' + Print.info "Number of cpus to be used set to #{arg}" + options[:max_cpu_cores] = arg + + when '--max-cpu-usage' + Print.info "Max CPU usage set to #{arg}" + options[:max_cpu_usage] = arg + else Print.err "Argument not valid: #{arg}" usage @@ -143,10 +164,10 @@ end case ARGV[0] when 'run', 'r' project_dir = default_project_dir unless project_dir - run(scenario, project_dir, gui_output) + run(scenario, project_dir, options) when 'build-project', 'p' project_dir = default_project_dir unless project_dir - build_config(scenario, project_dir, gui_output) + build_config(scenario, project_dir, options) when 'build-vms', 'v' if project_dir build_vms(project_dir)