lots of comments

atomic_red_team.rb

@@ -3,13 +3,15 @@ require 'yaml'
 require 'erb'
 require './attack_api'
 
-
 class AtomicRedTeam
   ATTACK_API = Attack.new
 
   # TODO- should these all be relative URLs?
   ROOT_GITHUB_URL = "https://github.com/redcanaryco/atomic-red-team"
   
+  #
+  # Returns a list of Atomic Tests in Atomic Red Team (as Hashes from source YAML) 
+  #
   def atomic_tests
     @atomic_tests ||= Dir["#{File.dirname(__FILE__)}/atomics/t*/t*.yaml"].sort.collect do |path| 
       atomic_yaml = YAML.load(File.read path)
@@ -18,9 +20,12 @@ class AtomicRedTeam
     end
   end
 
+  #
+  # Returns the individual Atomic Tests for a given identifer, passed as either a string (T1234) or an ATT&CK technique object
+  #
   def atomic_tests_for_technique(technique_or_technique_identifier)
     technique_identifier = if technique_or_technique_identifier.is_a? Hash
-      technique_or_technique_identifier.fetch('external_references', []).find {|refs| refs['source_name'] == 'mitre-attack'}['external_id'].downcase
+      ATTACK_API.technique_identifier_for_technique technique_or_technique_identifier
     else
       technique_or_technique_identifier
     end
@@ -30,8 +35,13 @@ class AtomicRedTeam
     end.to_h.fetch('atomic_tests', [])
   end
 
+  #
+  # Returns a Markdown formatted Github link to a technique. This will be to the edit page for 
+  # techniques that already have one or more Atomic Red Team tests, or the create page for
+  # techniques that have no existing tests.
+  #
   def github_link_to_technique(technique, include_identifier=false)
-    technique_identifier = technique.fetch('external_references', []).find {|refs| refs['source_name'] == 'mitre-attack'}['external_id'].downcase
+    technique_identifier = ATTACK_API.technique_identifier_for_technique(technique).downcase
     link_display = "#{"#{technique_identifier.upcase} " if include_identifier}#{technique['name']}"
 
     if File.exists? "#{File.dirname(__FILE__)}/atomics/#{technique_identifier}/#{technique_identifier}.md"

attack_api.rb

@@ -2,7 +2,14 @@
 require 'open-uri'
 require 'json'
 
+#
+# Attack is an API class that loads information about ATT&CK techniques from MITRE'S ATT&CK 
+# STIX representation. It makes it very simple to do common things with ATT&CK.
+#
 class Attack
+  # 
+  # Tactics as presented in the order that the ATT&CK matrics uses
+  #
   def ordered_tactics 
     [
       'initial-access',
@@ -19,6 +26,18 @@ class Attack
     ]
   end
 
+  # 
+  # Returns the technique identifier (T1234) for a Technique object
+  #
+  def technique_identifier_for_technique(technique)
+    technique.fetch('external_references', []).find do |refs| 
+      refs['source_name'] == 'mitre-attack'
+    end['external_id'].upcase
+  end
+
+  # 
+  # Returns a Technique object given a technique identifier (T1234)
+  #
   def technique_info(technique_id)
     techniques.find do |item| 
       item.fetch('external_references', []).find do |references|
@@ -27,16 +46,9 @@ class Attack
     end
   end
   
-  def techniques_by_tactic
-    techniques_by_tactic = Hash.new {|h, k| h[k] = []}
-    techniques.each do |technique|
-      technique.fetch('kill_chain_phases', []).select {|phase| phase['kill_chain_name'] == 'mitre-attack'}.each do |tactic|
-        techniques_by_tactic[tactic.fetch('phase_name')] << technique
-      end
-    end
-    techniques_by_tactic
-  end 
-
+  # 
+  # Returns the ATT&CK Matrix as a 2D array, in order by `ordered_tactics`
+  #
   def ordered_tactic_to_technique_matrix
     # make an 2d array of our techniques in the order our tactics appear
     all_techniques_in_tactic_order = []
@@ -54,17 +66,38 @@ class Attack
     all_techniques_in_tactic_order.transpose
   end
 
+  # 
+  # Returns a map of all [ ATT&CK Tactic name ] => [ List of ATT&CK techniques associated with that tactic]
+  #
+  def techniques_by_tactic
+    techniques_by_tactic = Hash.new {|h, k| h[k] = []}
+    techniques.each do |technique|
+      technique.fetch('kill_chain_phases', []).select {|phase| phase['kill_chain_name'] == 'mitre-attack'}.each do |tactic|
+        techniques_by_tactic[tactic.fetch('phase_name')] << technique
+      end
+    end
+    techniques_by_tactic
+  end 
+
+  #
+  # Returns a list of all ATT&CK techniques
+  #
   def techniques
     # pull out the attack pattern objects
