require 'prometheus/client'
require 'prometheus/client/formats/text'
require 'net/dns'
require 'yaml'

class	DnsblCollector
	attr_reader :lists, :prometheus, :configfile, :resolver, :metrics

	def initialize prometheus: nil, configfile: nil
		@configfile = configfile || 'config.yml'
		@prometheus = prometheus || Prometheus::Client.registry
		%i[dnsbl_listed dnsbl_query_error].each {|m| @prometheus.unregister m }
		@metrics = {
			dnsbl_listed: @prometheus.gauge( :dnsbl_listed, 'IP listed currently on Blacklist'),
			dnsbl_query_error: @prometheus.counter( :dnsbl_query_error, 'Errors while query for blacklist'),
		}
		load_config
	end

	def load_config
		config = YAML.load_file @configfile
		@resolver = config['resolver']
		@lists =
			config['blacklists'].map do |bl|
				case bl
				when String then {blacklist: bl}
				when Hash then bl
				else raise ConfigError, "Unexpected Element in blacklists: #{bl}"
				end
			end
	end

	def collect ip
		prefix = ip.send :_reverse
		ip = ip.to_s
		dns = Net::DNS::Resolver.new @resolver
		todo = @lists.sort{|_|rand}
		durs = {}
		10.times.map do
			Thread.new do
				while bl = @lists.pop
					b = Time.now
					begin
						r = dns.search( "#{prefix}.#{bl[:blacklist]}", Net::DNS::A).answer.empty?
						@metrics[:dnsbl_listed].set( {blacklist: bl[:blacklist], target: ip}, r ? 0 : 1)
					rescue Net::DNS::Resolver::NoResponseError
						@metrics[:dnsbl_query_error].increment( {blacklist: bl[:blacklist], target: ip})
					end
					durs[bl[:blacklist]] = Time.now-b
				end
			end
		end.each &:join
	end
end