Configure whether messages should be uniq or not

PI/input.py

@@ -20,14 +20,15 @@ class Input:
 
     """Implementation of base input class"""
 
-    def __init__(self, filename, maxMessages):
+    def __init__(self, filename, maxMessages, onlyUniq):
         """Import specified filename"""
 
         self.set = set()
         self.sequences = []
         self.index = 0
         self.maxMessages = maxMessages
         self.readMessages = 0
+        self.onlyUniq = onlyUniq
 
     def __iter__(self):
         self.index = 0
@@ -54,8 +55,8 @@ class Pcap(Input):
 
     """Handle the pcap file format"""
 
-    def __init__(self, filename, maxMessages, offset=14):
-        Input.__init__(self, filename, maxMessages)
+    def __init__(self, filename, maxMessages, onlyUniq, offset=14):
+        Input.__init__(self, filename, maxMessages, onlyUniq)
         self.pktNumber = 0
         self.offset = offset
 
@@ -121,7 +122,7 @@ def handler(self, hdr, pkt):
         l = len(self.set)
         self.set.add(seq)
 
-        if len(self.set) == l:
+        if len(self.set) == l and self.onlyUniq:
             return
 
         self.readMessages += 1
@@ -137,8 +138,8 @@ class ASCII(Input):
 
     """Handle newline delimited ASCII input files"""
 
-    def __init__(self, filename, maxMessages):
-        Input.__init__(self, filename, maxMessages)
+    def __init__(self, filename, maxMessages, onlyUniq):
+        Input.__init__(self, filename, maxMessages, onlyUniq)
 
         fd = open(filename, "r")
 
@@ -158,7 +159,7 @@ def __init__(self, filename, maxMessages):
             l = len(self.set)
             self.set.add(line)
 
-            if len(self.set) == l:
+            if len(self.set) == l and self.onlyUniq:
                 continue
 
             self.readMessages += 1
@@ -203,7 +204,7 @@ def consumeMessageBlock(self, data, connectionID, messageNumber, contentLength):
 	        # only insert uniq messages into the list of seuqences. avoid duplicate
 	        l = len(self.set)
 	        self.set.add(seq)
-	        if len(self.set) == l:
+	        if len(self.set) == l and self.onlyUniq:
        		     continue
 
                 self.readMessages += 1
@@ -216,10 +217,10 @@ def consumeMessageBlock(self, data, connectionID, messageNumber, contentLength):
 		self.mNumber += 1
       		self.sequences.append((self.mNumber, digitSeq))
 
-    def __init__(self, filename, maxMessages, messageDelimiter, fieldDelimiter):
+    def __init__(self, filename, maxMessages, onlyUniq, messageDelimiter, fieldDelimiter):
     	self.messageDelimiter = messageDelimiter
 	self.fieldDelimiter = fieldDelimiter
-        Input.__init__(self, filename, maxMessages)
+        Input.__init__(self, filename, maxMessages, onlyUniq)
 
         self.blockseparator = "******************************************"
         self.mNumber = 0

exampleconfig.yml

@@ -1,4 +1,5 @@
 format: <bro|pcap|ascii>
+onlyUniqMessages: [ False | True ]
 weight: [0, 1.0]
 graph: [ False | True ]
 maxMessages: [0 - ...]

main.py

@@ -30,6 +30,7 @@ def main():
     analysis = None
     config = None
     gnuplotFile = None
+    onlyUniq = False
 
     #
     # Parse command line options and do sanity checking on arguments
@@ -60,6 +61,9 @@ def main():
     else:
         weight = config['weight']
 
+    if 'onlyUniqMessages' in config:
+        onlyUniq = config['onlyUniqMessages']
+
     if not 'format' in  config:
         print "FATAL: No input format configured!"
         sys.exit(-1)
@@ -105,11 +109,11 @@ def main():
     #
     try:
         if format == "pcap":
-            sequences = PI.input.Pcap(file, maxMessages)
+            sequences = PI.input.Pcap(file, maxMessages, onlyUniq)
         elif format == "ascii":
-            sequences = PI.input.ASCII(file, maxMessages)
+            sequences = PI.input.ASCII(file, maxMessages, onlyUniq)
         elif format == "bro":
-            sequences = PI.input.Bro(file, maxMessages, messageDelimiter, fieldDelimiter)
+            sequences = PI.input.Bro(file, maxMessages, onlyUniq, messageDelimiter, fieldDelimiter)
         else:
             print "FATAL: Specify file format"
             sys.exit(-1)