#!/root/.rbenv/shims/ruby

# Install

# wget https://raw.githubusercontent.com/tomlobato/dns_one/master/util/dnsone_munin_ -O /usr/local/sbin/dnsone_munin_
# chmod 755 /usr/local/sbin/dnsone_munin_
# dnsone_munin_ install
# /etc/init.d/munin-node restart

require 'active_support/core_ext/hash/conversions'
require 'msgpack'
require 'ostruct'
require 'resolv'

class MuninDnsOne

    CACHE_FILE = '/tmp/dns_one_cache'
    CACHE_EXPIRE = 30

    GRAPHS = {
        rcode: {
            percent: true,
            vlabel: '%',
                fields: %w(not_imp refused yx_domain yxrr_set nxrr_set not_auth not_zone badvers no_error badsig badkey badtime badmode nx_domain badalg badname form_err serv_fail).sort
        },
        req_resource: {
            percent: true,
            vlabel: '%',
            fields: %w(a aaaa any cname hinfo minfo mx ns ptr soa txt wks)
        },
        cache: {
            percent: true,
            vlabel: '%',
            fields: {
                hit: 'Domain found in cache.',
                miss: 'Domain not found in cache, requiring a backend access.',
            },
        },
        requests: {
            vlabel: 'count',
            fields: {
                requests: 'requests',
            },
        }
    }

    def initialize argv
        @argv = argv
    end
        
    def cli 
        case @argv[0]
        when 'config'
            @param = get_params
            config
        when 'install'
            install
        else
            @param = get_params
            run
        end
    end

    private

    def graph
        GRAPHS[@param.graph]
    end

    def fields
        if graph[:fields].is_a? Hash
            graph[:fields]
        else
            _fields = {}            
            graph[:fields].each do |field|
                _fields[field.to_sym] = field
            end
            _fields
        end
    end

    def config
        puts <<-CONFIG
graph_category DnsOne
graph_title #{@param.graph}
graph_vlabel #{graph[:vlabel]}
graph_args --base 1000 -l 0

CONFIG

        fields.each_pair do |k, v|
            puts "#{k}.label #{k}"
        end

        exit 0
    end

    def run
        stats = get_stats[@param.graph.to_s]

        if graph[:percent]
            sum = stats.values.reduce(:+).to_f
        end 
        
        fields.keys.each do |k|
            value = stats[k] || 0.0
            if graph[:percent]
                value *= 100 / sum 
            end
            puts "#{k}.value #{ value }"
        end

        exit 0
    end

    # Util

    def get_params        
        filename = File.basename __FILE__

        if filename =~ /^dnsone_(#{ GRAPHS.keys.join '|' })/
            graph_key = $1
            OpenStruct.new(
                graph: graph_key.to_sym,
            )
        else
            error "Invalid graph type for filename #{filename}."
        end
    end

    def get_stats
        if false and File.exists?(CACHE_FILE) and
            (Time.now - File.stat(CACHE_FILE).ctime) < CACHE_EXPIRE
            MessagePack.unpack File.open(CACHE_FILE, 'rb').read
         else
            stat = fetch_stats
            File.open(CACHE_FILE, 'wb', 0600).write MessagePack.pack(stat)
            stat
        end
    end

    def fetch_stats
        stat = {}
        section = nil

        `dns_one stats`.each_line do |l|
            if l =~ /---\s*(\S+)\s*---/
                section = $1
            elsif l =~ /^(.*?)\s+(.*)$/
                key = $1.to_sym
                val = $2.strip.to_f
                stat[section] ||= {}
                stat[section][key] = val
            end
        end

        stat["requests"] = {
            requests: stat["cache"].values.reduce(:+).to_i
        }

        stat
    end

    def install
        # Links
        links = []
        GRAPHS.keys.each do |graph_key|
            links << graph_key
        end
        target = File.expand_path(__FILE__)
        links.each do |link|
            system "ln -s #{target} /etc/munin/plugins/dnsone_#{link}"
        end

        # munin_node_conf
        munin_node_conf = "[dnsone_*]\nuser root\n\n"
        File.open('/etc/munin/plugin-conf.d/munin-node', 'a').write munin_node_conf
    end

    def error msg
        STDERR.puts msg
        exit 1
    end
end

MuninDnsOne.new(ARGV).cli