mirror of https://github.com/0xERR0R/blocky.git
143 lines
3.2 KiB
Go
143 lines
3.2 KiB
Go
package resolver
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
|
|
"github.com/0xERR0R/blocky/stats"
|
|
"github.com/0xERR0R/blocky/util"
|
|
|
|
"github.com/jedib0t/go-pretty/table"
|
|
"github.com/miekg/dns"
|
|
)
|
|
|
|
// StatsResolver calculates query statistics
|
|
type StatsResolver struct {
|
|
NextResolver
|
|
recorders []*resolverStatRecorder
|
|
statsChan chan *statsEntry
|
|
}
|
|
|
|
type statsEntry struct {
|
|
request *Request
|
|
response *Response
|
|
}
|
|
|
|
type resolverStatRecorder struct {
|
|
aggregator *stats.Aggregator
|
|
fn func(*statsEntry) string
|
|
}
|
|
|
|
func newRecorder(name string, fn func(*statsEntry) string) *resolverStatRecorder {
|
|
return &resolverStatRecorder{
|
|
aggregator: stats.NewAggregator(name),
|
|
fn: fn,
|
|
}
|
|
}
|
|
|
|
func newRecorderWithMax(name string, max uint, fn func(*statsEntry) string) *resolverStatRecorder {
|
|
return &resolverStatRecorder{
|
|
aggregator: stats.NewAggregatorWithMax(name, max),
|
|
fn: fn,
|
|
}
|
|
}
|
|
|
|
func (r *StatsResolver) collectStats() {
|
|
for statsEntry := range r.statsChan {
|
|
for _, rec := range r.recorders {
|
|
rec.recordStats(statsEntry)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Resolve calculates query statistics
|
|
func (r *StatsResolver) Resolve(request *Request) (*Response, error) {
|
|
resp, err := r.next.Resolve(request)
|
|
|
|
if err == nil {
|
|
r.statsChan <- &statsEntry{
|
|
request: request,
|
|
response: resp,
|
|
}
|
|
}
|
|
|
|
return resp, err
|
|
}
|
|
|
|
// Configuration returns current configuraion
|
|
func (r *StatsResolver) Configuration() (result []string) {
|
|
result = append(result, "stats:")
|
|
for _, rec := range r.recorders {
|
|
result = append(result, fmt.Sprintf(" - %s", rec.aggregator.Name))
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
func (r *resolverStatRecorder) recordStats(e *statsEntry) {
|
|
r.aggregator.Put(r.fn(e))
|
|
}
|
|
|
|
// NewStatsResolver creates new instance of the resolver
|
|
func NewStatsResolver() ChainedResolver {
|
|
resolver := &StatsResolver{
|
|
statsChan: make(chan *statsEntry, 20),
|
|
recorders: createRecorders(),
|
|
}
|
|
|
|
go resolver.collectStats()
|
|
|
|
registerStatsTrigger(resolver)
|
|
|
|
return resolver
|
|
}
|
|
|
|
func (r *StatsResolver) printStats() {
|
|
logger := logger("stats_resolver")
|
|
|
|
w := logger.Writer()
|
|
defer w.Close()
|
|
|
|
logger.Info("******* STATS 24h *******")
|
|
|
|
for _, s := range r.recorders {
|
|
t := table.NewWriter()
|
|
t.SetOutputMirror(w)
|
|
t.SetTitle(s.aggregator.Name)
|
|
|
|
t.SetStyle(table.StyleLight)
|
|
|
|
util.IterateValueSorted(s.aggregator.AggregateResult(), func(k string, v int) {
|
|
t.AppendRow([]interface{}{fmt.Sprintf("%50s", k), v})
|
|
})
|
|
|
|
t.Render()
|
|
}
|
|
}
|
|
|
|
func createRecorders() []*resolverStatRecorder {
|
|
return []*resolverStatRecorder{
|
|
newRecorderWithMax("Top 20 queries", 20, func(e *statsEntry) string {
|
|
return util.ExtractDomain(e.request.Req.Question[0])
|
|
}),
|
|
newRecorderWithMax("Top 20 blocked queries", 20, func(e *statsEntry) string {
|
|
if e.response.RType == BLOCKED {
|
|
return util.ExtractDomain(e.request.Req.Question[0])
|
|
}
|
|
return ""
|
|
}),
|
|
newRecorder("Query count per client", func(e *statsEntry) string {
|
|
return strings.Join(e.request.ClientNames, ",")
|
|
}),
|
|
newRecorder("Reason", func(e *statsEntry) string {
|
|
return e.response.Reason
|
|
}),
|
|
newRecorder("Query type", func(e *statsEntry) string {
|
|
return dns.TypeToString[e.request.Req.Question[0].Qtype]
|
|
}),
|
|
newRecorder("Response type", func(e *statsEntry) string {
|
|
return dns.RcodeToString[e.response.Res.Rcode]
|
|
}),
|
|
}
|
|
}
|