#!/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, logging, lockfile, version check, etc)
#

require 'pty'
require 'yaml'
require 'logger'
require 'socket'
require 'fileutils'

DEBIAN_VERSION    = /^(jessie|8\.)/
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'
LOG_DIR           = '/var/log/leap'
DEPLOY_LOG        = '/var/log/leap/deploy.log'
SUMMARY_LOG       = '/var/log/leap/deploy-summary.log'
SUMMARY_LOG_1     = '/var/log/leap/deploy-summary.log.1'
APPLY_START_STR   = "STARTING APPLY"
APPLY_FINISH_STR  = "APPLY COMPLETE"
LANG              = "en_US.UTF-8"

def main
  if File.read('/etc/debian_version') !~ DEBIAN_VERSION
    log "ERROR: This operating system is not supported. The file /etc/debian_version must match #{DEBIAN_VERSION}."
    exit 1
  end
  process_command_line_arguments
  with_lockfile do
    @commands.each do |command|
      self.send(command)
    end
  end
end

def open_log_files
  FileUtils.mkdir_p(LOG_DIR)
  $logger = Logger.new(DEPLOY_LOG)
  $summary_logger = Logger.new(SUMMARY_LOG)
  [$logger, $summary_logger].each do |logger|
    logger.level = Logger::INFO
    logger.formatter = proc do |severity, datetime, progname, msg|
      "%s %s: %s\n" % [datetime.strftime("%b %d %H:%M:%S"), Socket.gethostname, msg]
    end
  end
end

def close_log_files
  $logger.close
  $summary_logger.close
end

def log(str, *args)
  str = str.strip
  $stdout.puts str
  $stdout.flush
  if $logger
    $logger.info(str)
    if args.include? :summary
      $summary_logger.info(str)
    end
  end
end

def process_command_line_arguments
  @commands = []
  @verbosity = 1
  @tags = DEFAULT_TAGS
  @info = {}
  @downgrade = false
  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 '--info'       then ARGV.shift; @info      = parse_info(ARGV.shift)
      when '--downgrade'  then ARGV.shift; @downgrade = true
      when /^-/           then usage("Unknown option: #{ARGV[0].inspect}")
      else break
    end
  end
  usage("No command given") unless @commands.any?
end

def apply
  platform_version_check! unless @downgrade
  log "#{APPLY_START_STR} {#{format_info(@info)}}", :summary
  exit_code = puppet_apply do |line|
    log line
  end
  log "#{APPLY_FINISH_STR} (#{exitcode_description(exit_code)}) {#{format_info(@info)}}", :summary
end

def set_hostname
  hostname = hiera_file['name']
  if hostname.nil? || hostname.empty?
    log('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
      log "Changed /etc/hostname to #{hostname}"
    else
      log "ERROR: failed to update /etc/hostname"
    end
  end

  # call /bin/hostname
  if current_hostname != hostname
    if run("/bin/hostname #{hostname}") == 0
      log "Changed hostname to #{hostname}"
    else
      log "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']['full']
  domain = hiera_file['domain']['full_suffix']
  Dir.chdir(PUPPET_DIRECTORY) do
    return run("LANG='#{LANG}' FACTER_fqdn='#{fqdn}' FACTER_domain='#{domain}' #{PUPPET_BIN} apply #{custom_parameters(options)} --modulepath='#{modulepath}' #{PUPPET_PARAMETERS} #{manifest}", &block)
  end
end

#
# parse the --info flag. example str: "key1: value1, key2: value2, ..."
#
def parse_info(str)
  str.split(', ').
    map {|i| i.split(': ')}.
    inject({}) {|h,i| h[i[0]] = i[1]; h}
rescue Exception => exc
  {"platform" => "INVALID_FORMAT"}
end

def format_info(info)
  info.to_a.map{|i|i.join(': ')}.join(', ')
end

#
# exits with a warning message if the last successful deployed
# platform was newer than the one we are currently attempting to
# deploy.
#
PLATFORM_RE = /\{.*platform: ([0-9\.]+)[ ,\}].*[\}$]/
def platform_version_check!
  return unless @info["platform"]
  new_version = @info["platform"].split(' ').first
  return unless new_version
  if File.exist?(SUMMARY_LOG) && File.size(SUMMARY_LOG) != 0
    file = SUMMARY_LOG
  elsif File.exist?(SUMMARY_LOG_1) && File.size(SUMMARY_LOG_1) != 0
    file = SUMMARY_LOG_1
  else
    return
  end
  most_recent_line = `tail '#{file}'`.split("\n").grep(PLATFORM_RE).last
  if most_recent_line
    prior_version = most_recent_line.match(PLATFORM_RE)[1]
    if Gem::Version.new(prior_version) > Gem::Version.new(new_version)
      log("ERROR: You are attempting to deploy platform v#{new_version} but this node uses v#{prior_version}.")
      log("       Run with --downgrade if you really want to deploy an older platform version.")
      exit(0)
    end
  end
end

#
# Return a ruby object representing the contents of the hiera yaml file.
#
def hiera_file
  unless File.exist?(HIERA_FILE)
    log("ERROR: hiera file '#{HIERA_FILE}' does not exist.")
    exit(1)
  end
  $hiera_contents ||= YAML.load_file(HIERA_FILE)
  return $hiera_contents
rescue Exception => exc
  log("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.
  --info           -- additional info to include in logs (e.g. 'user: alice, platform: 0.6.1')
  --downgrade      -- allow a deploy even if the platform version is older than previous deploy.
  ")
  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
    open_log_files
    yield
    remove_lockfile
    close_log_files
  rescue Errno::EEXIST
    log("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
    log("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)
  log(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()