/
main.go
127 lines (106 loc) · 2.75 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
package main
import (
"flag"
"fmt"
"os"
"github.com/Yara-Rules/yara-endpoint/server/config"
"github.com/Yara-Rules/yara-endpoint/server/database"
"github.com/Yara-Rules/yara-endpoint/server/router"
log "github.com/sirupsen/logrus"
)
var (
// Version is the bot version
Version = "v0.0.0"
// BuildID is the build identifier
BuildID = ""
// BuildDate is the compilation date
BuildDate = ""
// Log stuff
// logLevel
logLevel = "info"
// logOutput
logOutput = "-"
// logFormat
logFormat = "json"
configFileName = "yara-server.ini"
// ShowVersion show up the version
ShowVersion = false
DB *database.DataStore
)
func printVerion() {
fmt.Fprintf(os.Stdout, "Yara-Endpoint %s\n", Version)
fmt.Fprintf(os.Stdout, "Build ID %s\n", BuildID)
fmt.Fprintf(os.Stdout, "Build on %s\n", BuildDate)
os.Exit(0)
}
func init() {
flag.StringVar(&configFileName, "configFile", configFileName, "Configuration file")
flag.StringVar(&logLevel, "logLevel", logLevel, "Log level")
flag.StringVar(&logOutput, "logOutput", logOutput, "Log output file")
flag.StringVar(&logFormat, "logFormat", logFormat, "Log output format")
flag.BoolVar(&ShowVersion, "version", ShowVersion, "Show version")
}
func main() {
// Parse command line flags
flag.Parse()
validateFlags()
setLog()
log.Infof("** Yara-Endpoint Server %s **", Version)
// Read config
config.LoadConfig(configFileName)
// Setting up TCP Server
log.Info("Starting TCP Server")
srv := NewServer()
srv.Connect(config.CFG)
defer srv.Close()
log.Infof("TCP Server started and listening on %s:%d", config.CFG.TCPServer.Addr, config.CFG.TCPServer.Port)
// Setting up WEB Server
log.Info("Starting WEB Server")
m := router.NewMacaron(Version)
router.RegisterRoutes(m)
// Forking WEB Server
go m.Run(config.CFG.Webserver.Host, config.CFG.Webserver.Port)
log.Infof("WEB Server started and listening on %s:%d", config.CFG.Webserver.Host, config.CFG.Webserver.Port)
DB = database.NewDataStore(config.CFG)
defer DB.Close()
// Start TCP main loop
log.Info("Waiting for connections...")
for {
conn, err := srv.Accept()
if err != nil {
continue
}
log.Debugf("New connection accepted from %s", conn.RemoteAddr())
go srv.HandleClient()
}
}
func validateFlags() {
if ShowVersion {
printVerion()
}
}
func setLog() {
if logFormat == "json" {
log.SetFormatter(&log.JSONFormatter{})
} else {
log.SetFormatter(&log.TextFormatter{})
}
if logOutput == "-" {
log.SetOutput(os.Stdout)
} else {
out, err := os.OpenFile(logOutput, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
if err != nil {
fmt.Println(err)
os.Exit(-1)
}
log.SetOutput(out)
}
switch logLevel {
case "warn":
log.SetLevel(log.WarnLevel)
case "debug":
log.SetLevel(log.DebugLevel)
case "info":
log.SetLevel(log.InfoLevel)
}
}