diff --git a/atomic_red_team/atomic_red_team.rb b/atomic_red_team/atomic_red_team.rb index c0ae26c1..e282a021 100755 --- a/atomic_red_team/atomic_red_team.rb +++ b/atomic_red_team/atomic_red_team.rb @@ -62,7 +62,7 @@ class AtomicRedTeam end end - def validate_atomic_yaml!(yaml, used_guids_file) + def validate_atomic_yaml!(yaml, used_guids_file, unique_guid_array) raise("YAML file has no elements") if yaml.nil? raise('`attack_technique` element is required') unless yaml.has_key?('attack_technique') @@ -81,8 +81,10 @@ class AtomicRedTeam if atomic.has_key?('auto_generated_guid') guid = atomic["auto_generated_guid"].to_s - raise("`atomic_tests[#{i}].auto_generated_guid` element must be unique") unless is_unique_guid(guid, used_guids_file) raise("`atomic_tests[#{i}].auto_generated_guid` element not a proper guid") unless /[a-f0-9]{8}-[a-f0-9]{4}-4[a-f0-9]{3}-[89aAbB][a-f0-9]{3}-[a-f0-9]{12}/.match(guid) + raise("`atomic_tests[#{i}].auto_generated_guid` element must be unique") unless !unique_guid_array.include?(guid) + unique_guid_array << guid + add_guid_to_used_guid_file(guid, used_guids_file) end raise("`atomic_tests[#{i}].description` element is required") unless atomic.has_key?('description') @@ -150,30 +152,46 @@ class AtomicRedTeam end end + def record_used_guids!(yaml, used_guids_file) + return unless !yaml.nil? + + yaml['atomic_tests'].each_with_index do |atomic, i| + next unless atomic.has_key?('auto_generated_guid') + guid = atomic["auto_generated_guid"].to_s + add_guid_to_used_guid_file(guid, used_guids_file) + end + end + def generate_guids_for_yaml!(path, used_guids_file) text = File.read(path) - guid = get_unique_guid(used_guids_file) # add the "auto_generated_guid:" element after the "- name:" element if it isn't already there text.gsub!(/(?i)(^([ \t]*-[ \t]*)name:.*$(?!\s*auto_generated_guid))/) { |m| "#{$1}\n#{$2.gsub(/-/," ")}auto_generated_guid:"} # fill the "auto_generated_guid:" element in if it doesn't contain a guid - text.gsub!(/(?i)^([ \t]*auto_generated_guid:)(?!([ \t]*[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)}"} + text.gsub!(/(?i)^([ \t]*auto_generated_guid:)(?!([ \t]*[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 << text } end - def get_unique_guid(used_guids_file) + # generates a unique guid and records the guid as having been used by writing it to the used_guids_file + 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| - f.puts new_guid unless new_guid == '' - } + add_guid_to_used_guid_file(new_guid, used_guids_file) return new_guid end + # add guid to used guid file if it is the proper format and is not already in the file. raises an exception if guid isn't valid + def add_guid_to_used_guid_file(guid, used_guids_file) + open(used_guids_file, 'a') { |f| + raise("`atomic_tests[#{i}].executor.command` element is required") unless /[a-f0-9]{8}-[a-f0-9]{4}-4[a-f0-9]{3}-[89aAbB][a-f0-9]{3}-[a-f0-9]{12}/ =~ guid + f.puts guid unless !is_unique_guid(guid, used_guids_file) + } + end + def is_unique_guid(guid, used_guids_file) return !File.foreach(used_guids_file).grep(/#{guid}/).any? end diff --git a/bin/generate-guids.rb b/bin/generate-guids.rb index 74b76571..bceb7c44 100755 --- a/bin/generate-guids.rb +++ b/bin/generate-guids.rb @@ -12,7 +12,8 @@ fails = [] ATOMIC_RED_TEAM.atomic_test_paths.each do |path| begin print "Generating guids #{path}..." - YAML.load_file(path) + + ATOMIC_RED_TEAM.record_used_guids!(YAML.load_file(path), USED_GUIDS_FILE) AtomicRedTeam.new.generate_guids_for_yaml!(path, USED_GUIDS_FILE) oks << path diff --git a/bin/validate-atomics.rb b/bin/validate-atomics.rb index 1f229087..4c1b78f7 100755 --- a/bin/validate-atomics.rb +++ b/bin/validate-atomics.rb @@ -9,12 +9,12 @@ USED_GUIDS_FILE = "#{File.dirname(File.dirname(__FILE__))}/atomics/used_guids.tx oks = [] fails = [] +unique_guid_array = [] ATOMIC_RED_TEAM.atomic_test_paths.each do |path| begin print "Validating #{path}..." - YAML.load_file(path) - AtomicRedTeam.new.validate_atomic_yaml!(YAML.load_file(path), USED_GUIDS_FILE) + AtomicRedTeam.new.validate_atomic_yaml!(YAML.load_file(path), USED_GUIDS_FILE, unique_guid_array) oks << path puts "OK"