#!/opt/opscode/bin/knife exec
# This script is inspired by the below source:
# https://github.com/facebook/chef-utils/blob/master/chef-server-stats/chef-server-stats

require 'base64'
require 'date'
require 'time'
require 'digest/sha1'
require 'openssl'
require 'net/https'
require 'json'

def handle_error(error_message = "Error while collecting the chef data")
    abort('ERROR| ' + error_message.to_s)
end

begin
  if File.exists?('/opt/opscode/bin/chef-server-ctl')
    embedded_path = '/opt/opscode/embedded/bin'
    ctl_bin = '/opt/opscode/bin/chef-server-ctl'
  else
    handle_error 'Chef server type check failed, exiting'
  end


  cmd = "#{ctl_bin} service-list"
  service_list = Mixlib::ShellOut.new(cmd)
  s = service_list.run_command

  if s.exitstatus.zero?
    services = s.stdout.split("\n")
  else
    handle_error 'Failed to get service list, exiting'
  end

  services.map! { |service| service.chomp('*') }

rescue StandardError => msg
  handle_error msg
end

def get_api_counts(services)
  Chef::Config[:http_retry_count] = 0
  stats = {
    'server.nodes.count' => api.get('nodes').count.to_i,
    'server.cookbooks.count' => api.get('cookbooks').count.to_i,
    'server.roles.count' => api.get('roles').count.to_i,
    'server.environments.count' => api.get('environments').count.to_i,
  }
  return stats
end


def get_rabbitmq_stats(embedded_path)
  rabbit_bin = "#{embedded_path}/rabbitmqctl"
  return {} unless File.exists?(rabbit_bin)

  cmd = "PATH=\"#{embedded_path}/:$PATH\" #{rabbit_bin}" +
        ' list_queues -p /chef -n rabbit@localhost messages_ready'
  s = Mixlib::ShellOut.new(cmd).run_command
  if s.exitstatus.zero?
    lines = s.stdout.split(/\n/)
    sum = 0
    lines.each do |line|
      next unless /^\d+$/.match(line)
      sum += line.to_i
    end
    return { 'server.rabbitmq.messages.ready' => sum }
  else
    return {}
  end
end


def get_postgresql_stats(embedded_path)
  psql_bin = "#{embedded_path}/psql"
  return {} unless File.exists?(psql_bin)
  
  columns = %w{
    seq_scan
    seq_tup_read
    idx_scan
    idx_tup_fetch
    n_tup_ins
    n_tup_upd
    n_tup_del
    n_live_tup
    n_dead_tup
  }

  stats = {}
  q = "SELECT SUM(#{columns.join('), SUM(')}) FROM pg_stat_all_tables;"
  cmd = "su opscode-pgsql -c \"cd; #{psql_bin} -A -P tuples_only -U " +
        "opscode-pgsql -d opscode_chef -c '#{q}'\""

  s = Mixlib::ShellOut.new(cmd).run_command
  if s.exitstatus.zero?
    s.stdout.split('|').each do |value|
      stats["server.postgresql.#{columns.shift}"] = value.chomp.to_i
    end
  end

  q = "SELECT count(*) FROM pg_stat_activity WHERE datname = 'opscode_chef';"
  cmd = "su opscode-pgsql -c \"cd; #{psql_bin} -A -P tuples_only -U " +
        "opscode-pgsql -d opscode_chef -c \\\"#{q}\\\"\""

  s = Mixlib::ShellOut.new(cmd).run_command
  if s.exitstatus.zero?
    stats['server.postgresql.connection.count'] = s.stdout.chomp.to_i
  end

  return stats
end


def get_server_status
  uri = URI.parse('https://localhost/_status')
  http = Net::HTTP.new(uri.host, uri.port)
  http.use_ssl = true
  http.verify_mode = OpenSSL::SSL::VERIFY_NONE
  request = Net::HTTP::Get.new(uri.request_uri)
  response = http.request(request)
  status_hash = Chef::JSONCompat.from_json(response.body)
  if status_hash['status'] == 'pong'
    server_status = { 'server.status' => 1 }
  else
    server_status = { 'server.status' => 0 }
  end
  status = server_status

  if status.empty?
    status = { 'server.status' => 0 }
  end
  return status
