#!/usr/bin/ruby

#
# This is a wrapper script around the puppet command used by the LEAP platform.
#
# We do this in order to make it faster and easier to control puppet remotely
# (exit codes, lockfile, multiple manifests, etc)
#

PUPPET_BIN        = '/usr/bin/puppet'
PUPPET_DIRECTORY  = '/srv/leap'
PUPPET_PARAMETERS = '--color=false --detailed-exitcodes --libdir=puppet/lib --confdir=puppet'
SITE_MANIFEST     = 'puppet/manifests/site.pp'
SETUP_MANIFEST    = 'puppet/manifests/setup.pp'
DEFAULT_TAGS      = 'leap_base,leap_service'

def main
  process_command_line_arguments
  with_lockfile do
    @commands.each do |command|
      self.send(command)
    end
  end
end

def puts(str)
  $stdout.puts str
  $stdout.flush
end

def process_command_line_arguments
  @commands = []
  @verbosity = 1
  @tags = DEFAULT_TAGS
  loop do
    case ARGV[0]
      when 'apply'        then ARGV.shift; @commands << 'apply'
      when 'set_hostname' then ARGV.shift; @commands << 'set_hostname'
      when '--verbosity'  then ARGV.shift; @verbosity = ARGV.shift.to_i
      when '--force'      then ARGV.shift; remove_lockfile
      when '--tags'       then ARGV.shift; @tags      = ARGV.shift
      when /^-/           then usage("Unknown option: #{ARGV[0].inspect}")
      else break
    end
  end
  usage("No command given") unless @commands.any?
end

def apply
  exit_code = puppet_apply do |line|
    puts line
  end
  puts "Puppet apply complete (#{exitcode_description(exit_code)})."
end

def set_hostname
  exit_code = puppet_apply(:manifest => SETUP_MANIFEST, :tags => '') do |line|
    # todo: replace setup.pp with https://github.com/lutter/ruby-augeas
    # or try this: http://www.puppetcookbook.com/posts/override-a-facter-fact.html
    if (line !~ /Finished catalog run/ || @verbosity > 2) &&
       (line !~ /dnsdomainname: Name or service not known/) &&
       (line !~ /warning: Could not retrieve fact fqdn/)
      puts line
    end
  end
  if exit_code == 2
    puts "Hostname updated."
  elsif exit_code == 4 || exit_code == 6
    puts "ERROR: could not update hostname."
  elsif exit_code == 0 && @verbosity > 1
    puts "No change to hostname."
  end
end

#
# each line of output is yielded. the exit code is returned.
#
def puppet_apply(options={}, &block)
  options = {:verbosity => @verbosity, :tags => @tags}.merge(options)
  manifest = options[:manifest] || SITE_MANIFEST
  Dir.chdir(PUPPET_DIRECTORY) do
    return run("#{PUPPET_BIN} apply #{custom_parameters(options)} #{PUPPET_PARAMETERS} #{manifest}", &block)
  end
end

def custom_parameters(options)
  params = []
  if options[:tags] && options[:tags].chars.any?
    params << "--tags #{options[:tags]}"
  end
  if options[:verbosity]
    case options[:verbosity]
      when 3 then params << '--verbose'
      when 4 then params << '--verbose --debug'
      when 5 then params << '--verbose --debug --trace'
    end
  end
  params.join(' ')
end

def exitcode_description(code)
  case code
    when 0 then "no changes"
    when 1 then "failed"
    when 2 then "changes made"
    when 4 then "failed"
    when 6 then "changes and failures"
    else code
  end
end

def usage(s)
  $stderr.puts(s)
  $stderr.puts
  $stderr.puts("Usage: #{File.basename($0)} COMMAND [OPTIONS]")
  $stderr.puts
  $stderr.puts("COMMAND may be one or more of:
  set_hostname     -- set the hostname of this server
  apply            -- apply puppet manifests")
  $stderr.puts
  $stderr.puts("OPTIONS may be one or more of:
  --verbosity VERB -- set the verbosity level 0..5
  --tags TAGS      -- set the tags to pass through to puppet
  --force          -- run even when lockfile is present")
  exit(2)
end

##
## Simple lock file
##

require 'fileutils'
DEFAULT_LOCKFILE = '/tmp/puppet.lock'

def remove_lockfile(lock_file_path=DEFAULT_LOCKFILE)
  FileUtils.remove_file(lock_file_path, true)
end

def with_lockfile(lock_file_path=DEFAULT_LOCKFILE)
  begin
    File.open(lock_file_path, File::CREAT | File::EXCL | File::WRONLY) do |o|
      o.write(Process.pid)
    end
    yield
    remove_lockfile
  rescue Errno::EEXIST
    puts("ERROR: the lock file '#{lock_file_path}' already exists. Wait a minute for the process to die, or run with --force to ignore. Bailing out.")
    exit(1)
  rescue IOError => exc
    puts("ERROR: problem with lock file '#{lock_file_path}' (#{exc}). Bailing out.")
    exit(1)
  end
end

##
## simple pass through process runner (to ensure output is not buffered and return exit code)
## this only works under ruby 1.9
##

require "pty"

def run(cmd)
  puts cmd if @verbosity >= 3
  PTY.spawn("#{cmd}") do |output, input, pid|
    begin
      while line = output.gets do
        yield line
        #$stdout.puts line
        #$stdout.flush
      end
    rescue Errno::EIO
    end
    Process.wait(pid) # only works in ruby 1.9, required to capture the exit status.
  end
  status = $?.exitstatus
  #yield status if block_given?
  return status
rescue PTY::ChildExited
end

##
## RUN MAIN
##

Signal.trap("EXIT") do
  remove_lockfile # clean up the lockfile when process is terminated.
                  # this will remove the lockfile if ^C killed the process
                  # but only after the child puppet process is also dead (I think).
end

main()