#!/usr/bin/ruby.ruby3.4 

$:.unshift(File.join(File.dirname(__dir__), 'lib'))

require 'puppetdb/connection'
require 'optparse'
require 'net/http'
require 'net/https'
require 'ipaddr'

def get_puppet_config(setting_name, section = nil)
    cmd = %w[puppet config print]
    cmd << "--section=#{section}" if !section&.empty?
    cmd << setting_name
    `#{cmd*' '}`.chomp
end

def local?(ip_or_hostname)
    IPAddr.new(ip_or_hostname).loopback?
rescue
    ip_or_hostname.to_s.strip.downcase == 'localhost'
end

options = Struct.new(:host, :port, :ssl, :ca, :cert, :key, :facts, :print, :query, :timeout, keyword_init: true).new(
    host:  'localhost',
    port:  nil,
    ssl:   false,
    facts: nil,
    query: nil
)

ARGV.options do |o|
    o.on('-p', '--puppetdb=PUPPETDB',          "Host running PuppetDB (#{options.host})")
    o.on('-P', '--port=PORT',         Integer, "Port PuppetDB is running on (#{options.port})")
    o.on('-s', '--[no-]ssl',                   "Use SSL for PuppetDB (#{options.ssl})")
    o.on('-q', '--query=QUERY',                'Query String')
    o.on('-f', '--facts=FACT',          Array, 'Comma separated list of facts')
    o.on('-t', '--timeout=SECONDS',   Integer, 'PuppetDB read timeout')
    o.on(      '--print',                      'Only print query')
    o.on(      '--ca=CA',                      'CA certificate')
    o.on(      '--cert=CERT',                  'Client certificate')
    o.on(      '--key=KEY',                    'Client private key')
    o.on('-h', '--help',                       'Display this help and exit') { STDOUT << o; exit }
    o.on_tail <<~EOF

    If --port is not specified, the port defaults to 8080 if ssl == false, and to 8081 if ssl == true.

    If ssl == true, host refers to localhost (by hostname, IPv4, or IPv6) and the user is root,
    then --ca, --cert and --key are automatically provided with puppet’s host certificates under
    puppet’s ssldir.

    Example:
        #{$0} 'Package["mariadb"]'
    EOF
    o.parse!(into: options)
end

options.query ||= ARGV[0]

abort 'Please specify a query' unless options.query

options.port = options.ssl ? 8081 : 8080 if options.port.nil?
if options.ssl && local?(options.host) && Process.uid.zero?
    options.ca   = get_puppet_config(:localcacert) if options.ca.nil?
    options.cert = get_puppet_config(:hostcert)    if options.cert.nil?
    options.key  = get_puppet_config(:hostprivkey) if options.key.nil?
end

puppetdb          = PuppetDB::Connection.new(options.host, options.port)
http              = Net::HTTP.new(options.host, options.port)
http.use_ssl      = options.ssl
http.verify_mode  = OpenSSL::SSL::VERIFY_NONE
http.read_timeout = options.timeout if options.timeout
http.ca_file      = File.expand_path(options.ca) if options.ca
http.key          = OpenSSL::PKey::RSA.new(File.read(File.expand_path(options.key)), '') if options.key
http.cert         = OpenSSL::X509::Certificate.new(File.read(File.expand_path(options.cert))) if options.cert

if options.print
    query = puppetdb.parse_query(options.query)
    puts query.to_json
    puts query.to_yaml
elsif options.facts
    query = puppetdb.parse_query(options.query, :facts)
    facts = puppetdb.facts(options.facts, query, http)
    facts.each_value do |host|
        puts options.facts.map { |f| host[f] if host.include?(f) }*','
    end
else
    query   = puppetdb.parse_query(options.query)
    results = puppetdb.query(:nodes, query, http)
    hosts   = results.map { |host| host['name'] }
    hosts.each { |host| puts host }
end
