Fixed and moved out MessageType and Message to a separate file

code/Attack/MembersMgmtCommAttack.py

@@ -8,71 +8,13 @@ import os
 import sys
 
 import ID2TLib.Botnet.libbotnetcomm as lb
+import ID2TLib.Botnet.Message as Bmsg
 from Attack import BaseAttack
 from Attack.AttackParameters import Parameter as Param
 from Attack.AttackParameters import ParameterTypes
 from ID2TLib.Ports import PortSelectors
 import ID2TLib.Utility as Util
 
-class MessageType(Enum):
-    """
-    Defines possible botnet message types
-    """
-
-    TIMEOUT = 3
-    SALITY_NL_REQUEST = 101
-    SALITY_NL_REPLY = 102
-    SALITY_HELLO = 103
-    SALITY_HELLO_REPLY = 104
-
-    def is_request(mtype):
-        """
-        Checks whether the given message type is a request or not.
-        :param mtype: the message type to check
-        :return: True if it is a request, False otherwise
-        """
-        return mtype in {MessageType.SALITY_HELLO, MessageType.SALITY_NL_REQUEST}
-
-    def is_response(mtype):
-        """
-        Checks whether the given message type is a response or not.
-        :param mtype: the message type to check
-        :return: True if it is a response, False otherwise
-        """
-        return mtype in {MessageType.SALITY_HELLO_REPLY, MessageType.SALITY_NL_REPLY}
-
-class Message():
-    INVALID_LINENO = -1
-
-    """
-    Defines a compact message type that contains all necessary information.
-    """
-    def __init__(self, msg_id: int, src, dst, type_: MessageType, time: float, refer_msg_id: int=-1, line_no = -1):
-        """
-        Constructs a message with the given parameters.
-
-        :param msg_id: the ID of the message
-        :param src: something identifiying the source, e.g. ID or configuration
-        :param dst: something identifiying the destination, e.g. ID or configuration
-        :param type_: the type of the message
-        :param time: the timestamp of the message
-        :param refer_msg_id: the ID this message is a request for or reply to. -1 if there is no related message.
-        :param line_no: The line number this message appeared at in the original CSV file
-        """
-        self.msg_id = msg_id
-        self.src = src
-        self.dst = dst
-        self.type = type_
-        self.time = time
-        self.csv_time = time
-        self.refer_msg_id = refer_msg_id
-        self.line_no = line_no
-
-    def __str__(self):
-        str_ = "{0}. at {1}: {2}-->{3}, {4}, refer:{5} (line {6})".format(self.msg_id, self.time, self.src, self.dst, self.type, self.refer_msg_id, self.line_no)
-        return str_
-
-
 from ID2TLib import FileUtils, Generator
 from ID2TLib.IPv4 import IPAddress
 from ID2TLib.PcapAddressOperations import PcapAddressOperations
@@ -139,7 +81,7 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
 
         # create dict with MessageType values for fast name lookup
         self.msg_types = {}
-        for msg_type in MessageType:
+        for msg_type in Bmsg.MessageType:
             self.msg_types[msg_type.value] = msg_type
 
     def init_params(self):
@@ -233,7 +175,7 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
 
             # if the type of the message is a NL reply, determine the number of entries
             nl_size = 0
-            if msg.type == MessageType.SALITY_NL_REPLY:
+            if msg.type == Bmsg.MessageType.SALITY_NL_REPLY:
                 nl_size = randint(1, 25)    # what is max NL entries?
 
             # create suitable IP/UDP packet and add to packets list

code/ID2TLib/Botnet/CommunicationProcessor.py

@@ -1,7 +1,7 @@
 from lea import Lea
 from random import randrange
-from Attack.MembersMgmtCommAttack import MessageType
-from Attack.MembersMgmtCommAttack import Message
+from ID2TLib.Botnet.Message import Message
+from ID2TLib.Botnet.Message import MessageType
 
 
 class CommunicationProcessor:

code/ID2TLib/Botnet/Message.py

@@ -0,0 +1,61 @@
+from enum import Enum
+
+
+class MessageType(Enum):
+    """
+    Defines possible botnet message types
+    """
+
+    TIMEOUT = 3
+    SALITY_NL_REQUEST = 101
+    SALITY_NL_REPLY = 102
+    SALITY_HELLO = 103
+    SALITY_HELLO_REPLY = 104
+
+    def is_request(self):
+        """
+        Checks whether the given message type is a request or not.
+        :return: True if it is a request, False otherwise
+        """
+        return self in [MessageType.SALITY_HELLO, MessageType.SALITY_NL_REQUEST]
+
+    def is_response(self):
+        """
+        Checks whether the given message type is a response or not.
+        :return: True if it is a response, False otherwise
+        """
+        return self in [MessageType.SALITY_HELLO_REPLY, MessageType.SALITY_NL_REPLY]
+
+
+class Message:
+    INVALID_LINENO = -1
+
+    """
+    Defines a compact message type that contains all necessary information.
+    """
+
+    def __init__(self, msg_id: int, src, dst, type_: MessageType, time: float, refer_msg_id: int = -1, line_no=-1):
+        """
+        Constructs a message with the given parameters.
+
+        :param msg_id: the ID of the message
+        :param src: something identifying the source, e.g. ID or configuration
+        :param dst: something identifying the destination, e.g. ID or configuration
+        :param type_: the type of the message
+        :param time: the timestamp of the message
+        :param refer_msg_id: the ID this message is a request for or reply to. -1 if there is no related message.
+        :param line_no: The line number this message appeared at in the original CSV file
+        """
+        self.msg_id = msg_id
+        self.src = src
+        self.dst = dst
+        self.type = type_
+        self.time = time
+        self.csv_time = time
+        self.refer_msg_id = refer_msg_id
+        self.line_no = line_no
+
+    def __str__(self):
+        str_ = "{0}. at {1}: {2}-->{3}, {4}, refer:{5} (line {6})".format(self.msg_id, self.time, self.src, self.dst,
+                                                                          self.type, self.refer_msg_id, self.line_no)
+        return str_

code/ID2TLib/Generator.py

@@ -8,7 +8,7 @@ from numpy.random import bytes
 from random import getrandbits
 from scapy.layers.inet import IP, Ether, UDP, TCP
 from scapy.packet import Raw
-from Attack.MembersMgmtCommAttack import MessageType
+from ID2TLib.Botnet.Message import MessageType
 from . import IPv4 as ip