#!/opt/puppetlabs/puppet/bin/ruby

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

# This method exists to facilitate testing
def last_child_exit_status
  return $?
end

# This method exists to facilitate testing
def is_win?
  return !!File::ALT_SEPARATOR
end

def last_run_result(exitcode)
  return {"kind"             => "unknown",
          "time"             => "unknown",
          "transaction_uuid" => "unknown",
          "environment"      => "unknown",
          "status"           => "unknown",
          "exitcode"         => exitcode,
          "version"          => 1}
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)
  return 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(action_input)
  env_hash = {}
  action_input["env"].each do |entry|
    key, value = entry.split('=')
    env_hash[key] = value
  end

  return env_hash
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 get_result_from_report(exitcode, config, start_time)
  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

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

  last_run_report_yaml = {}

  begin
    last_run_report_yaml = YAML.load_file(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

  if start_time.to_i < last_run_report_yaml.time.to_i
    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
  end

  return run_result
end

def start_run(config, action_input)
  exitcode = DEFAULT_EXITCODE
  cmd = make_command_array(config, action_input)
  env = make_environment_hash(action_input)
  start_time = Time.now

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

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

  exitcode = run_result.exitstatus
  puppet_output = run_result.to_s

  if disabled?(puppet_output)
    return make_error_result(exitcode, Errors::Disabled,
                             "Puppet agent is disabled")
  elsif running?(puppet_output)
    return make_error_result(exitcode, Errors::AlreadyRunning,
                             "Puppet agent is already performing a run")
  end

  return get_result_from_report(exitcode, config, start_time)
end

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",
            }
          },
          :required => [:env, :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 !is_win?
      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"]

# 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|
      if flag.is_a?(String) and 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
