/
stdout.go
97 lines (83 loc) · 4.31 KB
/
stdout.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
package output
import (
"context"
"errors"
"fmt"
"os"
"github.com/mosajjal/dnsmonster/internal/util"
metrics "github.com/rcrowley/go-metrics"
log "github.com/sirupsen/logrus"
)
type stdoutConfig struct {
StdoutOutputType uint `long:"stdoutoutputtype" ini-name:"stdoutoutputtype" env:"DNSMONSTER_STDOUTOUTPUTTYPE" default:"0" description:"What should be written to stdout. options:\n;\t0: Disable Output\n;\t1: Enable Output without any filters\n;\t2: Enable Output and apply skipdomains logic\n;\t3: Enable Output and apply allowdomains logic\n;\t4: Enable Output and apply both skip and allow domains logic" choice:"0" choice:"1" choice:"2" choice:"3" choice:"4"`
StdoutOutputFormat string `long:"stdoutoutputformat" ini-name:"stdoutoutputformat" env:"DNSMONSTER_STDOUTOUTPUTFORMAT" default:"json" description:"Output format for stdout. options:json,csv, csv_no_header, gotemplate. note that the csv splits the datetime format into multiple fields" choice:"json" choice:"csv" choice:"csv_no_header" choice:"gotemplate"`
StdoutOutputGoTemplate string `long:"stdoutoutputgotemplate" ini-name:"stdoutoutputgotemplate" env:"DNSMONSTER_STDOUTOUTPUTGOTEMPLATE" default:"{{.}}" description:"Go Template to format the output as needed"`
StdoutOutputWorkerCount uint `long:"stdoutoutputworkercount" ini-name:"stdoutoutputworkercount" env:"DNSMONSTER_STDOUTOUTPUTWORKERCOUNT" default:"8" description:"Number of workers"`
outputChannel chan util.DNSResult
closeChannel chan bool
outputMarshaller util.OutputMarshaller
}
func init() {
c := stdoutConfig{}
if _, err := util.GlobalParser.AddGroup("stdout_output", "Stdout Output", &c); err != nil {
log.Fatalf("error adding output Module")
}
c.outputChannel = make(chan util.DNSResult, util.GeneralFlags.ResultChannelSize)
util.GlobalDispatchList = append(util.GlobalDispatchList, &c)
}
// initialize function should not block. otherwise the dispatcher will get stuck
func (stdConfig stdoutConfig) Initialize(ctx context.Context) error {
var err error
var header string
stdConfig.outputMarshaller, header, err = util.OutputFormatToMarshaller(stdConfig.StdoutOutputFormat, stdConfig.StdoutOutputGoTemplate)
if err != nil {
log.Warnf("Could not initialize output marshaller, removing output: %s", err)
return err
}
// print header to stdout
if header != "" {
fmt.Println(header)
}
if stdConfig.StdoutOutputType > 0 && stdConfig.StdoutOutputType < 5 {
log.Info("Creating Stdout Output Channel")
go stdConfig.Output(ctx)
} else {
// we will catch this error in the dispatch loop and remove any output from the registry if they don't have the correct output type
return errors.New("no output")
}
return err
}
func (stdConfig stdoutConfig) Close() {
// todo: implement this
<-stdConfig.closeChannel
}
func (stdConfig stdoutConfig) OutputChannel() chan util.DNSResult {
return stdConfig.outputChannel
}
func (stdConfig stdoutConfig) stdoutOutputWorker(ctx context.Context) {
stdoutSentToOutput := metrics.GetOrRegisterCounter("stdoutSentToOutput", metrics.DefaultRegistry)
stdoutSkipped := metrics.GetOrRegisterCounter("stdoutSkipped", metrics.DefaultRegistry)
for {
select {
case data := <-stdConfig.outputChannel:
for _, dnsQuery := range data.DNS.Question {
if util.CheckIfWeSkip(stdConfig.StdoutOutputType, dnsQuery.Name) {
stdoutSkipped.Inc(1)
continue
}
stdoutSentToOutput.Inc(1)
fmt.Fprint(os.Stdout, string(stdConfig.outputMarshaller.Marshal(data)))
fmt.Fprint(os.Stdout, "\n")
}
case <-ctx.Done():
log.Debug("exitting out of stdout output worker") //todo:remove
}
}
}
func (stdConfig stdoutConfig) Output(ctx context.Context) {
for i := 0; i < int(stdConfig.StdoutOutputWorkerCount); i++ { // todo: make this configurable
go stdConfig.stdoutOutputWorker(ctx)
}
}
// This will allow an instance to be spawned at import time
// var _ = stdoutConfig{}.initializeFlags()