#!/opt/puppetlabs/puppet/bin/ruby
# encoding: UTF-8

require 'json'
require 'yaml'
require 'puppet'

module Errors
  InvalidJson = "invalid_json"
  NoPuppetBin = "no_puppet_bin"
  NoLastRunReport = "no_last_run_report"
  InvalidLastRunReport = "invalid_last_run_report"
  AlreadyRunning = "agent_already_running"
  Disabled = "agent_disabled"
  FailedToStart = "agent_failed_to_start"
  NonZeroExit = "agent_exit_non_zero"
end

class ProcessingError < StandardError
  attr_reader :error_type

  def initialize(error_type, message = nil)
    super(message)
    @error_type = error_type
  end
end

DEFAULT_EXITCODE = -1

env_fix_up = nil

# use this way of defining the `get_env_fix_up` method so that it can "see"
# the `env_fix_up` variable
self.class.send(:define_method, :get_env_fix_up) do
  # If running in a C or POSIX locale, ask Puppet to use UTF-8
  base_env = {}
  if Encoding.default_external == Encoding::US_ASCII
    base_env = {"RUBYOPT" => "#{ENV['RUBYOPT']} -EUTF-8"}
  end

  # Prepare an environment fix-up to make up for its cleansing performed
  # by the Puppet::Util::Execution.execute function.
  # This fix-up is meant for running puppet under a non-root user;
  # puppet cannot find the user's HOME directory otherwise.
  env_fix_up ||= if Puppet.features.microsoft_windows? || Process.euid == 0
    # no environment fix-up is needed on windows or for root
    base_env
  else
    begin
      require 'etc'

      pwentry = Etc.getpwuid(Process.euid)

      {"USER"    => pwentry.name,
       "LOGNAME" => pwentry.name,
       "HOME"    => pwentry.dir}.merge base_env
    rescue => e
      # oh well ..., let's give it a try without the environment fix-up
      myname = File.basename($0)
      $stderr.puts "#{myname}: Could not fix environment for effective UID #{Process.euid}: #{e.message}"
      $stderr.puts "#{myname}: Expect puppet run problems"
      base_env
    end
  end
end

def last_run_result(exitcode)
  return {"kind"             => "unknown",
          "time"             => "unknown",
          "transaction_uuid" => "unknown",
          "environment"      => "unknown",
          "status"           => "unknown",
          "exitcode"         => exitcode,
          "version"          => 1}
end

def force_unicode(s)
  begin
    # Later comparisons assume UTF-8. Convert to that encoding now.
    s.encode(Encoding::UTF_8)
  rescue Encoding::InvalidByteSequenceError
    # Found non-native characters, hope it's a UTF-8 string. Since this is Puppet, and
    # incorrect characters probably means we're in a C or POSIX locale, this is usually safe.
    s.force_encoding(Encoding::UTF_8)
  end
end

def check_config_print(cli_arg, config)
  command_array = [config["puppet_bin"], "agent", "--configprint", cli_arg]
  process_output = Puppet::Util::Execution.execute(command_array,
                                                   {:custom_environment => get_env_fix_up(),
                                                    :override_locale => false})
  return force_unicode(process_output.to_s).chomp
end

def running?(run_result)
  !!(run_result =~ /Run of Puppet configuration client already in progress/)
end

def disabled?(run_result)
  !!(run_result =~ /disabled(.*?)Use 'puppet agent --enable' to re-enable/m)
end

def make_environment_hash()
  # NB: we're ignoring the `env` array for setting the environment
  return get_env_fix_up()
end

def make_command_array(config, action_input)
  cmd_array = [config["puppet_bin"], "agent"]
  cmd_array +=  action_input["flags"]
  return cmd_array
end

def make_error_result(exitcode, error_type, error_message)
  result = last_run_result(exitcode)
  result["error_type"] = error_type
  result["error"] = error_message
  return result
end

def parse_report(filename)
  # Read the report and drop Ruby objects first. We can't parse the Ruby objects
  # (because we don't have Puppet loaded) and don't need that data.
  # Psych::Nodes::Node#each iterates over each node in the parsed document tree.
  # YAML.parse_file returns a Psych::Nodes::Document, and #root returns the
  # root-level Node.
  data = YAML.parse_file(filename)
  data.root.each do |o|
    o.tag = nil if o.respond_to?(:tag=)
  end

  data.to_ruby
end

