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

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(params)
  env_hash = {}
  params["env"].each do |entry|
    key, value = entry.split('=')
    env_hash[key] = value
  end

  return env_hash
end

def make_command_array(config, params)
  cmd_array = [config["puppet_bin"], "agent"]
  cmd_array +=  params["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 = File.join(check_config_print("statedir",config),
                              "last_run_report.yaml")

  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

  run_result = last_run_result(exitcode)

  if exitcode != 0
    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, params)
  exitcode = DEFAULT_EXITCODE
  cmd = make_command_array(config, params)
  env = make_environment_hash(params)
  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

  if disabled?(run_result.to_s)
    return make_error_result(exitcode, Errors::Disabled,
                             "Puppet agent is disabled")
  elsif running?(run_result.to_s)
    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]
        },
        :output => {
          :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(params_and_config)
  if !params_and_config
    return make_error_result(DEFAULT_EXITCODE, Errors::InvalidJson,
                             "Invalid json parsed on STDIN. Cannot start run action")
  end

  params = params_and_config["params"]
  config = params_and_config["config"]
  params["flags"] = params["flags"] | ["--onetime", "--no-daemonize", "--verbose"]

  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

  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, params)
end

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

  if action == 'metadata'
    puts metadata.to_json
  else
    params_and_config = JSON.parse($stdin.read.chomp) rescue nil

    run_result = run(params_and_config)
    puts run_result.to_json

    if !run_result["error"].nil?
      exit 1
    end
  end
end
