chef-mix-fix

#!/usr/bin/env ruby
#
# Find Chef nodes which have not checked in with the server in some time, and
# attempt to force a chef-client run on them.
#

require 'rubygems'
require 'chef'
require 'net/ssh/multi'
require 'trollop'

#
# Author:: Adam Jacob (<[email protected]>)
# Author:: John Keiser (<[email protected]>)
# Copyright:: Copyright (c) 2012 Opscode, Inc.
# License:: Apache License, Version 2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

require 'chef/config'
require 'uri'
require 'chef/rest'
# These are needed so that JSON can inflate search results
require 'chef/node'
require 'chef/role'
require 'chef/environment'
require 'chef/data_bag'
require 'chef/data_bag_item'

class Chef
  class PartialSearch

    attr_accessor :rest

    def initialize(url=nil)
      @rest = ::Chef::REST.new(url || ::Chef::Config[:chef_server_url])
    end

    # Search Solr for objects of a given type, for a given query. If you give
    # it a block, it will handle the paging for you dynamically.
    def search(type, query='*:*', args={}, &block)
      raise ArgumentError, "Type must be a string or a symbol!" unless (type.kind_of?(String) || type.kind_of?(Symbol))

      sort = args.include?(:sort) ? args[:sort] : 'X_CHEF_id_CHEF_X asc'
      start = args.include?(:start) ? args[:start] : 0
      rows = args.include?(:rows) ? args[:rows] : 1000
      query_string = "search/#{type}?q=#{escape(query)}&sort=#{escape(sort)}&start=#{escape(start)}&rows=#{escape(rows)}"
      if args[:keys]
        response = @rest.post_rest(query_string, args[:keys])
        response_rows = response['rows'].map { |row| row['data'] }
      else
        response = @rest.get_rest(query_string)
        response_rows = response['rows']
      end
      if block
        response_rows.each { |o| block.call(o) unless o.nil?}
        unless (response["start"] + response_rows.length) >= response["total"]
          nstart = response["start"] + rows
          args_hash = {
            :keys => args[:keys],
            :sort => sort,
            :start => nstart,
            :rows => rows
          }
          search(type, query, args_hash, &block)
        end
        true
      else
        [ response_rows, response["start"], response["total"] ]
      end
    end

    def list_indexes
      response = @rest.get_rest("search")
    end

    private
      def escape(s)
        s && URI.escape(s.to_s)
      end
  end
end

# partial_search(type, query, options, &block)
#
# Searches for nodes, roles, etc. and returns the results.  This method may
# perform more than one search request, if there are a large number of results.
#
# ==== Parameters
# * +type+: index type (:role, :node, :client, :environment, data bag name)
# * +query+: SOLR query.  "*:*", "role:blah", "not role:blah", etc.  Defaults to '*:*'
# * +options+: hash with options:
# ** +:start+: First row to return (:start => 50, :rows => 100 means "return the
#               50th through 150th result")
# ** +:rows+: Number of rows to return.  Defaults to 1000.
# ** +:sort+: a SOLR sort specification.  Defaults to 'X_CHEF_id_CHEF_X asc'.
# ** +:keys+: partial search keys.  If this is not specified, the search will
#             not be partial.
#
# ==== Returns
#
# This method returns an array of search results.  Partial search results will
# be JSON hashes with the structure specified in the +keys+ option.  Other
# results include +Chef::Node+, +Chef::Role+, +Chef::Client+, +Chef::Environment+,
# +Chef::DataBag+ and +Chef::DataBagItem+ objects, depending on the search type.
#
# If a block is specified, the block will be called with each result instead of
# returning an array.  This method will not block if it returns
#
# If start or row is specified, and no block is given, the result will be a
# triple containing the list, the start and total:
#
#     [ [ row1, row2, ... ], start, total ]
#
# ==== Example
#
#     partial_search(:node, 'role:webserver',
#                    keys: {
#                      name: [ 'name' ],
#                      ip: [ 'amazon', 'ip', 'public' ]
#                    }
#     ).each do |node|
#       puts "#{node[:name]}: #{node[:ip]}"
#     end
#
def partial_search(type, query='*:*', *args, &block)
  # Support both the old (positional args) and new (hash args) styles of calling
  if args.length == 1 && args[0].is_a?(Hash)
    args_hash = args[0]
  else
    args_hash = {}
    args_hash[:sort] = args[0] if args.length >= 1
    args_hash[:start] = args[1] if args.length >= 2
    args_hash[:rows] = args[2] if args.length >= 3
  end
  # If you pass a block, or have the start or rows arguments, do raw result parsing
  if Kernel.block_given? || args_hash[:start] || args_hash[:rows]
    Chef::PartialSearch.new.search(type, query, args_hash, &block)
  # Otherwise, do the iteration for the end user
  else
    results = Array.new
    Chef::PartialSearch.new.search(type, query, args_hash) do |o|
      results << o
    end
    results
  end
