#!/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