/
main.go
127 lines (107 loc) · 3.53 KB
/
main.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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
// Vautour - A distributed & extensible web hunter
// Copyright (C) 2019 Quentin Machu & Vautour contributors
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.
package main
import (
"flag"
"fmt"
"github.com/quentin-m/vautour/src/pkg/formatter"
_ "github.com/quentin-m/vautour/src/modules/pastebin"
"github.com/quentin-m/vautour/src/pkg/vautour"
"github.com/quentin-m/vautour/src/pkg/version"
log "github.com/sirupsen/logrus"
"gopkg.in/yaml.v2"
"io/ioutil"
"os"
"runtime/pprof"
"strings"
_ "github.com/quentin-m/vautour/src/modules/redis"
_ "github.com/quentin-m/vautour/src/modules/elasticsearch"
_ "github.com/quentin-m/vautour/src/modules/yara"
_ "github.com/quentin-m/vautour/src/modules/mailer"
)
func main() {
// Parse command-line arguments.
flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
flagConfigPath := flag.String("config", "./config/vautour.yaml", "Load configuration from the specified file.")
flagLogLevel := flag.String("log-level", "info", "Define the logging level.")
flagCPUProfilePath := flag.String("cpu-profile", "", "Write a CPU profile to the specified file before exiting.")
flagVersion := flag.Bool("version", false, "Display the version of Vautour.")
flag.Parse()
if *flagVersion {
fmt.Printf("Vautour (%s)\n", version.Version)
return
}
// Configure Vautour.
configureLogger(flagLogLevel)
// Load configuration.
config, err := loadConfig(*flagConfigPath)
if err != nil {
log.WithError(err).Fatal("failed to load configuration")
}
// Enable CPU Profiling if specified
if *flagCPUProfilePath != "" {
defer stopCPUProfiling(startCPUProfiling(*flagCPUProfilePath))
}
// Start Vautour.
vautour.Boot(config)
}
// Initialize logging system
func configureLogger(flagLogLevel *string) {
logLevel, err := log.ParseLevel(strings.ToUpper(*flagLogLevel))
if err != nil {
log.WithError(err).Error("failed to set logger parser level")
}
log.SetLevel(logLevel)
log.SetOutput(os.Stdout)
log.SetFormatter(&formatter.JSONExtendedFormatter{ShowLn: true})
}
func startCPUProfiling(path string) *os.File {
f, err := os.Create(path)
if err != nil {
log.WithError(err).Fatal("failed to create profile file")
}
err = pprof.StartCPUProfile(f)
if err != nil {
log.WithError(err).Fatal("failed to start CPU profiling")
}
log.Info("started CPU profiling")
return f
}
func stopCPUProfiling(f *os.File) {
pprof.StopCPUProfile()
f.Close()
log.Info("stopped CPU profiling")
}
type namespacedConfig struct {
Vautour vautour.Config
}
func loadConfig(cfgPath string) (vautour.Config, error) {
cfg := vautour.Config{}
// Load configuration file if specified.
if cfgPath == "" {
return cfg, nil
}
yamlConfig, err := ioutil.ReadFile(os.ExpandEnv(cfgPath))
if err != nil {
log.WithError(err).Fatal("could not load configuration file")
}
// Parse the configuration.
var nCfg namespacedConfig
if err := yaml.Unmarshal([]byte(yamlConfig), &nCfg); err != nil {
return cfg, err
}
return nCfg.Vautour, nil
}