pxp-module-puppet

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

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

module Pxp
  class ModulePuppet
    module Errors
      InvalidJson = "invalid_json"
      NoPuppetBin = "no_puppet_bin"
      NoLastRunReport = "no_last_run_report"
      InvalidLastRunReport = "invalid_last_run_report"
      Disabled = "agent_disabled"
      Locked = "agent_locked"
      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

    attr_reader :config, :flags

    def self.handle_action(action)
      if action == 'metadata'
        puts metadata.to_json
      else
        result = create_runner($stdin.read.chomp)
        if result.is_a?(self)
          action_results = result.run
        else
          action_results = result
        end

        print action_results.to_json

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

    DEFAULT_EXITCODE = -1

    def get_env_fix_up
      # 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 self.last_run_result(exitcode)
      return {"time"             => "unknown",
              "transaction_uuid" => "unknown",
              "environment"      => "unknown",
              "status"           => "unknown",
              "metrics"          => {},
              "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, Encoding::UndefinedConversionError
        # 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 config_print(*keys)
      command_array = [config["puppet_bin"], "agent", "--configprint", keys.join(',')]
      process_output = Puppet::Util::Execution.execute(command_array,
                                                       {:custom_environment => get_env_fix_up(),
                                                        :override_locale => false})

      result = force_unicode(process_output.to_s)
      if keys.count == 1
        result.chomp
      else
        result.lines.inject({}) do |conf, line|
          key, value = line.chomp.split(' = ', 2)
          if key && value
            conf[key] = value
          end
          conf
        end
      end
    end

    def running?(lockfile)
      return File.exist?(lockfile)
    end

    def disabled?(lockfile)
      return File.exist?(lockfile)
    end

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

    def puppet_agent_command
      cmd_array = [config["puppet_bin"], "agent"]
      cmd_array += flags
      return cmd_array
    end

    def self.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 make_error_result(exitcode, error_type, error_message)
      self.class.make_error_result(exitcode, error_type, error_message)
    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 nest_metrics(metrics)
      metrics.fetch('resources', {}).fetch('values', {}).inject({}) do |result, (name,human_name,value)|
        result.merge(name => value)
      end
    end

    def get_result_from_report(last_run_report, exitcode, 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,
                                 "The Puppet run failed in an unexpected way")
      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 = self.class.last_run_result(exitcode)
      else
        run_result = make_error_result(exitcode, Errors::NonZeroExit,
                                       "Puppet agent exited with a non 0 exitcode")
      end

      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']
      run_result["metrics"] = nest_metrics(last_run_report_yaml['metrics'])

      return run_result
    end

    # Wait for the lockfile to be removed. If it hasn't after 10 minutes, give up.
    def wait_for_lockfile(lockfile, check_interval = 0.1, give_up_after = 10*60)
      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

    # Determine whether the configured puppet bin exists. This method mostly
    # exists for testing.
    def puppet_bin_present?
      File.exist?(config["puppet_bin"])
    end

    def get_start_time(last_run_report)
      File.mtime(last_run_report) if File.exist?(last_run_report)
    end

    def try_run(last_run_report)
      start_time = get_start_time(last_run_report)
      run_result = Puppet::Util::Execution.execute(puppet_agent_command, {:failonfail => false,
                                                                          :custom_environment => make_environment_hash(),
                                                                          :override_locale => false})
      return start_time, (run_result ? run_result.exitstatus : nil)
    end

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

      puppet_config = config_print('lastrunreport', 'agent_disabled_lockfile', 'agent_catalog_run_lockfile')
      last_run_report = puppet_config['lastrunreport']

      if last_run_report.nil? || last_run_report.empty?
        return make_error_result(DEFAULT_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.
      start_time, exitcode = try_run(last_run_report)
      if exitcode.nil?
        return make_error_result(DEFAULT_EXITCODE, Errors::FailedToStart,
                                 "Failed to start Puppet agent")
      end

      # If the run was successful, don't check for failure modes.
      if exitcode != 0
        if disabled?(puppet_config['agent_disabled_lockfile'] || '')
          return make_error_result(exitcode, Errors::Disabled,
                                   "Puppet agent is disabled")
        end

        # Check for a lockfile. If present, wait until it's removed and try running again.
        # There's a chance that our run finished with a real error rather than because Puppet was
        # already running, but another run started immediately after. Since we have no
        # language-agnostic way to tell, we accept that we might run twice in that case.
        # The run could also finish immediately after we tried, and the lockfile be absent.
        # In that case we'll fail with poor error reporting.
        lockfile = puppet_config['agent_catalog_run_lockfile'] || ''
        if running?(lockfile)
          wait_for_lockfile(lockfile)

          start_time, exitcode = try_run(last_run_report)
          if exitcode.nil?
            return make_error_result(DEFAULT_EXITCODE, Errors::FailedToStart,
                                     "Failed to start Puppet agent")
          end

          if exitcode != 0
            if disabled?(puppet_config['agent_disabled_lockfile'] || '')
              return make_error_result(exitcode, Errors::Disabled,
                                       "Puppet agent is disabled")
            end

            if running?(lockfile)
              return make_error_result(exitcode, Errors::Locked,
                                       "Puppet agent run is already in progress")
            end
          end
        end
      end

      return get_result_from_report(last_run_report, exitcode, start_time)
    end

    # TODO(ale): remove `env` from input before bumping to the next version
    def self.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"
                  }
                },
                :job => {
                  :type => "string"
                }
              },
              :required => [:flags]
            },
            :results => {
              :type => "object",
              :properties => {
                :time => {
                  :type => "string"
                },
                :transaction_uuid => {
                  :type => "string"
                },
                :metrics => {
                  :type => "object"
                },
                :environment => {
                  :type => "string"
                },
                :status => {
                  :type => "string"
                },
                :error_type => {
                  :type => "string"
                },
                :error => {
                  :type => "string"
                },
                :exitcode => {
                  :type => "number"
                },
                :version => {
                  :type => "number"
                }
              },
              :required => [:time, :transaction_uuid, :environment, :status,
                            :exitcode, :version]
            }
          }
        ],
        :configuration => {
          :type => "object",
          :properties => {
            :puppet_bin => {
              :type => "string"
            }
          }
        }
      }
    end

    def self.add_config_defaults(config)
      config = config.dup
      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 self.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 self.process_flags(action_input)
      flags = action_input['flags'] || []

      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

      flags |= DEFAULT_FLAGS

      if action_input.has_key?("job")
        flags += ["--job-id", action_input["job"]]
      end

      flags
    end

    def self.create_runner(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', 0640))
          $stderr.reopen(File.open(output_files["stderr"], 'w', 0640))
        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', 0640) 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 = add_config_defaults(args["configuration"] || {})

        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

        flags = process_flags(action_input)

        new(config, flags)
      rescue ProcessingError => e
        return make_error_result(DEFAULT_EXITCODE, e.error_type, e.message)
      end
    end

    def initialize(config, flags)
      @config = config
      @flags = flags
    end

  end
end

if __FILE__ == $0
  action = ARGV.shift || 'metadata'
  Pxp::ModulePuppet.handle_action(action)
end