end


def headers(opts={})
  path        = opts[:path]
  client_name = Chef::Config[:client_name]
  key_file    = Chef::Config[:client_key]
  body        = ""
  timestamp   = Time.now.utc.iso8601
  key         = OpenSSL::PKey::RSA.new(File.read(key_file))
  canonical   = "Method:GET\nHashed Path:#{encode(path)}\nX-Ops-Content-Hash:#{encode(body)}\nX-Ops-Timestamp:#{timestamp}\nX-Ops-UserId:#{client_name}"

  header_hash = {
    'Accept'             => 'application/json',
    'X-Ops-Sign'         => 'algorithm=sha1;version=1.0',
    'X-Ops-Userid'       => client_name,
    'X-Ops-Timestamp'    => timestamp,
    'X-Ops-Content-Hash' => encode(body),
    'X-Ops-Server-API-Info' => '1',
    'X-Ops-Reporting-Protocol-Version' => '0.1.0'
  }

  signature = Base64.encode64(key.private_encrypt(canonical))
  signature_lines = signature.split(/\n/)
  signature_lines.each_index do |idx|
    key = "X-Ops-Authorization-#{idx + 1}"
    header_hash[key] = signature_lines[idx]
  end
  
  header_hash
end


def encode(string)
  ::Base64.encode64(Digest::SHA1.digest(string)).chomp
end


def get_request(opts={})
  uri = URI.parse(opts[:path])
  
  if opts[:params] != nil
    params = opts[:params]
  else
    params = { }
  end
  
  uri.query = URI.encode_www_form(params)
  request  = Net::HTTP::Get.new(uri.to_s, headers(opts))

  http = Net::HTTP.new('localhost', opts.fetch(:port, 443))
  http.use_ssl     = true
  http.verify_mode = opts[:ssl_insecure] ? OpenSSL::SSL::VERIFY_NONE :  OpenSSL::SSL::VERIFY_PEER
  response = http.request(request)

  resp = if response.body != '' then JSON.parse(response.body) else {} end
  resp
end


def get_reports_status()
  result = get_request ({
      :path         => '/reports/status',
      :ssl_insecure => true
  })

  return_obj = {}
  return_obj["server.rest_api"] = if result["rest_api"]=='online' then 1 else 0 end
  return_obj["server.sql_db"] = if result["sql_db"]=='online' then 1 else 0 end
  return_obj["server.index"] = if result["index"]=='online' then 1 else 0 end
  
  return_obj
end


def get_reports_org_runs()
  time_now = DateTime.now.strftime('%s').to_i

  return_obj = {}
  
  params = { :until => time_now, :from => (time_now - (60*60)), :status => 'success', :granularity => 'hour', :rows => 9999 }
  result = get_request({
      :path         => Chef::Config[:chef_server_url] + '/reports/org/runs',
      :params       => params,
      :ssl_insecure => true
  })
  return_obj["server.run.success"] = if result['run_history'] != nil then result['run_history'].size else 0 end
  
  params = { :until => time_now, :from => (time_now - (60*60)), :status => 'failure', :granularity => 'hour', :rows => 9999 }
  result = get_request({
      :path         => Chef::Config[:chef_server_url] + '/reports/org/runs',
      :params       => params,
      :ssl_insecure => true
  })
  return_obj["server.run.failure"] = if result['run_history'] != nil then result['run_history'].size else 0 end

  params = { :until => time_now, :from => (time_now - (60*60)), :status => 'aborted', :granularity => 'hour', :rows => 9999 }
  result = get_request({
      :path         => Chef::Config[:chef_server_url] + '/reports/org/runs',
      :params       => params,
      :ssl_insecure => true
  })
  return_obj["server.run.aborted"] = if result['run_history'] != nil then result['run_history'].size else 0 end

  return return_obj
end

begin
  output = {}
  output.merge!(get_api_counts(services))
  output.merge!(get_rabbitmq_stats(embedded_path))
  output.merge!(get_postgresql_stats(embedded_path))
  output.merge!(get_server_status)
  output.merge!(get_reports_status())
  output.merge!(get_reports_org_runs())

  puts JSON.pretty_generate(output)

rescue  StandardError => msg
  handle_error msg
end

exit 0