Add rules statistics

4 files changed, 70 insertions, 39 deletions

diff --git a/connection_handler.go b/connection_handler.go
index 6b2b411..4e92ccf 100644
--- a/connection_handler.go
+++ b/connection_handler.go
@@ -256,16 +256,22 @@ func (ch *connectionHandlerImpl) UpdateStatistics(connection Connection) {
 	}
 	servicePort := connection.DestinationPort
 
+	updateDocument := UnorderedDocument{
+		fmt.Sprintf("connections_per_service.%d", servicePort):  1,
+		fmt.Sprintf("client_bytes_per_service.%d", servicePort): connection.ClientBytes,
+		fmt.Sprintf("server_bytes_per_service.%d", servicePort): connection.ServerBytes,
+		fmt.Sprintf("total_bytes_per_service.%d", servicePort):  connection.ClientBytes + connection.ServerBytes,
+		fmt.Sprintf("duration_per_service.%d", servicePort):     duration.Milliseconds(),
+	}
+
+	for _, ruleID := range connection.MatchedRules {
+		updateDocument[fmt.Sprintf("matched_rules.%s", ruleID.Hex())] = 1
+	}
+
 	var results interface{}
 	if _, err := ch.Storage().Update(Statistics).Upsert(&results).
-		Filter(OrderedDocument{{"_id", time.Unix(rangeStart*60, 0)}}).OneComplex(UnorderedDocument{
-		"$inc": UnorderedDocument{
-			fmt.Sprintf("connections_per_service.%d", servicePort):  1,
-			fmt.Sprintf("client_bytes_per_service.%d", servicePort): connection.ClientBytes,
-			fmt.Sprintf("server_bytes_per_service.%d", servicePort): connection.ServerBytes,
-			fmt.Sprintf("duration_per_service.%d", servicePort):     duration.Milliseconds(),
-		},
-	}); err != nil {
+		Filter(OrderedDocument{{"_id", time.Unix(rangeStart*60, 0)}}).
+		OneComplex(UnorderedDocument{"$inc": updateDocument}); err != nil {
 		log.WithError(err).WithField("connection", connection).Error("failed to update connection statistics")
 	}
 }
