From 60d3604efd7402bdb40e31936fef3ca29f6b65cf Mon Sep 17 00:00:00 2001
From: thomashaw <thomashaw@gmail.com>
Date: Wed, 16 Feb 2022 15:28:59 +0000
Subject: [PATCH] Dynamic generation of goal flags (and some cleanup, removing
 goal_flags etc.).

---
 lib/output/project_files_creator.rb           | 87 +++++++++++++------
 .../xml_alertaction_config_generator.rb       | 28 +++---
 lib/output/xml_marker_generator.rb            | 11 ++-
 .../goal_flag_hacktivity/secgen_metadata.xml  |  2 +-
 .../goal_message_host/secgen_metadata.xml     |  2 +-
 .../ctf/misc/hidden_file/secgen_metadata.xml  | 10 +--
 .../unix/ctf/pwn/symlinks/secgen_metadata.xml | 15 +---
 .../nfs/nfs_overshare/secgen_metadata.xml     | 14 +--
 .../examples/auto_grading/experiment_aaa.xml  |  9 --
 scenarios/examples/auto_grading_example.xml   |  5 --
 scenarios/tests/test_scenario.xml             |  5 --
 11 files changed, 87 insertions(+), 101 deletions(-)

diff --git a/lib/output/project_files_creator.rb b/lib/output/project_files_creator.rb
index e1795c0cb..41939953a 100644
--- a/lib/output/project_files_creator.rb
+++ b/lib/output/project_files_creator.rb
@@ -10,16 +10,16 @@ require 'librarian'
 require 'zip/zip'
 
 class ProjectFilesCreator
-# Creates project directory, uses .erb files to create a report and the vagrant file that will be used
-# to create the virtual machines
+  # Creates project directory, uses .erb files to create a report and the vagrant file that will be used
+  # to create the virtual machines
   @systems
   @currently_processing_system
   @scenario_networks
   @option_range_map
 
-# @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
+  # @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, options)
     @systems = systems
     @out_dir = out_dir
@@ -32,22 +32,23 @@ class ProjectFilesCreator
     @scenario = scenario
     @time = Time.new.to_s
     @options = options
-    @scenario_networks = Hash.new {|h, k| h[k] = 1}
+    @scenario_networks = Hash.new { |h, k| h[k] = 1 }
     @option_range_map = {}
-
+    @number_of_goals = -1
+    @extra_flags = []
     # Packer builder type
     @builder_type = @options.has_key?(:esxi_url) ? :vmware_iso : :virtualbox_iso
     resolve_interp_strings
   end
 
-# Generate all relevant files for the project
+  # Generate all relevant files for the project
   def write_files
     # when writing to a project that already contains a project, move everything out the way,
     # and keep the Vagrant config, so that existing VMs can be re-provisioned/updated
     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%S")}"
       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))
         if f =~ /\.vagrant/
@@ -181,7 +182,16 @@ class ProjectFilesCreator
         # Get the config json object from the alert_actioner
         aa_confs = JSON.parse(system.get_module('analysis_alert_action_server').received_inputs['aaa_config'][0])['aa_configs']
         xml_aa_conf_file = "#{aa_conf_dir}#{@out_dir.split('/')[-1]}.xml"
-        xml_aa_conf_generator = XmlAlertActionConfigGenerator.new(@systems, @scenario, @time, aa_confs, @options)
+
+        # Calculate the number of goals in the scenario and generate flags to insert into the alert action and hints XML generators
+        n_goals = get_total_number_of_goals
+
+        i = 0
+        (1..n_goals).each { |_|
+          @extra_flags << "flag{#{SecureRandom.hex}}"
+        }
+
+        xml_aa_conf_generator = XmlAlertActionConfigGenerator.new(@systems, @scenario, @time, aa_confs, @options, @extra_flags)
         xml = xml_aa_conf_generator.output
         Print.std "AlertActioner: Creating alert_actioner configuration file: #{xml_aa_conf_file}"
         write_data_to_file(xml, xml_aa_conf_file)
@@ -209,7 +219,7 @@ class ProjectFilesCreator
     # Create the marker xml file
     x2file = "#{@out_dir}/#{FLAGS_FILENAME}"
 
-    xml_marker_generator = XmlMarkerGenerator.new(@systems, @scenario, @time)
+    xml_marker_generator = XmlMarkerGenerator.new(@systems, @scenario, @time, @extra_flags)
     xml = xml_marker_generator.output
     Print.std "Creating flags and hints file: #{x2file}"
     write_data_to_file(xml, x2file)