def get_result_from_report(last_run_report, exitcode, config, start_time)
  if !File.exist?(last_run_report)
    return make_error_result(exitcode, Errors::NoLastRunReport,
                             "#{last_run_report} doesn't exist")
  end

  if start_time && File.mtime(last_run_report) == start_time
    return make_error_result(exitcode, Errors::NoLastRunReport,
                             "#{last_run_report} was not written")
  end

  last_run_report_yaml = {}

  begin
    last_run_report_yaml = parse_report(last_run_report)
  rescue => e
    return make_error_result(exitcode, Errors::InvalidLastRunReport,
                             "#{last_run_report} could not be loaded: #{e}")
  end

  if exitcode == 0
    run_result = last_run_result(exitcode)
  else
    run_result = make_error_result(exitcode, Errors::NonZeroExit,
                                   "Puppet agent exited with a non 0 exitcode")
  end

  run_result["kind"] = last_run_report_yaml['kind']
  run_result["time"] = last_run_report_yaml['time']
  run_result["transaction_uuid"] = last_run_report_yaml['transaction_uuid']
  run_result["environment"] = last_run_report_yaml['environment']
  run_result["status"] = last_run_report_yaml['status']

  return run_result
end

# Wait for the lockfile to be removed. If it hasn't after 30 seconds, give up.
def wait_for_lockfile(lockfile = '', check_interval = 0.1, give_up_after = 30)
  number_of_tries = give_up_after / check_interval
  count = 0
  while File.exist?(lockfile) && count < number_of_tries
    sleep check_interval
    count += 1
  end
end

def start_run(config, action_input)
  exitcode = DEFAULT_EXITCODE
  cmd = make_command_array(config, action_input)
  env = make_environment_hash()

  last_run_report = check_config_print("lastrunreport", config)
  if last_run_report.empty?
    return make_error_result(exitcode, Errors::NoLastRunReport,
                             "could not find the location of the last run report")
  end

  # Initially ignore the lockfile. It might be out-dated, so we give Puppet a chance
  # to clean it up and run.
  lockfile = ''
  while true
    wait_for_lockfile(lockfile)

    start_time = File.mtime(last_run_report) if File.exist?(last_run_report)

    run_result = Puppet::Util::Execution.execute(cmd, {:failonfail => false,
                                                       :custom_environment => env,
                                                       :override_locale => false})

    if !run_result
      return make_error_result(exitcode, Errors::FailedToStart,
                               "Failed to start Puppet agent")
    end

    exitcode = run_result.exitstatus
    puppet_output = force_unicode(run_result.to_s)

    if disabled?(puppet_output)
      return make_error_result(exitcode, Errors::Disabled,
                               "Puppet agent is disabled")
    elsif running?(puppet_output)
      # Puppet is already running, attempt to wait until it finishes and retry.
      lockfile = check_config_print('agent_catalog_run_lockfile', config)
      if lockfile.empty?
        # Can't identify lockfile, just report an error.
        return make_error_result(exitcode, Errors::AlreadyRunning,
                                 "Puppet agent is already performing a run; it also appears to be" +
                                 " misconfigured, agent_catalog_run_lockfile is an empty string")
      end
      next
    end

    return get_result_from_report(last_run_report, exitcode, config, start_time)
  end
end

# TODO(ale): remove `env` from input before bumping to the next version
def metadata()
  return {
    :description => "PXP Puppet module",
    :actions => [
      { :name        => "run",
        :description => "Start a Puppet run",
        :input       => {
          :type      => "object",
          :properties => {
            :env => {
              :type => "array",
            },
            :flags => {
              :type => "array",
              :items => {
                :type => "string"
              }
            }
          },
          :required => [:flags]
        },
        :results => {
          :type => "object",
          :properties => {
            :kind => {
              :type => "string"
            },
            :time => {
              :type => "string"
            },
            :transaction_uuid => {
              :type => "string"
            },
            :environment => {
              :type => "string"
            },
            :status => {
              :type => "string"
            },
            :error_type => {
              :type => "string"
            },
            :error => {
              :type => "string"
            },
            :exitcode => {
              :type => "number"
            },
            :version => {
              :type => "number"
            }
          },
          :required => [:kind, :time, :transaction_uuid, :environment, :status,
                        :exitcode, :version]
        }
      }
    ],
    :configuration => {
      :type => "object",
      :properties => {
        :puppet_bin => {
          :type => "string"
        }
      }
    }
  }
end

def run(config, action_input)
  if !File.exist?(config["puppet_bin"])
    return make_error_result(DEFAULT_EXITCODE, Errors::NoPuppetBin,
                             "Puppet executable '#{config["puppet_bin"]}' does not exist")
  end

  return start_run(config, action_input)
end

def get_configuration(args)
  config = args["configuration"] || {}

  if config["puppet_bin"].nil? || config["puppet_bin"].empty?
    if !Puppet.features.microsoft_windows?
      config["puppet_bin"] = "/opt/puppetlabs/bin/puppet"
    else
      module_path = File.expand_path(File.dirname(__FILE__))
      puppet_bin = File.join(module_path, '..', '..', 'bin', 'puppet.bat')
      config["puppet_bin"] = File.expand_path(puppet_bin)
    end
  end

  config
