Fix aggregation of the Dst/Src Ip

inverse-inc · Sep 19, 2024 · 74b3f75 · 74b3f75
1 parent d9db155
commit 74b3f75
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/go/cron/aggregator_test.go b/go/cron/aggregator_test.go
@@ -73,6 +73,14 @@ func TestAggregator(t *testing.T) {
 		t.Fatalf("Not aggreated DestPort")
 	}
 
+	if ne[0].DestIp != netip.AddrFrom4([4]byte{1, 1, 1, 2}) {
+		t.Fatalf("Not aggreated DestIp")
+	}
+
+	if ne[0].SourceIp != netip.AddrFrom4([4]byte{1, 1, 1, 1}) {
+		t.Fatalf("Not aggreated SrcIp")
+	}
+
 	events = []*PfFlows{
 		{
 			Flows: &[]PfFlow{

diff --git a/go/cron/pfflow.go b/go/cron/pfflow.go
@@ -124,6 +124,22 @@ func (f *PfFlow) CalculatedDstPort() int {
 	return int(f.DstPort)
 }
 
+func (f *PfFlow) CalculatedSrcIp() netip.Addr {
+	if f.BiFlow == 2 {
+		return f.DstIp
+	}
+
+	return f.SrcIp
+}
+
+func (f *PfFlow) CalculatedDstIp() netip.Addr {
+	if f.BiFlow == 2 {
+		return f.SrcIp
+	}
+
+	return f.DstIp
+}
+
 func (f *PfFlow) ToNetworkEvent() *NetworkEvent {
 	if f.DstMac == "00:00:00:00:00:00" && f.SrcMac == "00:00:00:00:00:00" {
 		return nil
@@ -136,8 +152,8 @@ func (f *PfFlow) ToNetworkEvent() *NetworkEvent {
 
 	return &NetworkEvent{
 		EventType:           NetworkEventTypeSuccessful,
-		SourceIp:            f.SrcIp,
-		DestIp:              f.DstIp,
+		SourceIp:            f.CalculatedSrcIp(),
+		DestIp:              f.CalculatedDstIp(),
 		DestPort:            f.CalculatedDstPort(),
 		IpProtocol:          ipProto,
 		IpVersion:           IpVersionIpv4,