96 lines
3.0 KiB
Ruby
96 lines
3.0 KiB
Ruby
#!/usr/bin/env ruby
|
|
# -*- coding: binary -*-
|
|
|
|
# Pre-requisites:
|
|
# Run the following command to fetch the latest MITRE ATT&CK data and add it to a JSON file called mitre_attack.json:
|
|
# curl -s 'https://raw.githubusercontent.com/mitre/cti/master/enterprise-attack/enterprise-attack.json' | jq '[.objects[] | select(.type == "attack-pattern")]' > mitre_attack.json
|
|
|
|
# This script generates Ruby constants for MITRE ATT&CK techniques and sub-techniques.
|
|
# It reads a JSON file (typically mitre_attack.json) containing MITRE ATT&CK data,
|
|
# extracts and formats technique names and IDs, sorts and groups them so that base
|
|
# techniques appear before their sub-techniques, and writes the result as a Ruby module.
|
|
#
|
|
# Usage:
|
|
# ruby tools/dev/generate_mitre_attack_technique_constants.rb
|
|
|
|
require 'json'
|
|
require 'fileutils'
|
|
|
|
# Required to handle transliteration of Non-ASCII characters in technique names, example: "T1186_PROCESS_DOPPELGÄNGING"
|
|
require 'i18n'
|
|
I18n.config.available_locales = :en
|
|
|
|
class MitreAttackConstantsGenerator
|
|
def initialize(input_file, output_file)
|
|
@input_file = input_file
|
|
@output_file = output_file
|
|
end
|
|
|
|
def run
|
|
data = read_json
|
|
constants = extract_constants(data)
|
|
grouped_constants = group_constants(constants)
|
|
write_output(grouped_constants)
|
|
end
|
|
|
|
private
|
|
|
|
def read_json
|
|
JSON.parse(File.read(@input_file))
|
|
end
|
|
|
|
def extract_constants(data)
|
|
constants = []
|
|
data.each do |technique|
|
|
next unless technique['type'] == 'attack-pattern'
|
|
|
|
description = technique['name']
|
|
id_entry = technique['external_references'].select { |hash| hash['external_id'] }
|
|
id = id_entry.first['external_id']
|
|
const_id = id.gsub('.', '_')
|
|
const_description = I18n.transliterate(description).upcase.gsub(/[^A-Z0-9]+/, '_').gsub(/^_+|_+$/, '')
|
|
const = "#{const_id}_#{const_description}"
|
|
constants << const
|
|
end
|
|
constants.sort
|
|
end
|
|
|
|
def group_constants(constants)
|
|
constants
|
|
.group_by { |const| const[/T\d{4}/] }
|
|
.values
|
|
.map do |group|
|
|
group.sort_by { |const| [const[/T\d{4}_(\d{3})/, 1].to_i] }
|
|
end
|
|
end
|
|
|
|
def write_output(grouped_constants)
|
|
output = []
|
|
output << '# frozen_string_literal: true'
|
|
output << ''
|
|
output << 'module Msf'
|
|
output << ' module Mitre'
|
|
output << ' module Attack'
|
|
output << " # This file was auto-generated by #{__FILE__} please do not manually edit it"
|
|
output << ' module Technique'
|
|
grouped_constants.each_with_index do |group, idx|
|
|
group.each do |const|
|
|
output << " #{const} = '#{const.match(/T\d{4}(?:_\d{3})?/).to_s.gsub('_', ".")}'"
|
|
end
|
|
output << '' unless idx == grouped_constants.size - 1
|
|
end
|
|
output << ' end'
|
|
output << ' end'
|
|
output << ' end'
|
|
output << 'end'
|
|
output << ''
|
|
|
|
FileUtils.mkdir_p(File.dirname(@output_file))
|
|
File.write(@output_file, output.join("\n"))
|
|
end
|
|
end
|
|
|
|
# Example usage:
|
|
generator = MitreAttackConstantsGenerator.new('mitre_attack.json', 'lib/msf/core/mitre/attack/technique.rb')
|
|
generator.run
|