#!/usr/bin/env ruby

require 'optparse'
require 'regexp-examples'
require 'rcmd'

=begin rdoc
# This program is for executing a command on multiple nodes at the same time and accepts the following arguments.
#
# * -u, --username    | Username for the connections
# * -n, --nodes       | comma seperated list of nodes. '-' to use piped list
# * -t, --threads     | Number of threads to run
# * -c, --command     | Command to be executed
# * -q, --quiet       | Boolean flag for suppressing stdout of executed commands
# * -v, --version     | Print version string and exit

=end
# Set default options
options = { :threads => 4, :nodes => nil, :environment => nil, :user => 'root', :password => false,
           :command => nil, :quiet => false, :version => false, :debug => false, :expression => nil,
           :database => false, :config => false, :type => nil, :host => nil, :os => nil ,
           :timeout => 2, :hosts_file => File.expand_path('~/.ssh/known_hosts'), :keys_only => false }

host_list = []

# Setup options
opts = OptionParser.new #:nodoc:

opts.on('-u username', '--username username', String, "Username for SSH connections") { |v| options[:user] = v }
opts.on('-n nodes', '--nodes x,y,z', Array, "Comma seperated list of nodes. use '-' for a space seperated list piped from another command") { |v| options[:nodes] = v }
opts.on('-r regex', '--regexp regex', Regexp, "Use Regex to build host list (ruby regexp)") { |v| options[:expression] = v}
opts.on('-t threads', '--threads threads', Integer, "Number of threads to run") { |v| options[:threads] = v }
opts.on('-c <command>', '--command <command>', String, "Quoted string containing the command to be run") { |v| options[:command] = v }
opts.on('-q', '--quiet', "Suppress stdout of commands.  stderr will still be displayed") { |v| options[:quiet] = v }
opts.on('-v', '--version', "Print what version of the command is in use") { |v| options[:version] = v}
opts.on('-D', '--debug', "Print debug information") { |v| options[:debug] = v }
opts.on('--timeout Seconds', Integer, "Number of seconds to wait before timeing out the session") { |v| options[:timeout] = v }
opts.on('--hosts-file FILE', String, "Filename to use as the SSH hosts file (known_hosts)") { |v| options[:hosts_file] = File.file?(File.expand_path(v)) ? File.expand_path(v) : File.expand_path('~/.ssh/known_hosts') }
opts.on('--keys_only', "Use only ssh keys for authentication. (No Passwords)") { |v| options[:keys_only] = v } 
opts.separator " "
opts.separator "Database Options"
opts.on('-C', '--create-config', "Create template dbconfig file in ~/.rcmd") { |v| options[:config] = v}
opts.on('-T server-type ', '--type server-type', "Database query on server type") { |v| options[:type] = v}
opts.on('-H hostname-pattern', '--host hostname-pattern', "Database query on hostname (sql like query)") { |v| options[:host] = v}
opts.on('-O operating-system', '--os operating-system', "Database query on Operating System string") { |v| options[:os] = v}

# Process options
begin
  opts.parse!(ARGV)
  # need a smarter way then this ugly hack
  unless options[:version] or options[:expression] or options[:config] or options[:type] or options[:host] or options[:os]
      mandatory = [:command, :nodes]
      missing = mandatory.select{ |param| options[param].nil? }
      raise OptionParser::MissingArgument, missing.join(', ') unless missing.empty?
  end
  rescue OptionParser::ParseError => e
    $stderr.puts e
    $stderr.puts opts
    exit
end

# Process Regex for node list generation
if options[:expression]
    options[:nodes] = Regexp.new(options[:expression]).examples(max_group_results: 500)
end

if options[:config]
  Rcmd::DB.create_config
  puts "Configuration file created.  Please adjust accordingly"
  exit
end

# Process database options
if options[:host]
  begin 
    Rcmd::DB.db_connect
    options[:nodes] = Rcmd::DB.query_by_hostname(options[:host])
  rescue e
    puts "Error:: #{e}"
    exit
  ensure
    Rcmd::DB.connection.disconnect!
  end
elsif options[:type]
  begin  
    Rcmd::DB.db_connect
    options[:nodes] = Rcmd::DB.query_by_type(options[:type])
  rescue e
    puts "Error:: #{e}"
    exit
  ensure
    Rcmd::DB.connection.disconnect!
  end
elsif options[:os]
  begin
    Rcmd::DB.db_connect
    options[:nodes] = Rcmd::DB.query_by_os(options[:os])
  rescue e
    puts "Error:: #{e}"
    exit
  ensure
    Rcmd::DB.connection.disconnect!
  end
end

# Print version and exit if version option chosen
if options[:version]
  $stdout.puts "Version: #{Rcmd::VERSION}"
  exit
end

# Populate the host list
if options[:nodes][0] == '-'
  ARGF.read.split().each do |h|
    host_list << h
  end
else
  host_list = options[:nodes]
end

# abort if host_list is empty
abort("No hosts to run command on") if host_list.nil?

# Set needed variables
Rcmd.host_list= host_list
Rcmd.nthreads= options[:threads]
Rcmd.user= options[:user]
Rcmd.quiet= options[:quiet]
Rcmd.command= options[:command]
Rcmd.debug= options[:debug]
Rcmd.timeout = options[:timeout]
Rcmd.keys_only = options[:keys_only]
Rcmd.hosts_file = options[:host_file]

# Execute command on provided hosts
Rcmd.run_command