1
#!/usr/bin/env ruby
2
# -*- coding: binary -*-
3

4
# Pre-requisites:
5
# Run the following command to fetch the latest MITRE ATT&CK data and add it to a JSON file called mitre_attack.json:
6
#   curl -s 'https://raw.githubusercontent.com/mitre/cti/master/enterprise-attack/enterprise-attack.json' | jq '[.objects[] | select(.type == "attack-pattern")]' > mitre_attack.json
7

8
# This script generates Ruby constants for MITRE ATT&CK techniques and sub-techniques.
9
# It reads a JSON file (typically mitre_attack.json) containing MITRE ATT&CK data,
10
# extracts and formats technique names and IDs, sorts and groups them so that base
11
# techniques appear before their sub-techniques, and writes the result as a Ruby module.
12
#
13
# Usage:
14
#   ruby tools/dev/generate_mitre_attack_technique_constants.rb
15

16
require 'json'
17
require 'fileutils'
18

19
# Required to handle transliteration of Non-ASCII characters in technique names, example: "T1186_PROCESS_DOPPELGÄNGING"
20
require 'i18n'
21
I18n.config.available_locales = :en
22

23
class MitreAttackConstantsGenerator
24
  def initialize(input_file, output_file)
25
    @input_file = input_file
26
    @output_file = output_file
27
  end
28

29
  def run
30
    data = read_json
31
    constants = extract_constants(data)
32
    grouped_constants = group_constants(constants)
33
    write_output(grouped_constants)
34
  end
35

36
  private
37

38
  def read_json
39
    JSON.parse(File.read(@input_file))
40
  end
41

42
  def extract_constants(data)
43
    constants = []
44
    data.each do |technique|
45
      next unless technique['type'] == 'attack-pattern'
46

47
      description = technique['name']
48
      id_entry = technique['external_references'].select { |hash| hash['external_id'] }
49
      id = id_entry.first['external_id']
50
      const_id = id.gsub('.', '_')
51
      const_description = I18n.transliterate(description).upcase.gsub(/[^A-Z0-9]+/, '_').gsub(/^_+|_+$/, '')
52
      const = "#{const_id}_#{const_description}"
53
      constants << const
54
    end
55
    constants.sort
56
  end
57

58
  def group_constants(constants)
59
    constants
60
      .group_by { |const| const[/T\d{4}/] }
61
      .values
62
      .map do |group|
63
      group.sort_by { |const| [const[/T\d{4}_(\d{3})/, 1].to_i] }
64
    end
65
  end
66

67
  def write_output(grouped_constants)
68
    output = []
69
    output << '# frozen_string_literal: true'
70
    output << ''
71
    output << 'module Msf'
72
    output << '  module Mitre'
73
    output << '    module Attack'
74
    output << "      # This file was auto-generated by #{__FILE__} please do not manually edit it"
75
    output << '      module Technique'
76
    grouped_constants.each_with_index do |group, idx|
77
      group.each do |const|
78
        output << "        #{const} = '#{const.match(/T\d{4}(?:_\d{3})?/).to_s.gsub('_', ".")}'"
79
      end
80
      output << '' unless idx == grouped_constants.size - 1
81
    end
82
    output << '      end'
83
    output << '    end'
84
    output << '  end'
85
    output << 'end'
86
    output << ''
87

88
    FileUtils.mkdir_p(File.dirname(@output_file))
89
    File.write(@output_file, output.join("\n"))
90
  end
91
end
92

93
# Example usage:
94
generator = MitreAttackConstantsGenerator.new('mitre_attack.json', 'lib/msf/core/mitre/attack/technique.rb')
95
generator.run
96

97