end

def interval_to_string(curtime, checktime)
  hour = 60 * 60
  day = 24 * 60 * 60

  days = (curtime - checktime) / day
  hours = (curtime - checktime - (days * day)) / hour
  mins = (curtime - checktime - (days * day) - (hours * hour)) / 60

  timearr = []
  timearr << "#{days}d" if days > 0
  timearr << "#{hours}h" if hours > 0
  timearr << "#{mins}m" if mins > 0

  timearr.join(' ')
end

def find_mia_nodes(curtime, hours)
  mianodes = []
  threshold = hours * (60 * 60)
  partial_search(:node, '*:*',
    keys: {
      name: [ 'name' ],
      platform: [ 'platform' ],
      platform_version: [ 'platform_version' ],
      ohai_time: [ 'ohai_time' ],
      ip: [ 'ipaddress' ]
    }).each do |node|
    if (curtime - node['ohai_time'].to_i) >= threshold
      # print node.to_s + "\n"
      mianodes << {
        'name' => node['name'].downcase,
        'os' => "#{node['platform']} #{node['platform_version']}",
        'ip' => node['ip'],
        'time' => interval_to_string(curtime, node['ohai_time'].to_i)
      } unless node['ip'].nil?
    end
  end

  mianodes
end

# get account credentials, db to update, and the permissions filename from the command line
whoami = `whoami`.chomp
opts = Trollop::options do
  opt :username, 'Unix username', :type => :string, :short => 'u', :required => false, :default => whoami
  opt :winuser, 'Windows username', :type => :string, :short => 'w', :required => false, :default => whoami
  opt :connections, 'Concurrent connections', :type => :integer, :short => 'c', :required => false, :default => 20
  opt :hours, 'Hours since last run', :type => :integer, :short => 'h', :required => false, :default => 7
  opt :ssh_copy_id, 'Invoke ssh-copy-id', :type => :boolean, :short => 's', :required => false, :default => true
  opt :knife, 'Path to knife.rb', :type => :string, :short => 'k', :required => false, :default => "#{ENV['HOME']}/.chef/knife.rb"
  opt :winpass, 'Windows password', :type => :string, :short => 'P', :required => false
  opt :windows, 'Windows support', :type => :boolean, :default => true
  opt :loglevel, 'Log level to pass to chef', :type => :string, :short => 'l', :required => false, :default => 'info'
end

# setup chef config
Chef::Config.from_file(opts[:knife])

# get current time
curtime = Time.now.to_i
ignore_nodes = []
sshnodes = []
winnodes = []

find_mia_nodes(curtime, opts[:hours]).each do |node|
  unless node['name'].end_with?('vcloud')
    print "  >> #{node['name']} last ran #{node['time']} "
    print "(OS = #{node['os']}, IP = #{node['ip']}) <<\n"
    if node['os'].start_with?('windows')
      winnodes << node['name']
    else
      sshnodes << node['ip']
    end
  end
end

#
# Use WinRM to rerun the chef-client... this often requires bouncing the WinRM service.
#
if opts[:windows]
  nodes = winnodws - ignore_nodes
  nodes.each do |node|
    cmd = "knife winrm -m #{node} \"chef-client -c c:/chef/client.rb\" -x #{opts[:winuser]}"
    cmd = "#{cmd} -P #{opts[:winpass]}" unless opts[:winpass].nil?
    puts cmd
    system(cmd)
  end
end

#
# Use Net::SSH::Multi to run chef-client over SSH in parallel.
#
nodes = sshnodes - ignore_nodes
puts ''
Net::SSH::Multi.start(:concurrent_connections => opts[:connections], :on_error => :warn) do |session|
  # define the servers we want to use
  nodes.each do |node|
    puts "Connecting to #{opts[:username]}@#{node}"
    session.use "#{opts[:username]}@#{node}"
  end

  # execute commands on all servers
  cmd = "sudo pkill -9 chef-client; sudo chef-client -l #{opts[:loglevel]}; ps auxw | grep chef"
  session.exec cmd

  # run the aggregated event loop
  session.loop
end