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

require 'pty'
require 'yaml'

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'
SITE_MODULES      = 'puppet/modules'
CUSTOM_MODULES    = ':files/puppet/modules'
DEFAULT_TAGS      = 'leap_base,leap_service'
HIERA_FILE        = '/etc/leap/hiera.yaml'

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
  hostname = hiera_file['name']
  if hostname.nil? || hostname.empty?
    puts('ERROR: "name" missing from hiera file')
    exit(1)
  end
  current_hostname_file = File.read('/etc/hostname') rescue nil
  current_hostname = `/bin/hostname`.strip

  # set /etc/hostname
  if current_hostname_file != hostname
    File.open('/etc/hostname', 'w', 0611, :encoding => 'ascii') do |f|
      f.write hostname
    end
    if File.read('/etc/hostname') == hostname
      puts "Changed /etc/hostname to #{hostname}"
    else
      puts "ERROR: failed to update /etc/hostname"
    end
  end

  # call /bin/hostname
  if current_hostname != hostname
    if run("/bin/hostname #{hostname}") == 0
      puts "Changed hostname to #{hostname}"
    else
      puts "ERROR: call to `/bin/hostname #{hostname}` returned an error."
    end
  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
  modulepath = options[:module_path] || SITE_MODULES + CUSTOM_MODULES
  fqdn = hiera_file['domain']['name']
  domain = hiera_file['domain']['full_suffix']
  Dir.chdir(PUPPET_DIRECTORY) do
    return run("FACTER_fqdn='#{fqdn}' FACTER_domain='#{domain}' #{PUPPET_BIN} apply #{custom_parameters(options)} --modulepath='#{modulepath}' #{PUPPET_PARAMETERS} #{manifest}", &block)
  end
end

#
# Return a ruby object representing the contents of the hiera yaml file.
#
def hiera_file
  unless File.exists?(HIERA_FILE)
    puts("ERROR: hiera file '#{HIERA_FILE}' does not exist.")
    exit(1)
  end
  $hiera_contents ||= YAML.load_file(HIERA_FILE)
  return $hiera_contents
rescue Exception => exc
  puts("ERROR: problem reading hiera file '#{HIERA_FILE}' (#{exc})")
  exit(1)
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
##

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