#!/usr/bin/env ruby

#
# = Synopsis
#
# Trigger a puppetd run on a set of hosts.
#
# = Usage
#
#   puppetrun [-a|--all] [-c|--class <class>] [-d|--debug] [-f|--foreground]
#       [-h|--help] [--host <host>] [--no-fqdn] [--ignoreschedules]
#       [-t|--tag <tag>] [--test] [-p|--ping]
#
# = Description
#
# This script can be used to connect to a set of machines running +puppetd+
# and trigger them to run their configurations.  The most common usage would
# be to specify a class of hosts and a set of tags, and +puppetrun+ would
# look up in LDAP all of the hosts matching that class, then connect to
# each host and trigger a run of all of the objects with the specified tags.
#
# If you are not storing your host configurations in LDAP, you can specify
# hosts manually.
#
# You will most likely have to run +puppetrun+ as root to get access to
# the SSL certificates.
#
# +puppetrun+ reads +puppetmaster+'s configuration file, so that it can copy
# things like LDAP settings.
#
# = Usage Notes
#
# +puppetrun+ is useless unless +puppetd+ is listening.  See its documentation
# for more information, but the gist is that you must enable +listen+ on the
# +puppetd+ daemon, either using +--listen+ on the command line or adding
# 'listen: true' in its config file.  In addition, you need to set the daemons
# up to specifically allow connections by creating the +namespaceauth+ file,
# normally at '/etc/puppet/namespaceauth.conf'.  This file specifies who has
# access to each namespace; if you create the file you must add every namespace
# you want any Puppet daemon to allow -- it is currently global to all Puppet
# daemons.
#
# An example file looks like this::
#
#     [fileserver]
#         allow *.madstop.com
# 
#     [puppetmaster]
#         allow *.madstop.com
# 
#     [puppetrunner]
#         allow culain.madstop.com
#
# This is what you would install on your Puppet master; non-master hosts could
# leave off the 'fileserver' and 'puppetmaster' namespaces.
#
# Expect more documentation on this eventually.
#
# = Options
#
# Note that any configuration parameter that's valid in the configuration file
# is also a valid long argument.  For example, 'ssldir' is a valid configuration
# parameter, so you can specify '--ssldir <directory>' as an argument.
#
# See the configuration file documentation at
# http://reductivelabs.com/projects/puppet/reference/configref.html for
# the full list of acceptable parameters. A commented list of all
# configuration options can also be generated by running puppetmasterdd with
# '--genconfig'.
#
#
# all::
#   Connect to all available hosts.  Requires LDAP support at this point.
#
# class::
#   Specify a class of machines to which to connect.  This only works if you
#   have LDAP configured, at the moment.
#
# debug::
#   Enable full debugging.
#
# foreground::
#   Run each configuration in the foreground; that is, when connecting to a host,
#   do not return until the host has finished its run.  The default is false.
#
# help::
#   Print this help message
#
# host::
#   A specific host to which to connect.  This flag can be specified more
#   than once.
#
# ignoreschedules::
#   Whether the client should ignore schedules when running its configuration.
#   This can be used to force the client to perform work it would not normally
#   perform so soon.  The default is false.
#
# parallel::
#   How parallel to make the connections.  Parallelization is provided by forking
#   for each client to which to connect.  The default is 1, meaning serial execution.
#
# tag::
#   Specify a tag for selecting the objects to apply. Does not work with the
#   --test option.
#
#
# test::
#   Print the hosts you would connect to but do not actually connect. This
#   option requires LDAP support at this point.
#
# ping::
#   
#   Do a ICMP echo against the target host. Skip hosts that don't respond to ping.
#
# = Example
#
#   sudo puppetrun -p 10 --host host1 --host host2 -t remotefile -t webserver
#
# = Author
#
# Luke Kanies
#
# = Copyright
#
# Copyright (c) 2005 Reductive Labs, LLC
# Licensed under the GNU Public License

[:INT, :TERM].each do |signal|
    trap(signal) do
        $stderr.puts "Cancelling"
        exit(1)
    end
end

begin
    require 'rubygems'
rescue LoadError
    # Nothing; we were just doing this just in case
end

begin
    require 'ldap'
rescue LoadError
    $stderr.puts "Failed to load ruby LDAP library. LDAP functionality will not be available"
end

require 'puppet'
require 'puppet/network/client'
require 'puppet/util/ldap/connection'
require 'getoptlong'

flags = [
    [ "--all",      "-a",	        GetoptLong::NO_ARGUMENT ],
    [ "--tag",      "-t",           GetoptLong::REQUIRED_ARGUMENT ],
    [ "--class",	"-c",			GetoptLong::REQUIRED_ARGUMENT ],
    [ "--foreground", "-f",	        GetoptLong::NO_ARGUMENT ],
    [ "--debug",	"-d",			GetoptLong::NO_ARGUMENT ],
    [ "--help",		"-h",			GetoptLong::NO_ARGUMENT ],
    [ "--host",         			GetoptLong::REQUIRED_ARGUMENT ],
    [ "--parallel", "-p",			GetoptLong::REQUIRED_ARGUMENT ],
    [ "--ping", "-P",			GetoptLong::NO_ARGUMENT ],
    [ "--no-fqdn",  "-n",			GetoptLong::NO_ARGUMENT ],
    [ "--test",                     GetoptLong::NO_ARGUMENT ],
    [ "--version",  "-V",           GetoptLong::NO_ARGUMENT ]
]

