forked from jaeles-project/jaeles
/
server.go
124 lines (114 loc) · 3.19 KB
/
server.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
package cmd
import (
"fmt"
"github.com/panjf2000/ants"
"os"
"path"
"path/filepath"
"sync"
"github.com/hktalent/jaeles/core"
"github.com/hktalent/jaeles/database"
"github.com/hktalent/jaeles/libs"
"github.com/hktalent/jaeles/server"
"github.com/hktalent/jaeles/utils"
"github.com/spf13/cobra"
)
func init() {
var serverCmd = &cobra.Command{
Use: "server",
Short: "Start API server",
Long: libs.Banner(), RunE: runServer,
}
serverCmd.Flags().String("host", "127.0.0.1", "IP address to bind the server")
serverCmd.Flags().String("port", "5000", "Port")
serverCmd.Flags().BoolP("no-auth", "A", false, "Turn off authenticated on API server")
serverCmd.SetHelpFunc(ServerHelp)
RootCmd.AddCommand(serverCmd)
}
func runServer(cmd *cobra.Command, _ []string) error {
if Options.NoDB {
fmt.Fprintf(os.Stderr, "Can't run Jaeles Server without DB\n")
os.Exit(-1)
}
SelectSign()
// prepare DB stuff
if Options.Server.Username != "" {
database.CreateUser(Options.Server.Username, Options.Server.Password)
}
// reload signature
SignFolder, _ := filepath.Abs(path.Join(Options.RootFolder, "base-signatures"))
allSigns := utils.GetFileNames(SignFolder, ".yaml")
if allSigns != nil {
for _, signFile := range allSigns {
database.ImportSign(signFile)
}
}
database.InitConfigSign()
var wg sync.WaitGroup
p, _ := ants.NewPoolWithFunc(Options.Concurrency, func(i interface{}) {
CreateRunner(i)
wg.Done()
}, ants.WithPreAlloc(true))
defer p.Release()
result := make(chan libs.Record)
go func() {
for {
record := <-result
utils.InforF("[Receive] %v %v \n", record.OriginReq.Method, record.OriginReq.URL)
for _, signFile := range Options.SelectedSigns {
sign, err := core.ParseSign(signFile)
if err != nil {
utils.ErrorF("Error loading sign: %v\n", signFile)
continue
}
// filter signature by level
if sign.Level > Options.Level {
continue
}
// parse sign as list or single
var url string
if sign.Type != "fuzz" {
url = record.OriginReq.URL
} else {
fuzzSign := sign
fuzzSign.Requests = []libs.Request{}
for _, req := range sign.Requests {
core.ParseRequestFromServer(&record, req, sign)
// override the original if these field defined in signature
if req.Method == "" {
req.Method = record.OriginReq.Method
}
if req.URL == "" {
req.URL = record.OriginReq.URL
}
if len(req.Headers) == 0 {
req.Headers = record.OriginReq.Headers
}
if req.Body == "" {
req.Body = record.OriginReq.Body
}
fuzzSign.Requests = append(fuzzSign.Requests, req)
}
url = record.OriginReq.URL
sign = fuzzSign
}
// single routine
wg.Add(1)
job := libs.Job{URL: url, Sign: sign}
_ = p.Invoke(job)
}
}
}()
host, _ := cmd.Flags().GetString("host")
port, _ := cmd.Flags().GetString("port")
Options.Server.NoAuth, _ = cmd.Flags().GetBool("no-auth")
bind := fmt.Sprintf("%v:%v", host, port)
Options.Server.Bind = bind
utils.GoodF("Start API server at %v", fmt.Sprintf("http://%v/", bind))
server.InitRouter(Options, result)
wg.Wait()
if utils.DirLength(Options.Output) == 0 {
os.RemoveAll(Options.Output)
}
return nil
}