# vim: set noet sw=2 ts=2 sts=2:
require 'prometheus/client'
require 'prometheus/client/formats/text'
require 'net/dns'
require 'ostruct'

class DnsblCollector
	class <<self
		attr_reader :blacklists, :resolver, :registry

		def load_config file = nil
			config = YAML.load_file file || './config.yml'
			self.blacklists = config['blacklists']
			self.resolver = config['resolver']
		end

		def blacklists= lists
			@blacklists =
				lists.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 resolver= resolver_cnf
			@resolver =
				case resolver_cnf
				when Net::DNS::Resolver then resolver_cnf
				when Hash then Net::DNS::Resolver.new resolver_cnf
				else
					raise ConfigError, "DnsblCollector.resolver= expects a Net::DNS::Resolver or a resolver-config as Hash"
				end
		end

		def registry= registry
			@registry = registry
			@registry.counter :dnsbl_query_error_count, docstring: 'Errors while query for blacklist', labels: %i[blacklist target]
			@registry.counter :dnsbl_query_blacklist_duration_sum, docstring: 'Total duration for querieng this blacklist', labels: %i[blacklist]
			@registry.counter :dnsbl_query_blacklist_count, docstring: 'Total count of queries this blacklist', labels: %i[blacklist]
		end
	end

	self.registry = Prometheus::Client::Registry.new

	attr_reader :lists, :registry, :configfile, :resolver, :metrics

	def initialize registry: nil
		@registry = registry || Prometheus::Client::Registry.new
		@registry.gauge :dnsbl_listed, docstring: 'IP listed currently on Blacklist', labels: %i[blacklist target]
		@metrics = OpenStruct.new @registry.instance_variable_get( :@metrics).merge( self.class.registry.instance_variable_get( :@metrics))
	end

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