# Add all of the config parameters as valid options.
Puppet.settings.addargs(flags)

result = GetoptLong.new(*flags)

options = {
    :ignoreschedules => false,
    :foreground => false,
    :parallel => 1,
    :ping => false,
    :debug => false,
    :test => false,
    :all => false,
    :verbose => true,
    :fqdn => true
}

hosts = []
classes = []
tags = []

Puppet::Util::Log.newdestination(:console)

begin
    result.each { |opt,arg|
        case opt
            when "--version"
                puts "%s" % Puppet.version
                exit
            when "--ignoreschedules"
                options[:ignoreschedules] = true
            when "--no-fqdn"
                options[:fqdn] = false
            when "--all"
                options[:all] = true
            when "--test"
                options[:test] = true
            when "--tag"
                tags << arg
            when "--class"
                classes << arg
            when "--host"
                hosts << arg
            when "--help"
                if Puppet.features.usage?
                    RDoc::usage && exit
                else
                    puts "No help available unless you have RDoc::usage installed"
                    exit
                end
            when "--parallel"
                begin
                    options[:parallel] = Integer(arg)
                rescue
                    $stderr.puts "Could not convert %s to an integer" % arg.inspect
                    exit(23)
                end
            when "--ping"
                 options[:ping] = true 
            when "--foreground"
                options[:foreground] = true
            when "--debug"
                options[:debug] = true
            else
                Puppet.settings.handlearg(opt, arg)
        end
    }
rescue GetoptLong::InvalidOption => detail
    $stderr.puts "Try '#{$0} --help'"
    exit(1)
end

if options[:debug]
    Puppet::Util::Log.level = :debug
else
    Puppet::Util::Log.level = :info
end

# Now parse the config
Puppet.parse_config

if Puppet[:node_terminus] == "ldap" and (options[:all] or classes)
    if options[:all]
        hosts = Puppet::Node.search("whatever").collect { |node| node.name }
        puts "all: %s" % hosts.join(", ")
    else
        hosts = []
        classes.each do |klass|
            list = Puppet::Node.search("whatever", :class => klass).collect { |node| node.name }
            puts "%s: %s" % [klass, list.join(", ")]

            hosts += list
        end
    end
elsif ! classes.empty?
    $stderr.puts "You must be using LDAP to specify host classes"
    exit(24)
end

if tags.empty?
    tags = ""
else
    tags = tags.join(",")
end

children = {}

# If we get a signal, then kill all of our children and get out.
[:INT, :TERM].each do |signal|
    trap(signal) do
        Puppet.notice "Caught #{signal}; shutting down"
        children.each do |pid, host|
            Process.kill("INT", pid)
        end

        waitall

        exit(1)
    end
end

if options[:test]
    puts "Skipping execution in test mode"
    exit(0)
end

todo = hosts.dup

failures = []

# Now do the actual work
go = true
while go
    # If we don't have enough children in process and we still have hosts left to
    # do, then do the next host.
    if children.length < options[:parallel] and ! todo.empty?
        host = todo.shift
        pid = fork do
            if options[:ping]
              out = %x{ping -c 1 #{host}}
              unless $? == 0
                  $stderr.print "Could not contact %s\n" % host
                  next
              end
            end
            client = Puppet::Network::Client.runner.new(
                :Server => host,
                :Port => Puppet[:puppetport]
            )

            print "Triggering %s\n" % host
            begin
                result = client.run(tags, options[:ignoreschedules], options[:foreground])
            rescue => detail
                $stderr.puts "Host %s failed: %s\n" % [host, detail]
                exit(2)
            end
            
            case result
            when "success": exit(0)
            when "running":
                $stderr.puts "Host %s is already running" % host
                exit(3)
            else
                $stderr.puts "Host %s returned unknown answer '%s'" % [host, result]
                exit(12)
            end
        end
        children[pid] = host
    else
        # Else, see if we can reap a process.
        begin
            pid = Process.wait

            if host = children[pid]
                # Remove our host from the list of children, so the parallelization
                # continues working.
                children.delete(pid)
                if $?.exitstatus != 0
                    failures << host
                end
                print "%s finished with exit code %s\n" % [host, $?.exitstatus]
            else
                $stderr.puts "Could not find host for PID %s with status %s" %
                    [pid, $?.exitstatus]
            end
        rescue Errno::ECHILD
            # There are no children left, so just exit unless there are still
            # children left to do.
            next unless todo.empty?

            if failures.empty?
                puts "Finished"
                exit(0)
            else
                puts "Failed: %s" % failures.join(", ")
                exit(3)
            end
        end
    end
end