-
Notifications
You must be signed in to change notification settings - Fork 415
/
virus_detection_statistics.go
107 lines (92 loc) · 2.94 KB
/
virus_detection_statistics.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
package virus_detection
import (
"context"
"time"
"github.com/bytedance/Elkeid/server/manager/infra"
"github.com/bytedance/Elkeid/server/manager/infra/ylog"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/bson/primitive"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
)
type VirusScanTaskStatistics struct {
RunningTaskNum int `json:"running_task_num" bson:"running_task_num"`
FininshedTaskNum int `json:"finished_task_num" bson:"finished_task_num"`
LastScanTime int64 `json:"last_scan_time" bson:"last_scan_time"`
}
type VirusTaskStatAggresData struct {
Status string `json:"status" bson:"status"`
Count int `json:"count" bson:"count"`
LastScanTime int64 `json:"last_scan_time" bson:"last_scan_time"`
}
func UpdateVirusTaskStatistics(c context.Context) {
var rsp VirusScanTaskStatistics
matchQuery := bson.D{primitive.E{
Key: "$match",
Value: bson.D{primitive.E{
Key: "action",
Value: bson.D{primitive.E{
Key: "$in",
Value: VirusTaskActionList,
}},
}},
}}
groupQuery := bson.D{primitive.E{
Key: "$group",
Value: bson.D{
primitive.E{Key: "_id", Value: "$task_status"},
primitive.E{Key: "count", Value: bson.D{primitive.E{Key: "$sum", Value: 1}}},
primitive.E{Key: "last_scan_time", Value: bson.D{primitive.E{Key: "$max", Value: "$create_time"}}},
},
}}
projectQuery := bson.D{primitive.E{
Key: "$project",
Value: bson.D{
primitive.E{Key: "status", Value: "$_id"},
primitive.E{Key: "count", Value: 1},
primitive.E{Key: "last_scan_time", Value: 1},
}}}
pipeline := mongo.Pipeline{
matchQuery,
groupQuery,
projectQuery,
}
collection := infra.MongoClient.Database(infra.MongoDatabase).Collection(infra.AgentTaskCollection)
opts := options.Aggregate().SetMaxTime(15 * time.Second)
cursor, err := collection.Aggregate(context.TODO(), pipeline, opts)
if err != nil {
ylog.Errorf("func UpdateVirusTaskStatistics Aggregate run error", err.Error())
return
}
var outList []VirusTaskStatAggresData
err = cursor.All(c, &outList)
if err != nil {
ylog.Errorf("func UpdateVirusTaskStatistics Aggregate decode error", err.Error())
return
}
rsp.LastScanTime = 0
for _, one := range outList {
switch one.Status {
case "running":
rsp.RunningTaskNum = rsp.RunningTaskNum + one.Count
case "created":
rsp.RunningTaskNum = rsp.RunningTaskNum + one.Count
case "finished":
rsp.FininshedTaskNum = rsp.FininshedTaskNum + one.Count
}
if one.LastScanTime > rsp.LastScanTime {
rsp.LastScanTime = one.LastScanTime
}
}
// update db
statCol := infra.MongoClient.Database(infra.MongoDatabase).Collection(infra.VirusDetectionTaskStatCollectionV1)
option := &options.UpdateOptions{}
option.SetUpsert(true)
setValue := bson.M{"$set": rsp}
_, err = statCol.UpdateOne(c, bson.M{}, setValue, option)
if err != nil {
ylog.Errorf("func UpdateVirusTaskStatistics update Statistics error", err.Error())
return
}
return
}