From 9883cd346f694ad09aac839f9ddc4a25df0e0b0a Mon Sep 17 00:00:00 2001
From: Emiliano Ciavatta
Date: Fri, 3 Apr 2020 12:47:23 +0200
Subject: Add connection_handler and stream_handler

---
 pcap_importer.go | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

(limited to 'pcap_importer.go')

diff --git a/pcap_importer.go b/pcap_importer.go
index 9428b29..c6260de 100644
--- a/pcap_importer.go
+++ b/pcap_importer.go
@@ -24,7 +24,7 @@ const importedPcapsCollectionName = "imported_pcaps"
 
 
 type PcapImporter struct {
-	storage     *Storage
+	storage     Storage
 	streamPool  *tcpassembly.StreamPool
 	assemblers  []*tcpassembly.Assembler
 	sessions    map[string]context.CancelFunc
@@ -36,10 +36,11 @@ type PcapImporter struct {
 type flowCount [2]int
 
 
-func NewPcapImporter(storage *Storage, serverIp string) *PcapImporter {
+func NewPcapImporter(storage Storage, serverIp net.IP) *PcapImporter {
+	serverEndpoint := layers.NewIPEndpoint(serverIp)
 	streamFactory := &BiDirectionalStreamFactory{
 		storage: storage,
-		serverIp: serverIp,
+		serverIp: serverEndpoint,
 	}
 	streamPool := tcpassembly.NewStreamPool(streamFactory)
 
@@ -50,7 +51,7 @@ func NewPcapImporter(storage *Storage, serverIp string) *PcapImporter {
 		sessions:    make(map[string]context.CancelFunc),
 		mAssemblers: sync.Mutex{},
 		mSessions:   sync.Mutex{},
-		serverIp:    layers.NewIPEndpoint(net.ParseIP(serverIp)),
+		serverIp:    serverEndpoint,
 	}
 }
 
-- 
cgit v1.2.3-70-g09d2