diff --git a/atomic_red_team/atomic_red_team.rb b/atomic_red_team/atomic_red_team.rb index bad5c204..57546cc6 100755 --- a/atomic_red_team/atomic_red_team.rb +++ b/atomic_red_team/atomic_red_team.rb @@ -152,13 +152,35 @@ class AtomicRedTeam end end - def generate_guids_for_yaml!(path) - guids = [] + def generate_guids_for_yaml!(path, used_guids_file) text = File.read(path) - content = text.gsub(/(?i)^\s*guid:(?!(\s*[a-f0-9]{8}-[a-f0-9]{4}-4[a-f0-9]{3}-[89aAbB][a-f0-9]{3}-[a-f0-9]{12})).*$/) { |m| "auto_generated_guid: #{SecureRandom.uuid}"} + guid = get_unique_guid(used_guids_file) + # add the "auto_generated_guid:" element after the "- name:" element if it isn't already there + c = text.gsub(/(?i)(^(\s*-\s*)name:.*$(?!\s*auto_generated_guid))/) { |m| "#{$1}\n auto_generated_guid:"} + # fill the "auto_generated_guid:" element in if it doesn't contain a guid + content = c.gsub(/(?i)^(\s*auto_generated_guid:)(?!(\s*[a-f0-9]{8}-[a-f0-9]{4}-4[a-f0-9]{3}-[89aAbB][a-f0-9]{3}-[a-f0-9]{12})).*$/) { |m| "#{$1} #{get_unique_guid(used_guids_file)}"} + File.open(path, "w") { |file| file << content } end + def get_unique_guid(used_guids_file) + new_guid = '' + 20.times do |i| # if it takes more than 20 tries to get a unique guid, there must be something else going on + new_guid = SecureRandom.uuid + break unless !is_unique_guid(new_guid, used_guids_file) + end + # add this new unique guid to the used guids file + open(used_guids_file, 'a') { |f| + puts used_guids_file + f.puts new_guid unless new_guid == '' + } + return new_guid + end + + def is_unique_guid(guid, used_guids_file) + return !File.foreach(used_guids_file).grep(/#{guid}/).any? + end + # # Validates that the arguments (specified in "#{arg}" format) in a string # match the input_arguments for a test diff --git a/atomic_red_team/spec.yaml b/atomic_red_team/spec.yaml index e7bfe631..80bf8f6b 100644 --- a/atomic_red_team/spec.yaml +++ b/atomic_red_team/spec.yaml @@ -35,7 +35,7 @@ atomic_tests: # This is the first atomic test # - name: Short name of the test that titles how it tests the technique. # Example: "SourceRecorder via cmd.exe" - auto_generated_guid: # uniquely identifies this test, leave it blank and it will be auto-generated during the Pull Request process on GitHub + auto_generated_guid: # This key and/or it's value will be added by the CI build after submitting a Pull Request description: | Long form description of the test. Markdown is supported so you can **bold** items, create diff --git a/bin/generate-guids.rb b/bin/generate-guids.rb index 368e2afc..53c8dcd0 100644 --- a/bin/generate-guids.rb +++ b/bin/generate-guids.rb @@ -4,16 +4,16 @@ require 'yaml' require 'atomic_red_team' ATOMIC_RED_TEAM = AtomicRedTeam.new -ATOMIC_TEST_TEMPLATE = "#{File.dirname(File.dirname(__FILE__))}/atomic_red_team/atomic_test_template.yaml" +USED_GUIDS_FILE = "#{File.dirname(File.dirname(__FILE__))}/atomic_red_team/used_guids.txt" oks = [] fails = [] ATOMIC_RED_TEAM.atomic_test_paths.each do |path| begin - print "Validating #{path}..." + print "Generating guids #{path}..." YAML.load_file(path) - AtomicRedTeam.new.generate_guids_for_yaml! path + AtomicRedTeam.new.generate_guids_for_yaml!(path, USED_GUIDS_FILE) oks << path puts "OK"