1
1
package engine
2
2
3
3
import (
4
+ "bytes"
4
5
"context"
6
+ "fmt"
5
7
"github.com/hktalent/51pwnPlatform/lib"
6
8
"github.com/hktalent/51pwnPlatform/pkg/models"
7
9
"github.com/hktalent/ProScan4all/lib/util"
8
10
"github.com/hktalent/ProScan4all/pocs_go"
9
11
"github.com/hktalent/jaeles/cmd"
12
+ jsoniter "github.com/json-iterator/go"
10
13
"github.com/panjf2000/ants/v2"
14
+ "github.com/ulule/deepcopier"
15
+ "io/ioutil"
11
16
"log"
17
+ "net"
18
+ "net/url"
12
19
"os"
13
20
"os/signal"
21
+ "strconv"
22
+ "strings"
14
23
"sync"
24
+ "time"
15
25
)
16
26
27
+ var json = jsoniter .ConfigCompatibleWithStandardLibrary
28
+
17
29
// 引擎对象,全局单实例
18
30
type Engine struct {
19
31
Context * context.Context // 上下文
20
32
Wg * sync.WaitGroup // Wg
21
33
Pool int // 线程池
22
34
PoolFunc * ants.PoolWithFunc // 线程调用
23
35
EventData chan * models.EventData // 数据队列
36
+ NodeId string `json:"node_id"` // 分布式引擎节点的id,除非系统更换,docker重制,否则始终一致
37
+ LimitTask int `json:"limit_task"` // 当前节点任务并发数的限制
38
+ SyTask int `json:"sy_task"` // 剩余task
39
+ DtServer string `json:"dt_server"` // 获取任务、提交任务状态的server
24
40
caseScanFunc sync.Map
25
41
}
26
42
27
- var G_Engine * Engine
43
+ var GEngine * Engine
44
+
45
+ // 获取分布式任务
46
+ // /api/v1.0/syncResult/task/
28
47
29
48
// 创建引擎
30
49
// 默认每个 goroutine 占用 8KB 内存
@@ -34,7 +53,17 @@ func NewEngine(c *context.Context, pool int) *Engine {
34
53
if nil != util .G_Engine {
35
54
return util .G_Engine .(* Engine )
36
55
}
37
- x1 := & Engine {Context : c , Wg : & sync.WaitGroup {}, Pool : pool , EventData : make (chan * models.EventData , pool )}
56
+
57
+ x1 := & Engine {
58
+ Context : c ,
59
+ Wg : & sync.WaitGroup {},
60
+ Pool : pool ,
61
+ DtServer : util .GetVal ("DtServer" ),
62
+ EventData : make (chan * models.EventData , pool ),
63
+ LimitTask : util .GetValAsInt ("LimitTask" , 4 ),
64
+ }
65
+ x1 .SyTask = x1 .LimitTask // 初始化剩余任务等于最大任务数
66
+ x1 .initNodeId ()
38
67
p , err := ants .NewPoolWithFunc (pool , func (i interface {}) {
39
68
defer x1 .Wg .Done ()
40
69
x1 .DoEvent (i .(* models.EventData ))
@@ -44,13 +73,115 @@ func NewEngine(c *context.Context, pool int) *Engine {
44
73
}
45
74
x1 .PoolFunc = p
46
75
util .G_Engine = x1
47
- G_Engine = x1
76
+ GEngine = x1
48
77
util .EngineFuncFactory = x1 .EngineFuncFactory
49
78
util .SendEvent = x1 .SendEvent
50
79
log .Println ("Engine init ok" )
51
80
return x1
52
81
}
53
82
83
+ func (e * Engine ) initNodeId () {
84
+ dirname , err := os .Getwd ()
85
+ szP := dirname + "/.DistributedId"
86
+ if nil == err {
87
+ if util .FileExists (szP ) {
88
+ data , err := ioutil .ReadFile (szP )
89
+ if nil == err {
90
+ e .NodeId = strings .TrimSpace (string (data ))
91
+ }
92
+ }
93
+ }
94
+ if "" == e .NodeId {
95
+ e .NodeId = util .GenUuid ()
96
+ ioutil .WriteFile (szP , []byte (e .NodeId ), os .ModePerm )
97
+ }
98
+ }
99
+
100
+ // "https://dt.51pwn.com/api/v1.0/syncResult/task/%d"
101
+ // curl -v -XPOST -d '{"Num":22,"task_ids":"","node_id":"xx","task_num":443}' https://127.0.0.1:8081/api/v1.0/syncResult/task/33
102
+ // 结果反馈 /api/v1.0/syncResult/task/%d
103
+ // 获取、确认分布式任务,Distributed Tasks
104
+ func (e * Engine ) GetTask (okTaskIds string ) {
105
+ if resp , err := util .DoPost (fmt .Sprintf (e .DtServer , e .LimitTask ), map [string ]string {
106
+ "Content-Type" : "application/json" ,
107
+ }, strings .NewReader (`{"Num":` + strconv .Itoa (e .SyTask )+ `,"task_ids":"` + okTaskIds + `","node_id":"` + e .NodeId + `","task_num":` + strconv .Itoa (e .LimitTask )+ `}` )); nil == err && nil != resp {
108
+ defer resp .Body .Close ()
109
+ var n1 = models.EventData {}
110
+ if data , err := ioutil .ReadAll (resp .Body ); nil == err {
111
+ if err := json .Unmarshal (data , & n1 ); nil == err {
112
+ e .SendEvent (& n1 , n1 .EventType )
113
+ }
114
+ }
115
+ }
116
+ }
117
+
118
+ // 获取公共ip
119
+ func (r * Engine ) GetPublicIP () ([]string , error ) {
120
+ ifas , err := net .Interfaces ()
121
+ if err != nil {
122
+ return nil , err
123
+ }
124
+ var as []string
125
+
126
+ for _ , ifa := range ifas {
127
+ a := ifa .HardwareAddr .String ()
128
+ // have mac
129
+ if a != "" {
130
+ addrs , err := ifa .Addrs ()
131
+ // get Ip error
132
+ if nil != err {
133
+ continue
134
+ }
135
+ for _ , addr := range addrs {
136
+ switch v := addr .(type ) {
137
+ case * net.IPNet :
138
+ if v .IP .IsPrivate () {
139
+ continue
140
+ }
141
+ as = append (as , v .IP .String ())
142
+ case * net.IPAddr :
143
+ if v .IP .IsPrivate () {
144
+ continue
145
+ }
146
+ as = append (as , v .IP .String ())
147
+ }
148
+ }
149
+ }
150
+ }
151
+ return as , nil
152
+ }
153
+ func (e * Engine ) generateTaskId (s string ) string {
154
+ return util .GetSha1 (s )
155
+ }
156
+
157
+ // 发送任务
158
+ // 只发送非私有网络的任务
159
+ func (e * Engine ) SendTask (s string ) {
160
+ szUrl := fmt .Sprintf (e .DtServer , e .LimitTask )
161
+ if oU , err := url .Parse (szUrl ); nil == err {
162
+ szUrl = strings .Join ([]string {oU .Scheme , "://" , oU .Host , "/api/v1.0/alipay_task" }, "" )
163
+ szSendData := ""
164
+ sW := util .Base64Encode (s )
165
+ szTaskId := e .generateTaskId (s )
166
+ szSendData = "task_id=" + szTaskId + "&" + "scan_web=" + sW
167
+ base64Str := util .GetSig (szSendData , prvKey )
168
+ m1 := map [string ]string {"task_id" : szTaskId , "op" : "0" , "data_sign" : base64Str }
169
+ data , _ := json .Marshal (& m1 )
170
+
171
+ if resp , err := util .DoPost (fmt .Sprintf (e .DtServer , e .LimitTask ), map [string ]string {
172
+ "Content-Type" : "application/json" ,
173
+ }, bytes .NewReader (data )); nil == err && nil != resp {
174
+ defer resp .Body .Close ()
175
+ var n1 = models.EventData {}
176
+ if data , err := ioutil .ReadAll (resp .Body ); nil == err {
177
+ if err := json .Unmarshal (data , & n1 ); nil == err {
178
+ e .SendEvent (& n1 , n1 .EventType )
179
+ }
180
+ }
181
+ }
182
+ }
183
+ }
184
+
54
185
func (e * Engine ) EngineFuncFactory (nT int64 , fnCbk interface {}) {
55
186
e .RegCaseScanFunc (nT , fnCbk )
56
187
}
@@ -83,7 +214,7 @@ func (e *Engine) DoCase(ed *models.EventData) util.EngineFuncType {
83
214
func (e * Engine ) SendEvent (evt * models.EventData , argsTypes ... int64 ) {
84
215
for _ , i := range argsTypes {
85
216
var n1 = models.EventData {}
86
- util . DeepCopy (evt , & n1 )
217
+ deepcopier . Copy (evt ). To ( n1 )
87
218
n1 .EventType = i
88
219
e .EventData <- & n1
89
220
}
@@ -114,6 +245,9 @@ func (x1 *Engine) Running() {
114
245
signal .Notify (c , os .Interrupt )
115
246
//nMax := 120 // 等xxx秒都没有消息进入就退出
116
247
//nCnt := 0
248
+ // 每10秒获取一次任务
249
+ c1Task := time .NewTicker (5 * time .Second )
250
+ c2Task := time .NewTicker (15 * time .Second )
117
251
for {
118
252
select {
119
253
case <- util .Ctx_global .Done ():
@@ -146,8 +280,11 @@ func (x1 *Engine) Running() {
146
280
})
147
281
}(x1 )
148
282
}
149
- default :
283
+ case <- c1Task .C :
284
+ x1 .GetTask ("" )
285
+ case <- c2Task .C :
150
286
util .DoDelayClear (x1 .Wg ) // panic: sync: WaitGroup misuse: Add called concurrently with Wait
287
+ default :
151
288
//util.DoSleep()
152
289
}
153
290
}
@@ -168,3 +305,35 @@ func init() {
168
305
})
169
306
})
170
307
}
308
+
309
+ // 发送方 的签名key
310
+ var prvKey = []byte (`-----BEGIN RSA PRIVATE KEY-----
311
+ Proc-Type: 4,ENCRYPTED
312
+ DEK-Info: DES-EDE3-CBC,0B3A74436D1F0AAE
313
+
314
+ hzXeizI3DX5udIFmtfBpIEbQYz4ConOmdD/Vel2ppj6EG8PLI3oirlH7eRKxsCtU
315
+ khRVyYhUhb9II1jKF0tu7glnZabHnTGFAeo5nEXjTl/dp5Of6eJIaNcWlc7nXTft
316
+ koodNonRFZtGe7cHI5+WM4AvjWEXztuyPKCa0Zepz7k77IdxxQ+gzIonbni0OGeh
317
+ ze2kkBZMgCnS2LNqhk5zAhb0ATCTMLfd1FbsJXieXQAyCZQBER6J8hXvvZP2oZyq
318
+ izou+5BsT9/W/4crq0glYSc2SMSWUFj5sSSPgVj7dV8KiRueD1ybm13B0N5XKEoY
319
+ YK+IfdVBu6NspMfW8b+mp4JoAChA360d12Zyrg4J9gtOvoR4eOhtm4c0D0GPsubS
320
+ XjK8Jvp41QtLRz9trNshrXI8/3z5bc26zHBLbQ6lRJSA9Q5Guc70/8FHxPOik+SZ
321
+ 57gMsG0OuxvUfoIif5dwwtYh5dWYktE+Ii/FnFH/X3wROiq+D4ZWI0dKNED7fFry
322
+ RYmLJK+Bn7BjbzC+ZWwmKgMpmZyKF1/AB7031rB77z5Zq4Ksk+F5UEGLA4287CRT
323
+ 6vYY7eKpkRnZ2QHI5fdQ+fZ/A40n3NO1letf2MXB6Fxcz1P2DMYGJVVHfNSurj32
324
+ F2fUckcHe9Kvy5FCXwui7aXZUkhbREAnAiKFHeRwlyjYUwxeo9QZUHt06hUKaY4c
325
+ HJOLBzpjErtRfYLtGLADzaKAPe+fV+FqBWquoOG3/3aoz9oiiyxIj2a9D9ORLsy+
326
+ e3QgdgiAluQ2QMqdNeYO7POWXjasaqZ8XVanSCHn3Tw5GdEq6naWz4cGxaJiXHV1
327
+ PTiH2g/KEgu+L6b2xnwvEmpOKD33DEkB5xlnqUUFzAksbpL8l/sk9LPRjbjHl2Nu
328
+ yL9myaJgpbPhw12Ika97VJp6ooH7Qy2WRGJ67FXGBkXXpTcItQzqqs6ZIFdwadq6
329
+ Z1jNv+Wiq/o8IuVZys0a/LJlYYKrnHvKVl6LQrmcd+SWUgwbbKUBytsMDPB6OwHB
330
+ Qrf3flcVIYvgqS+R1745JeFK/kxI1vtYlyNlveiAi9yRtDVnw+0DWrY03kWNsGfg
331
+ 9swsfO+/nHUxf81hC2g/Carrkdz7BLrsMKZnHlNVVFTNsHeELpmlGIO4VxOJzjel
332
+ nuz4sITlXjBPsZernQuIbJ7GYqDv8Zb/dsW47BqIcl0PQ5FLOjJYBcIjpOnU1tLn
333
+ e6pShBS4KWK/YegJdo+SxDvqLl66fdn58s1TlaZfgQic6P/mSzHgBYImb7rIlrUk
334
+ aEWbEs4rAi0i8cwlg313ASK35E5enKM0C9uPaqnmFUQlT8X9SD+ELB7qHGRaXjcr
335
+ rd1HBuFu2bxJm6Tcfmy4bf+6QYW5czg1mJpGjvM9zCVHDtBxVj71XvVM6MLIruYn
336
+ zvNHq6ia8y1XUfkCxE5pzb0ap0LSS2XIEZxdRUCapGLAg4GNiA3Zkq4aDt8s7rGJ
337
+ fARzsx7PrOF3TgCxF97GZhRU6chMK8YAChRfwqsg0Mpw2plqiYa9v99KrRwPdzJo
338
+ 7J8M8tAQhZB8YzG0U4Dsvb6odc8OYAFJTPpFvNjyQGgcjWudp6vo0YbK54z/z5s4
339
+ -----END RSA PRIVATE KEY-----` )
0 commit comments