diff --git a/frontend/src/components/Timeline.js b/frontend/src/components/Timeline.js
index 615203f..6b8806f 100644
--- a/frontend/src/components/Timeline.js
+++ b/frontend/src/components/Timeline.js
@@ -102,23 +102,24 @@ class Timeline extends Component {
         ports.forEach(s => urlParams.append("ports", s));
 
         const metrics = (await backend.get("/api/statistics?" + urlParams)).json;
+        if (metrics.length === 0) {
+            return;
+        }
+
         const zeroFilledMetrics = [];
         const toTime = m => new Date(m["range_start"]).getTime();
-
-        if (metrics.length > 0) {
-            let i = 0;
-            for (let interval = toTime(metrics[0]); interval <= toTime(metrics[metrics.length - 1]); interval += minutes) {
-                if (interval === toTime(metrics[i])) {
-                    const m = metrics[i++];
-                    m["range_start"] = new Date(m["range_start"]);
-                    zeroFilledMetrics.push(m);
-                } else {
-                    const m = {};
-                    m["range_start"] = new Date(interval);
-                    m[metric] = {};
-                    ports.forEach(p => m[metric][p] = 0);
-                    zeroFilledMetrics.push(m);
-                }
+        let i = 0;
+        for (let interval = toTime(metrics[0]); interval <= toTime(metrics[metrics.length - 1]); interval += minutes) {
+            if (interval === toTime(metrics[i])) {
+                const m = metrics[i++];
+                m["range_start"] = new Date(m["range_start"]);
+                zeroFilledMetrics.push(m);
+            } else {
+                const m = {};
+                m["range_start"] = new Date(interval);
+                m[metric] = {};
+                ports.forEach(p => m[metric][p] = 0);
+                zeroFilledMetrics.push(m);
             }
         }
 
@@ -127,6 +128,7 @@ class Timeline extends Component {
             columns: ["time"].concat(ports),
             points: zeroFilledMetrics.map(m => [m["range_start"]].concat(ports.map(p => m[metric][p] || 0)))
         });
+
         const start = series.range().begin();
         const end = series.range().end();
         start.setTime(start.getTime() - minutes);
diff --git a/rules_manager.go b/rules_manager.go
index 636fc74..0e6c3d1 100644
--- a/rules_manager.go
+++ b/rules_manager.go
@@ -121,16 +121,24 @@ func LoadRulesManager(storage Storage, flagRegex string) (RulesManager, error) {
 
 	// if there are no rules in database (e.g. first run), set flagRegex as first rule
 	if len(rulesManager.rules) == 0 {
-		if _, err := rulesManager.AddRule(context.Background(), Rule{
-			Name:  "flag",
-			Color: "#E53935",
-			Notes: "Mark connections where the flag is stolen",
-			Patterns: []Pattern{
-				{Regex: flagRegex, Direction: DirectionToClient},
-			},
-		}); err != nil {
-			return nil, err
-		}
+		go func() {
+			_, _ = rulesManager.AddRule(context.Background(), Rule{
+				Name:  "flag_out",
+				Color: "#e53935",
+				Notes: "Mark connections where the flags are stolen",
+				Patterns: []Pattern{
+					{Regex: flagRegex, Direction: DirectionToClient, Flags: RegexFlags{Utf8Mode: true}},
+				},
+			})
+			_, _ = rulesManager.AddRule(context.Background(), Rule{
+				Name:  "flag_in",
+				Color: "#43A047",
+				Notes: "Mark connections where the flags are placed",
+				Patterns: []Pattern{
+					{Regex: flagRegex, Direction: DirectionToServer, Flags: RegexFlags{Utf8Mode: true}},
+				},
+			})
+		}()
 	} else {
 		if err := rulesManager.generateDatabase(rules[len(rules)-1].ID); err != nil {
 			return nil, err
diff --git a/statistics_controller.go b/statistics_controller.go
index 1714c0b..57c7d95 100644
--- a/statistics_controller.go
+++ b/statistics_controller.go
@@ -29,26 +29,29 @@ type StatisticRecord struct {
 	ConnectionsPerService map[uint16]int   `json:"connections_per_service,omitempty" bson:"connections_per_service"`
 	ClientBytesPerService map[uint16]int   `json:"client_bytes_per_service,omitempty" bson:"client_bytes_per_service"`
 	ServerBytesPerService map[uint16]int   `json:"server_bytes_per_service,omitempty" bson:"server_bytes_per_service"`
+	TotalBytesPerService  map[uint16]int   `json:"total_bytes_per_service,omitempty" bson:"total_bytes_per_service"`
 	DurationPerService    map[uint16]int64 `json:"duration_per_service,omitempty" bson:"duration_per_service"`
+	MatchedRules          map[RowID]int64  `json:"matched_rules,omitempty" bson:"matched_rules"`
 }
 
 type StatisticsFilter struct {
 	RangeFrom time.Time `form:"range_from"`
 	RangeTo   time.Time `form:"range_to"`
 	Ports     []uint16  `form:"ports"`
+	RulesIDs  []RowID   `form:"rules_ids"`
 	Metric    string    `form:"metric"`
 }
 
 type StatisticsController struct {
-	storage Storage
-	metrics []string
+	storage         Storage
+	servicesMetrics []string
 }
 
 func NewStatisticsController(storage Storage) StatisticsController {
 	return StatisticsController{
 		storage: storage,
-		metrics: []string{"connections_per_service", "client_bytes_per_service",
-			"server_bytes_per_service", "duration_per_service"},
+		servicesMetrics: []string{"connections_per_service", "client_bytes_per_service",
+			"server_bytes_per_service", "total_bytes_per_service", "duration_per_service"},
 	}
 }
 
@@ -62,19 +65,31 @@ func (sc *StatisticsController) GetStatistics(context context.Context, filter St
 		query = query.Filter(OrderedDocument{{"_id", UnorderedDocument{"$gt": filter.RangeTo}}})
 	}
 	for _, port := range filter.Ports {
-		for _, metric := range sc.metrics {
+		for _, metric := range sc.servicesMetrics {
 			if filter.Metric == "" || filter.Metric == metric {
 				query = query.Projection(OrderedDocument{{fmt.Sprintf("%s.%d", metric, port), 1}})
 			}
 		}
+
 	}
 	if filter.Metric != "" && len(filter.Ports) == 0 {
-		for _, metric := range sc.metrics {
+		for _, metric := range sc.servicesMetrics {
 			if filter.Metric == metric {
 				query = query.Projection(OrderedDocument{{metric, 1}})
 			}
 		}
 	}
+	for _, ruleID := range filter.RulesIDs {
+		if filter.Metric == "" || filter.Metric == "matched_rules" {
+			query = query.Projection(OrderedDocument{{fmt.Sprintf("matched_rules.%s", ruleID.Hex()), 1}})
+		}
+
+	}
+	if filter.Metric != "" && len(filter.RulesIDs) == 0 {
+		if filter.Metric == "matched_rules" {
+			query = query.Projection(OrderedDocument{{"matched_rules", 1}})
+		}
+	}
 
 	if err := query.All(&statisticRecords); err != nil {
 		log.WithError(err).WithField("filter", filter).Error("failed to retrieve statistics")