-    attack_json.fetch("objects").select do |item| 
+    attack_stix.fetch("objects").select do |item| 
       item.fetch('type') == 'attack-pattern' && item.fetch('external_references', []).select do |references|
         references['source_name'] == 'mitre-attack'
       end
     end
   end
 
-  def attack_json
-    @attack_json ||= begin
+  private
+
+  #
+  # Returns the complete ATT&CK STIX collection parsed into a Hash
+  #
+  def attack_stix
+    @attack_stix ||= begin
       # load the full attack library
       local_attack_json_to_try = "#{File.dirname(__FILE__)}/enterprise-attack.json"
       if File.exists? local_attack_json_to_try

generate_atomic_docs.rb

@@ -1,5 +1,4 @@
 #! /usr/bin/env ruby
-require 'yaml'
 require 'erb'
 require './attack_api'
 require './atomic_red_team'
@@ -8,6 +7,9 @@ class AtomicRedTeamDocs
   ATTACK_API = Attack.new
   ATOMIC_RED_TEAM = AtomicRedTeam.new
 
+  #
+  # Generates all the documentation used by Atomic Red Team
+  #
   def generate_all_the_docs!
     oks = []
     fails = []
@@ -15,7 +17,7 @@ class AtomicRedTeamDocs
     ATOMIC_RED_TEAM.atomic_tests.each do |atomic_yaml|
       begin
         print "Generating docs for #{atomic_yaml['atomic_yaml_path']}"
-        generate_docs! atomic_yaml, atomic_yaml['atomic_yaml_path'].gsub(/.yaml/, '.md')
+        generate_technique_docs! atomic_yaml, atomic_yaml['atomic_yaml_path'].gsub(/.yaml/, '.md')
         
         oks << atomic_yaml['atomic_yaml_path']
         puts "OK"
@@ -25,13 +27,16 @@ class AtomicRedTeamDocs
       end
     end
 
-    generate_attack_matrix!
-    generate_index!
+    generate_attack_matrix! "#{File.dirname(__FILE__)}/atomics/matrix.md"
+    generate_index! "#{File.dirname(__FILE__)}/atomics/index.md"
     
     return oks, fails
   end
 
-  def generate_docs!(atomic_yaml, output_doc_path)
+  #
+  # Generates Markdown documentation for a specific technique from its YAML source
+  #
+  def generate_technique_docs!(atomic_yaml, output_doc_path)
     technique = ATTACK_API.technique_info(atomic_yaml.fetch('attack_technique'))
     technique['identifier'] = atomic_yaml.fetch('attack_technique').upcase
 
@@ -41,8 +46,11 @@ class AtomicRedTeamDocs
     print " => #{output_doc_path} => "
     File.write output_doc_path, generated_doc
   end
-
-  def generate_attack_matrix!
+  
+  #
+  # Generates a Markdown ATT&CK documentation matrix for all techniques
+  #
+  def generate_attack_matrix!(output_doc_path)
     result = "| #{ATTACK_API.ordered_tactics.join(' | ')} |\n"
     result += "|#{'-----|' * ATTACK_API.ordered_tactics.count}\n"
 
@@ -54,10 +62,13 @@ class AtomicRedTeamDocs
       end
       result += "| #{row_values.join(' | ')} |\n"
     end
-    File.write "#{File.dirname(__FILE__)}/atomics/matrix.md", result
+    File.write output_doc_path, result
   end
 
-  def generate_index!
+  #
+  # Generates a master Markdown index of ATT&CK Tactic -> Technique -> Atomic Tests
+  #
+  def generate_index!(output_doc_path)
     result = ''
 
     ATTACK_API.techniques_by_tactic.each do |tactic, techniques|
@@ -71,10 +82,13 @@ class AtomicRedTeamDocs
       result += "\n"
     end
 
-    File.write "#{File.dirname(__FILE__)}/atomics/index.md", result
+    File.write output_doc_path, result
   end
 end
 
+#
+# MAIN
+#
 oks, fails = AtomicRedTeamDocs.new.generate_all_the_docs!
 puts
 puts "Generated docs for #{oks.count} techniques, #{fails.count} failures"