#!/usr/bin/env ruby
# Copyright 2015 Adaptavist.com Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

require 'fileutils'
require 'yaml'
require 'set'
require 'pathname'
require 'deep_merge'
require 'docopt'
require 'colorize'
require 'facter'

doc = <<DOCOPT
Adaptavist puppet runner

Usage:
  puppet-runner (prepare|all) [-c CONFIG_DIR] [-t TEMPLATES] [-d DESTINATION_DIR] [-f FACTS_DEST] [-s SERVERNAME] [-p PUPPET_APPLY] [-r PUPPETFILE_CONFIG] [-o PUPPETFILE_OUTPUT_PATH] [-e EYAML_KEY_PATH]
  puppet-runner start [-p PUPPET_APPLY]
  puppet-runner -h | --help

Options:
  -h --help                                                                 Show this screen.
  -s SERVERNAME --servername SERVERNAME                                     Custom identification of server, hostname fact if not provided
  -c CONFIG_DIR --config_dir CONFIG_DIR                                     Hiera configuration directory, must ontain <hostname>.yaml and <hostname>_facts.yaml
  -d DESTINATION_DIR --dest_dir DESTINATION_DIR                             Directory for result hiera config.
  -t TEMPLATES --templates TEMPLATES                                        Directory containing templates and defaults folder with functionality templates and default facts
  -f FACTS_DEST --facts_dest_dir FACTS_DEST                                 Destination directory to store result facts
  -p PUPPET_APPLY --puppet_apply PUPPET_APPLY                               Custom puppet apply command to run
  -r PUPPETFILE_CONFIG --puppetfile_config puppetfile_config                Puppetfile composition config file
  -o PUPPETFILE_OUTPUT_PATH --puppetfile_output_path PUPPETFILE_OUTPUT_PATH Result Puppetfile path
  -e EYAML_KEY_PATH --eyaml_key_path EYAML_KEY_PATH                         Path to eyaml encryption key pair
Commands:
  all           Runs the following commands prepare, start 
  start         Runs puppet apply 
  prepare       Creates result hiera config as a composition of functionalities based on config, merges provided facts with defaults
DOCOPT


def path_join_glob(dir, file_pattern)
    "#{dir}/#{file_pattern}"
end

def debug(msg)
    puts msg.green
end

def warning(msg)
    puts msg.yellow
end

def extract_value_from_hash(input)
  res = {}
  if input
    res = input.map{|key, val| 
      if val.is_a?(Hash)
        value = val["value"]
      else
        value = val
      end
      {key => value}
    }
  end
  res
end

def extract_comment_from_hash(input)
  res = {}
  if input
    res = input.map{|key, val|

      if val != nil
        if val.is_a?(Hash)
          value = val["comment"]
        else
          value = nil
        end
      else
        value = nil
      end
      {key => value}
    }
  end
  res
end

begin
    options = Docopt::docopt(doc)
rescue Docopt::Exit => e
    abort(e.message.red)
end

stop_apply = false