end

DEFAULT_FLAGS = ["--onetime", "--no-daemonize", "--verbose"]

DEFAULT_FLAGS_NAMES = ["onetime", "daemonize", "verbose"]

WHITELISTED_FLAGS_NAMES = [
  "color", "configtimeout",
  "debug","disable_warnings",
  "environment", "evaltrace",
  "filetimeout",
  "graph",
  "http_connect_timeout", "http_debug", "http_keepalive_timeout", "http_read_timeout",
  "log_level",
  "noop",
  "ordering",
  "pluginsync",
  "show_diff", "skip_tags", "splay", "strict_environment_mode",
  "tags", "trace",
  "use_cached_catalog", "usecacheonfailure",
  "waitforcert"]

# All flags and valid arguments to them should be caught by this.
#  It was constructed for the following argument types:
#  "timeout": /\A\d+[mhdy]?\Z/,
#  "environment": /\A[a-z0-9_]+\Z/,
#  "tag": /\A[a-z0-9_][a-z0-9_:\.\-]*\Z/,
#  "ordering": /\Atitle-hash|manifest|random\Z/
VALID_FLAG_REGEX = /\A[a-zA-Z0-9_:,\.\-]+\Z/

# This asserts that the flag has a valid prefix
def get_flag_name(flag)
  if flag.start_with?("--no-")
    flag[5..-1]
  elsif flag.start_with?("--")
    flag[2..-1]
  else
    raise "Assertion error: we're here by mistake"
  end
end

def get_action_input(args)
  action_input = args["input"]
  unless action_input.is_a?(Hash)
    raise ProcessingError.new(Errors::InvalidJson, "The json received on STDIN did not contain a valid 'input' key: #{args.to_s}")
  end

  if action_input.has_key?("flags")
    action_input["flags"].each do |flag|
      flag = flag.strip
      unless flag =~ VALID_FLAG_REGEX
        raise ProcessingError.new(Errors::InvalidJson, "The json received on STDIN contained characters not present in valid flags: #{flag}")
      end

      if flag.start_with?("--")
        flag_name = get_flag_name(flag)

        if DEFAULT_FLAGS_NAMES.include?(flag_name)
          unless DEFAULT_FLAGS.include?(flag)
            raise ProcessingError.new(Errors::InvalidJson, "The json received on STDIN overrides a default setting with: #{flag}")
          end
          next
        end

        unless WHITELISTED_FLAGS_NAMES.include?(flag_name)
          raise ProcessingError.new(Errors::InvalidJson, "The json received on STDIN included a non-permitted flag: #{flag}")
        end
      end
    end
  end

  action_input["flags"] = action_input["flags"] | DEFAULT_FLAGS
  action_input
end

def action_run(input)
  begin
    args = JSON.parse(input)
  rescue
    return make_error_result(DEFAULT_EXITCODE, Errors::InvalidJson,
                             "Invalid json received on STDIN: #{input}")
  end
  unless args.is_a?(Hash)
    return make_error_result(DEFAULT_EXITCODE, Errors::InvalidJson,
                             "The json received on STDIN was not a hash: #{args.to_s}")
  end

  output_files = args["output_files"]
  if output_files
    begin
      $stdout.reopen(File.open(output_files["stdout"], 'w'))
      $stderr.reopen(File.open(output_files["stderr"], 'w'))
    rescue => e
      print make_error_result(DEFAULT_EXITCODE, Errors::InvalidJson,
                              "Could not open output files: #{e.message}").to_json
      exit 5 # this exit code is reserved for problems with opening
             # of the output_files
    end

    at_exit do
      status = if $!.nil?
        0
      elsif $!.is_a?(SystemExit)
        $!.status
      else
        1
      end

      # flush the stdout/stderr before writing the exitcode
      # file to avoid pxp-agent reading incomplete output
      $stdout.fsync
      $stderr.fsync
      begin
        File.open(output_files["exitcode"], 'w') do |f|
          f.puts(status)
        end
      rescue => e
        print make_error_result(DEFAULT_EXITCODE, Errors::InvalidJson,
                                "Could not open exit code file: #{e.message}").to_json
        exit 5 # this exit code is reserved for problems with opening
               # of the output_files
      end
    end
  end

  begin
    config = get_configuration(args)
    action_input = get_action_input(args)
  rescue ProcessingError => e
    return make_error_result(DEFAULT_EXITCODE, e.error_type, e.message)
  end

  run(config, action_input)
end

if __FILE__ == $0
  action = ARGV.shift || 'metadata'

  if action == 'metadata'
    puts metadata.to_json
  else
    action_results = action_run($stdin.read.chomp)
    print action_results.to_json

    unless action_results["error"].nil?
      exit 1
    end
  end
end