@@ -223,10 +233,10 @@ class ProjectFilesCreator
 
     # zip up the CTFd export
     begin
-      Zip::ZipFile.open(ctfdfile, Zip::ZipFile::CREATE) {|zipfile|
+      Zip::ZipFile.open(ctfdfile, Zip::ZipFile::CREATE) { |zipfile|
         zipfile.mkdir("db")
         ctfd_files.each do |ctfd_file_name, ctfd_file_content|
-          zipfile.get_output_stream("db/#{ctfd_file_name}") {|f|
+          zipfile.get_output_stream("db/#{ctfd_file_name}") { |f|
             f.print ctfd_file_content
           }
         end
@@ -264,8 +274,8 @@ class ProjectFilesCreator
   end
 
 
-# Goal string interpolation for the whole system
-# prior to calling the rule generator multiple times
+  # Goal string interpolation for the whole system
+  # prior to calling the rule generator multiple times
   def resolve_interp_strings
     @systems.each do |system|
       system.module_selections.each do |module_selection|
@@ -277,8 +287,8 @@ class ProjectFilesCreator
     end
   end
 
-# @param [Object] template erb path
-# @param [Object] filename file to write to
+  # @param [Object] template erb path
+  # @param [Object] filename file to write to
   def template_based_file_write(template, filename)
     template_out = ERB.new(File.read(template), 0, '<>-')
 
@@ -292,9 +302,9 @@ class ProjectFilesCreator
     end
   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.
+  # 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
@@ -311,7 +321,7 @@ class ProjectFilesCreator
       else
         # Remove options_ips that have already been used
         options_ips = @options[:ip_ranges]
-        options_ips.delete_if {|ip| @option_range_map.has_value? ip}
+        options_ips.delete_if { |ip| @option_range_map.has_value? ip }
         @option_range_map[scenario_ip_range] = options_ips.first
         ip_range = options_ips.first
       end
@@ -336,14 +346,14 @@ class ProjectFilesCreator
     split_ip.join('.')
   end
 
-# Replace 'network' with 'snoop' where the system name contains snoop
+  # Replace 'network' with 'snoop' where the system name contains snoop
   def get_ovirt_network_name(system_name, network_name)
     split_name = network_name.split('-')
     split_name[1] = 'snoop' if system_name.include? 'snoop'
     split_name.join('-')
   end
 
-# Determine how much memory the system requires for Vagrantfile
+  # Determine how much memory the system requires for Vagrantfile
   def resolve_memory(system)
     if @options.has_key? :memory_per_vm
       memory = @options[:memory_per_vm]
@@ -363,10 +373,31 @@ class ProjectFilesCreator
     memory
   end
 
-# Returns binding for erb files (access to variables in this classes scope)
-# @return binding
-  def get_binding
-    binding
+  def get_total_number_of_goals
+    if @number_of_goals == -1
+      n = 0
+      @systems.each do |system|
+        # calculate number of system goals
+        if system.goals != []
+          n = n + system.goals.size
+        end
+        # calculate number of module goals on this system
+        system.module_selections.each do |module_selection|
+          if module_selection.goals != []
+            n = n + module_selection.goals.size
+          end
+        end
+      end
+      @number_of_goals = n
+      Print.info("Number of goals " + @number_of_goals.to_s)
+    end
+    @number_of_goals
   end
 
-end
+    # Returns binding for erb files (access to variables in this classes scope)
+    # @return binding
+    def get_binding
+      binding
+    end
+
+  end
diff --git a/lib/output/xml_alertaction_config_generator.rb b/lib/output/xml_alertaction_config_generator.rb
index 4014e2072..eed6689be 100644
--- a/lib/output/xml_alertaction_config_generator.rb
+++ b/lib/output/xml_alertaction_config_generator.rb
@@ -9,14 +9,14 @@ class XmlAlertActionConfigGenerator
   # @param [Object] scenario the scenario file used to generate
   # @param [Object] time the current time as a string
   # @param [Array[Hash]] the alert_actioner configuration settings (list of aa_conf JSON hashes)
-  def initialize(systems, scenario, time, aa_confs, options)
+  def initialize(systems, scenario, time, aa_confs, options, goal_flags)
     @systems = systems
     @scenario = scenario
     @time = time
     @aa_confs = aa_confs
     @options = options
     @alert_actions = []
-    @goal_flags = []
+    @goal_flags = goal_flags.clone
   end
 
   # outputs a XML AlertActioner configuration file
@@ -30,12 +30,12 @@ class XmlAlertActionConfigGenerator
     Print.info 'AlertActioner: Creating alert actions from aa_conf.'
     @aa_confs.each do |aa_conf|
       if aa_conf['mapping_type']
+        Print.info("**** Generating AlertActions for #{aa_conf['mapping_type']}")
         case aa_conf['mapping_type']
-        when 'all_goal_flags_to_hacktivity'
-          all_goal_flags_to_hacktivity(aa_conf)
-
-        when 'all_goal_messages_to_host'
-          all_goal_message_host(aa_conf)
+        when 'hacktivity_flags'
+          generate_hacktivity_flags(aa_conf)
+        when 'message_host'
+          generate_message_host(aa_conf)
         else
           Print.err("AlertActioner Config: Invalid mapping type #{aa_conf['mapping_type']}")
           exit(1)
@@ -49,7 +49,7 @@ class XmlAlertActionConfigGenerator
     end
   end
 
-  def all_goal_message_host(aa_conf)
+  def generate_message_host(aa_conf)
     @systems.each do |system|
       system.module_selections.each do |module_selection|
         module_name = module_selection.module_path_end
@@ -72,14 +72,11 @@ class XmlAlertActionConfigGenerator
     end
   end
 
-  def all_goal_flags_to_hacktivity(aa_conf)
-    Print.info("**** sending all_goal_flags_to_hacktivity ****")
+  def generate_hacktivity_flags(aa_conf)
     auto_grader_hostname = get_auto_grader_hostname
     Print.info("auto_grader_hostname: " + auto_grader_hostname)
     Print.info("systems.size: " + @systems.size.to_s)
 
-    @goal_flags = @goal_flags + $datastore['goal_flags']
-
     @systems.each do |system|
       Print.info("System goals: " + system.goals.to_s)
       if system.goals != []
@@ -97,7 +94,6 @@ class XmlAlertActionConfigGenerator
   end
 
   def get_web_alertactions(aa_conf, name, goals, hostname, auto_grader_hostname)
-    Print.info("**** get_web_alertactions() ****")
     alert_actions = []
 
     # Validate whether there are an equal number of goals and goal_flags + warn / error here if not...
@@ -119,8 +115,8 @@ class XmlAlertActionConfigGenerator
           }
         end
     else
-      Print.err("goals: " + goals)
-      Print.err("goal_flags: " + @goal_flags)
+      Print.err("goals: " + goals.to_s)
+      Print.err("goal_flags: " + @goal_flags.to_s)
     end
     alert_actions
   end
@@ -138,7 +134,7 @@ class XmlAlertActionConfigGenerator
         xml.comment "#{@time}"
         xml.comment "Based on a fulfilment of scenario: #{@scenario}"
 
-        @alert_actions.each {|alert_action|
+        @alert_actions.each { |alert_action|
           xml.alertaction {
             xml.alert_name alert_action['alert_name']
             case alert_action['action_type']
diff --git a/lib/output/xml_marker_generator.rb b/lib/output/xml_marker_generator.rb
index a0fc84a88..471080fc0 100644
--- a/lib/output/xml_marker_generator.rb
+++ b/lib/output/xml_marker_generator.rb
@@ -6,10 +6,11 @@ class XmlMarkerGenerator
   # @param [Object] systems the list of systems
   # @param [Object] scenario the scenario file used to generate
   # @param [Object] time the current time as a string
-  def initialize(systems, scenario, time)
+  def initialize(systems, scenario, time, extra_flags)
     @systems = systems
     @scenario = scenario
     @time = time
+    @extra_flags = extra_flags.clone
   end
 
   # outputs a XML marker file that can be used to mark flags and provide hints
@@ -30,6 +31,14 @@ class XmlMarkerGenerator
             xml.system_name system.name
             xml.platform system.module_selections.first.attributes['platform'].first
 
+            # add extra flags to the auto_grading_server
+            if system.name == 'auto_grading_server'
+              extra_flags_n = @extra_flags.size
+              (1..extra_flags_n).each { |_|
+                xml.flag(@extra_flags.pop)
+              }
+            end
+
             system.module_selections.each { |selected_module|
 
               # start by finding a flag, and work the way back providing hints
diff --git a/modules/generators/structured_content/alert_actioner_config/goal_flag_hacktivity/secgen_metadata.xml b/modules/generators/structured_content/alert_actioner_config/goal_flag_hacktivity/secgen_metadata.xml
index 8cd83b615..168ba63c0 100644
--- a/modules/generators/structured_content/alert_actioner_config/goal_flag_hacktivity/secgen_metadata.xml
+++ b/modules/generators/structured_content/alert_actioner_config/goal_flag_hacktivity/secgen_metadata.xml
@@ -23,7 +23,7 @@
 
   <!-- Mapping types tell the project_file_creator what to do.-->
   <default_input into="mapping_type">
-    <value>all_goal_flags_to_hacktivity</value>
+    <value>hacktivity_flags</value>
   </default_input>
 
   <output_type>json</output_type>
diff --git a/modules/generators/structured_content/alert_actioner_config/goal_message_host/secgen_metadata.xml b/modules/generators/structured_content/alert_actioner_config/goal_message_host/secgen_metadata.xml
index a0802f415..a2acc860a 100644
--- a/modules/generators/structured_content/alert_actioner_config/goal_message_host/secgen_metadata.xml
+++ b/modules/generators/structured_content/alert_actioner_config/goal_message_host/secgen_metadata.xml
@@ -38,7 +38,7 @@
 
   <!-- Mapping types tell the project_file_creator what to do.-->
   <default_input into="mapping_type">
-    <value>all_goal_messages_to_host</value>
+    <value>message_host</value>
   </default_input>
 
   <output_type>json</output_type>
diff --git a/modules/vulnerabilities/unix/ctf/misc/hidden_file/secgen_metadata.xml b/modules/vulnerabilities/unix/ctf/misc/hidden_file/secgen_metadata.xml
index 12eb11c78..717a37eac 100644
--- a/modules/vulnerabilities/unix/ctf/misc/hidden_file/secgen_metadata.xml
+++ b/modules/vulnerabilities/unix/ctf/misc/hidden_file/secgen_metadata.xml
@@ -24,14 +24,12 @@
 
   <goals>
     <read_file>
-      <file_path>#{['flag_path'][0]}</file_path>
+      <file_path>#{['file_path_to_leak'][0]}</file_path>
     </read_file>
   </goals>
 
   <read_fact>file_path_to_leak</read_fact>
   <read_fact>strings_to_leak</read_fact>
-  <read_fact>flag_path</read_fact>
-  <read_fact>goal_flags</read_fact>
 
   <default_input into="file_path_to_leak">
     <value>/.secret_file</value>
@@ -39,12 +37,6 @@
   <default_input into="strings_to_leak">
     <generator type="flag_generator"/>
   </default_input>
-  <default_input into="flag_path">
-    <value>#{['file_path_to_leak'][0]}</value>
-  </default_input>
-  <default_input into="goal_flags">
-    <value>#{['strings_to_leak'][0]}</value>
-  </default_input>
 
   <hint>Try reading the man page for ls with 'man ls'. Linux hidden files have names which begin with a dot. e.g.
     .file
diff --git a/modules/vulnerabilities/unix/ctf/pwn/symlinks/secgen_metadata.xml b/modules/vulnerabilities/unix/ctf/pwn/symlinks/secgen_metadata.xml
index c0cbb478a..20fea17c8 100644
--- a/modules/vulnerabilities/unix/ctf/pwn/symlinks/secgen_metadata.xml
+++ b/modules/vulnerabilities/unix/ctf/pwn/symlinks/secgen_metadata.xml
@@ -20,16 +20,15 @@
 
   <goals>
     <read_file>
-      <file_path>#{['flag_path'][0]}</file_path>
+      <file_path>/home/#{['accounts'][0]['username']}/#{['accounts'][0]['leaked_filenames'][0]}</file_path>
     </read_file>
     <read_file>
-      <file_path>#{['flag_path'][1]}</file_path>
+      <file_path>/etc/shadow</file_path>
     </read_file>
   </goals>
 
   <read_fact>username</read_fact>
   <read_fact>flag_path</read_fact>
-  <read_fact>goal_flags</read_fact>
 
   <default_input into="accounts">
     <generator type="account">
@@ -49,16 +48,6 @@
     </generator>
   </default_input>
 
-  <default_input into="flag_path">
-    <value>/etc/shadow</value>
-    <value>/home/#{['accounts'][0]['username']}/#{['accounts'][0]['leaked_filenames'][0]}</value>
-  </default_input>
-
-  <default_input into="goal_flags">
-    <value>#{['accounts'][0]['strings_to_leak'][0]}</value>
-    <value>#{['accounts'][0]['strings_to_leak'][1]}</value>
-  </default_input>
-
   <requires>
     <module_path>utilities/unix/system/accounts</module_path>
   </requires>
diff --git a/modules/vulnerabilities/unix/nfs/nfs_overshare/secgen_metadata.xml b/modules/vulnerabilities/unix/nfs/nfs_overshare/secgen_metadata.xml
index e052da8e4..47d0131fb 100644
--- a/modules/vulnerabilities/unix/nfs/nfs_overshare/secgen_metadata.xml
+++ b/modules/vulnerabilities/unix/nfs/nfs_overshare/secgen_metadata.xml
@@ -16,7 +16,7 @@
 
   <goals>
     <read_file>
-      <file_path>#{['flag_path']}</file_path>
+      <file_path>#{['storage_directory'][0]}/#{['leaked_filenames'][0]}</file_path>
     </read_file>
   </goals>
 
@@ -26,7 +26,6 @@
   <read_fact>storage_directory</read_fact>
   <!-- Auto grading -->
   <read_fact>flag_path</read_fact>
-  <read_fact>goal_flags</read_fact>
 
   <!--if an input is not specified in the scenario-->
   <default_input into="strings_to_leak">
@@ -51,17 +50,6 @@
     <generator type="storage_directory_generator"/>
   </default_input>
 
-  <!-- Auto-grading goals -->
-  <default_input into="flag_path">
-    <value>#{['storage_directory'][0]}/#{['leaked_filenames'][0]}</value>
-  </default_input>
-
-  <default_input into="goal_flags">
-    <generator type="flag_generator"/>
-    <generator type="flag_generator"/>
-  </default_input>
-
-
   <!--For testing, you can define default values for whatever values you want through to factor:-->
   <!--<default_input into="harmless_unset_variable">-->
     <!--<value>Plain text from the metadata default, destined for harmless_unset_variable...</value>-->
diff --git a/scenarios/examples/auto_grading/experiment_aaa.xml b/scenarios/examples/auto_grading/experiment_aaa.xml
index a1b5052b9..788353130 100644
--- a/scenarios/examples/auto_grading/experiment_aaa.xml
+++ b/scenarios/examples/auto_grading/experiment_aaa.xml
@@ -144,14 +144,6 @@
     <system_name>auto_grading_server</system_name>
     <base distro="Debian 10" />
 
-    <input into_datastore="goal_flags">
-      <value>flag{static_flag_1}</value>
-      <value>flag{static_flag_2}</value>
-      <value>flag{static_flag_3}</value>
-      <value>flag{static_flag_4}</value>
-      <value>flag{static_flag_5}</value>
-    </input>
-
     <utility module_path=".*handy_cli_tools.*"/>
 
     <service module_path=".*analysis_alert_action_server">
@@ -173,5 +165,4 @@
     </build>
   </system>
 
-
 </scenario>
\ No newline at end of file
diff --git a/scenarios/examples/auto_grading_example.xml b/scenarios/examples/auto_grading_example.xml
index 3a543aace..2fb5424f5 100644
--- a/scenarios/examples/auto_grading_example.xml
+++ b/scenarios/examples/auto_grading_example.xml
@@ -113,11 +113,6 @@
       </read_file>
     </goals>
 
-    <input into_datastore="goal_flags">
-      <generator type="flag_generator"/>
-      <generator type="flag_generator"/>
-    </input>
-
     <utility module_path=".*handy_cli_tools.*"/>
 
     <utility module_path=".*parameterised_accounts">
diff --git a/scenarios/tests/test_scenario.xml b/scenarios/tests/test_scenario.xml
index 930adeaa3..a7bbd5cb4 100644
--- a/scenarios/tests/test_scenario.xml
+++ b/scenarios/tests/test_scenario.xml
@@ -79,11 +79,6 @@
     <system_name>auto_grading_server</system_name>
     <base distro="Debian 10" />
 
-    <input into_datastore="goal_flags">
-      <value>flag{static_flag_1}</value>
-      <value>flag{static_flag_1}</value>
-    </input>
-
     <utility module_path=".*handy_cli_tools.*"/>
 
     <service module_path=".*analysis_alert_action_server">