if options['all'] || options['prepare']
    input_dir = options["--config_dir"] || options["-c"] 
    dest_dir = options["--dest_dir"] || options["-d"]
    facts_dest_dir = options["--facts_dest_dir"] || options["-f"]
    templates = options["--templates"] || options["-t"]
    puppetfile_config_path = options["--puppetfile_config"] || options["-r"]
    puppetfile_output_path = options["--puppetfile_output_path"] || options["-o"]
    eyaml_key_path = options["--eyaml_key_path"] || options["-e"] || "/etc/puppet/config"
    hostname = options["--servername"] || options["-s"] || Facter.value("hostname")
    puts "Hostname #{hostname}"

    config_file_path = path_join_glob(input_dir, hostname+".yaml")
    templates_dir = path_join_glob(templates, "templates")
    def_facts_dir = path_join_glob(templates, "defaults")

    debug "Reading #{config_file_path}"
    if File.file? config_file_path and File.directory? templates_dir and File.directory? def_facts_dir and File.directory? dest_dir and File.directory? facts_dest_dir and File.file? puppetfile_config_path
        config = YAML.load_file(config_file_path)
    else
        abort "Can not find config file #{config_file_path}. \
         or #{templates_dir}. \
         or #{def_facts_dir}. \
         or #{dest_dir}. \
         or #{facts_dest_dir}.".red
    end

    functionalities = config["functionalities"]

    output_file_path = path_join_glob(dest_dir, "#{hostname}.eyaml")
    output_encrypted_facts_file_path = "/tmp/#{hostname}_facts.eyaml"
    output_facts_file_path = path_join_glob(facts_dest_dir, "#{hostname}_facts.yaml")
    if File.file? output_file_path
      FileUtils.rm output_file_path
    end
    if File.file? output_encrypted_facts_file_path
      FileUtils.rm output_encrypted_facts_file_path 
    end
    if File.file? output_facts_file_path
      FileUtils.rm output_facts_file_path
    end

    debug "Writing to #{output_file_path}"
    result_template = {}
    result_default_facts = {}
    prefixed_required_facts = Set.new
    prefixed_facts_comments = {}
    puppetfile_config = YAML.load_file(puppetfile_config_path) || {}
    puppetfile_dependencies = []
      # functionalities:
      # # In honor of Henry... 
      #   1_app:
      #     - confluence: "conf1"
      #     - confluence: "conf2"
      #     - jira
      #   2_database:
      #     - mysql

      functionalities.keys.sort.each do |key|
        next unless functionalities[key]
          
        functionalities[key].each do |to_add|
          
          if to_add.is_a?(Hash)
            template_to_add = path_join_glob(templates_dir, "#{to_add.keys[0]}.yaml")
            facts_to_add = path_join_glob(def_facts_dir, "#{to_add.keys[0]}.yaml")
          else
            template_to_add = path_join_glob(templates_dir, "#{to_add}.yaml")
            facts_to_add = path_join_glob(def_facts_dir, "#{to_add}.yaml")
          end

          debug "Adding template #{template_to_add}"
          debug "Adding facts #{facts_to_add}"
          if File.file? template_to_add and File.file? facts_to_add
            # prefix is defined, must replace
            data = YAML.load_file(template_to_add) || {}
            default_facts = YAML.load_file(facts_to_add) || {}
            prefixes = data["prefixes"] || []
            required_facts = data["required_facts"] || []
            puppetfile_parts = data["dependencies"] || []
            # merge dependencies 
            puppetfile_dependencies = puppetfile_dependencies | puppetfile_parts

            data_as_string = data.to_s
            facts_as_string = extract_value_from_hash(default_facts).to_s

            fact_comments_as_string = extract_comment_from_hash(default_facts).to_s
            
            if to_add.is_a?(Hash)
              # if prefixes are not defined skip replacing
              if prefixes
          
                # in case of hash, replace each otherwise replace all with prefix
                if to_add.values[0].is_a?(Hash)
                  to_add.values[0].keys.each do |prefix_key|
                    
                    prefixes.each do |prefix|
                      if prefix == prefix_key
                        replace_prefixes_with = to_add.values[0][prefix_key]
                        debug "will substiture: #{prefix} with #{replace_prefixes_with}"
                        data_as_string = data_as_string.gsub(/\%{::#{prefix}/, "\%{::#{replace_prefixes_with}")
                        facts_as_string = facts_as_string.gsub(/#{prefix}/, "#{replace_prefixes_with}")
                        fact_comments_as_string = fact_comments_as_string.gsub(/#{prefix}/, "#{replace_prefixes_with}")
                        prefixed_required_facts = prefixed_required_facts.merge(required_facts.map! { |item| item.gsub(/#{prefix}/, "#{replace_prefixes_with}") })
                      end
                    end
                  end

                else 
                  replace_prefixes_with = to_add.values[0]
                  prefixes.each do |prefix|
                    debug "will substitute: #{prefix} with #{replace_prefixes_with}"
                    data_as_string = data_as_string.gsub(/\%{::#{prefix}/, "\%{::#{replace_prefixes_with}")
                    facts_as_string = facts_as_string.gsub(/#{prefix}/, "#{replace_prefixes_with}")
                    fact_comments_as_string = fact_comments_as_string.gsub(/#{prefix}/, "#{replace_prefixes_with}")
                    prefixed_required_facts = prefixed_required_facts.merge(required_facts.map! { |item| item.gsub(/#{prefix}/, "#{replace_prefixes_with}") })
                  end
                end
              end

              template = eval (data_as_string)
              default_facts_prefixed = eval (facts_as_string)
              default_fact_comments = eval (fact_comments_as_string)
            else
              template = YAML.load_file(template_to_add)
              plain_facts = YAML.load_file(facts_to_add)
              default_facts_prefixed = extract_value_from_hash(plain_facts)
              default_fact_comments = extract_comment_from_hash(plain_facts)  
              prefixed_required_facts = prefixed_required_facts.merge(required_facts)
            end
            result_template.deep_merge!(template)
            # default_facts_prefixed is Array of hashes as the result of map, this will create hash from it
            result_default_facts.merge!(default_facts_prefixed.reduce Hash.new, :merge)
            prefixed_facts_comments.merge!(default_fact_comments.reduce Hash.new, :merge)
          else
            abort "Can not find template in templates folder #{template_to_add} or #{facts_to_add}".red
          end
        end
      end
      # Write results
      File.open(output_file_path, 'w+') do |output_file|
        YAML.dump(result_template, output_file)
      end
      custom_facts_path = path_join_glob(input_dir, "#{hostname}_facts.yaml")
      custom_facts = YAML.load_file(custom_facts_path) || {}
      File.open(output_encrypted_facts_file_path, 'w+') do |output_file|
        output_result_default_facts = result_default_facts.deep_merge!(custom_facts, {:merge_hash_arrays => true}).to_yaml
        prefixed_facts_comments.each do |pattern, replacement|
          if replacement != nil
            output_result_default_facts.gsub!(/^#{pattern}/, "\##{replacement}\n#{pattern}")
          end
        end
        output_file.write(output_result_default_facts)
      end
      
      # decrypt facts file because Puppet doesn't appear to be able to read encrypted facts
      require 'hiera/backend/eyaml/plugins'
      require 'hiera/backend/eyaml/encryptors/pkcs7'
      require 'hiera/backend/eyaml/subcommands/decrypt'
      require 'hiera/backend/eyaml/options'
      Hiera::Backend::Eyaml::Encryptors::Pkcs7.register
      options = {
          :eyaml=>output_encrypted_facts_file_path, 
          :pkcs7_public_key =>"#{eyaml_key_path}/public_key.pkcs7.pem",
          :pkcs7_private_key=>"#{eyaml_key_path}/private_key.pkcs7.pem" }
      Hiera::Backend::Eyaml::Options.set(Hiera::Backend::Eyaml::Subcommands::Decrypt.validate options)

      # manually ensure multi-line encrypted values are output correctly into the new decrypted yaml file
      # this is just modifed source from http://www.rubydoc.info/gems/hiera-eyaml/2.0.8/Hiera%2FBackend%2FEyaml%2FSubcommands%2FDecrypt.execute
      File.open(output_facts_file_path, 'w') do |output_file|
        parser = Hiera::Backend::Eyaml::Parser::ParserFactory.encrypted_parser
        tokens = parser.parse(Hiera::Backend::Eyaml::Options[:input_data])
        case Hiera::Backend::Eyaml::Options[:source]
          when :eyaml
            decrypted = tokens.map{ |token|
              decrypted_value = token.to_decrypted
              encryption_indicator = 'DEC::PKCS7['
              four_spaces = '    '
              multiline_value = if decrypted_value.include? encryption_indicator and decrypted_value.include? "\n" then
                                  "|\n#{four_spaces}" + decrypted_value.gsub("\n", "\n#{four_spaces}")
                                else
                                  decrypted_value
                                end
              multiline_value.gsub(encryption_indicator, '').gsub(']!', '')
            }
          else
            decrypted = tokens.map{ |token|
              case token.class.name
                when /::EncToken$/
                  token.plain_text
                else
                  token.match
              end
            }
        end
        output_file.write(decrypted.join)
      end
      
      File.open(puppetfile_output_path, 'w+') do |output_file|
        header = "#!/usr/bin/env ruby\n\n"
        output_file.write(header)
        puppetfile_dependencies.each do |pup|
          dep = puppetfile_config[pup]
          if dep
            res = "mod \"#{dep['name']}\", \n" +
             "  :#{dep['repo_type']} => '#{dep['repo']}',\n" +
             "  :#{dep['ref_type']} => '#{dep['ref_value']}'\n\n"
             output_file.write(res)
          else 
            warning "Can not find configuration for module #{pup} in config of puppet modules!"
          end
        end
      end

      nil_value_facts = result_default_facts.select{|key, val| val == nil}

      # Check that all required prefixed facts are present
        if (prefixed_required_facts && !prefixed_required_facts.empty?) or !nil_value_facts.empty?
        not_provided_required_facts = prefixed_required_facts - custom_facts.keys
        
        if !not_provided_required_facts.empty? or !nil_value_facts.empty?
          not_provided_required_facts.merge(nil_value_facts.keys)
          warning "You have to provide all required fields, they will default to empty string and puppet will fail: "
          not_provided_required_facts.each do |f|
            warning "#{f}"
          end
          warning "Puppet apply will not run as it will fail without those facts provided!"
          stop_apply = true
        end
      end      
end


# start puppet 
if (options['start'] || options['all']) && !stop_apply
    require 'puppet'
    modulefile_definition = Gem::Version.new(Puppet.version) > Gem::Version.new('4.0.0') ? '--modulepath /etc/puppet/modules' : ''
    puppet_command = "sudo su -c 'source /usr/local/rvm/scripts/rvm; puppet apply /etc/puppet/manifests/site.pp --confdir=/etc/puppet  --verbose --detailed-exitcodes #{modulefile_definition}'"
    to_execute = options["--puppet_apply"] || options["-p"] || puppet_command
    debug "Running #{to_execute}"
    `#{to_execute}`
    exit_code = $?.exitstatus
    if exit_code != 2
      raise "execute_puppet exit status: #{exit_code}"
    end
end