1require 'chef'
 2require 'chef-config/workstation_config_loader'
 3
 4class ChefSearch
 5  attr_reader :chef_environment
 6
 7  def initialize(chef_environment = fetch(:chef_env_name))
 8    @chef_environment = chef_environment
 9    load_configuration
10  end
11
12  def find_by(role, &filter)
13    hosts, _, count = search(:node, query_for(role), filter_result: filters)
14    hostnames = hostnames_from(hosts, &filter)
15    puts "Using #{hostnames.count}/#{count} hosts for role: #{role}"
16    hostnames.each { |x| puts "  #{x}" }
17    hostnames
18  end
19
20  private
21
22  def filters
23    { 'fqdn' => ['fqdn'], 'ip' => [ 'ec2', 'public_ipv4' ] }
24  end
25
26  def query_for(role)
27    [
28      "chef_environment:#{chef_environment}",
29      "roles:#{role}",
30    ].join(" AND ")
31  end
32
33  def search(*params)
34    Chef::Search::Query.new.search(*params)
35  end
36
37  def hostnames_from(hosts, &filter)
38    hostnames = hosts.map { |x| x['ip'] }.compact.sort
39    if filter.nil?
40      hostnames
41    else
42      hostnames = hostnames.find_all do |hostname|
43        yield hostname
44      end
45    end
46  end
47
48  def load_configuration(system_knife_rb = '/etc/chef/client.pem')
49    if File.readable?(system_knife_rb)
50      puts "Loading chef config from #{system_knife_rb}"
51      Chef::Config.from_file(system_knife_rb)
52    else
53      @loader = ChefConfig::WorkstationConfigLoader.new(nil, nil)
54      puts "Loading chef config from #{@loader.config_location}"
55      @loader.load
56    end
57  end
58end