3
0

Message.py 2.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  1. from enum import Enum
  2. class MessageType(Enum):
  3. """
  4. Defines possible botnet message types
  5. """
  6. TIMEOUT = 3
  7. SALITY_NL_REQUEST = 101
  8. SALITY_NL_REPLY = 102
  9. SALITY_HELLO = 103
  10. SALITY_HELLO_REPLY = 104
  11. def is_request(self):
  12. """
  13. Checks whether the given message type is a request or not.
  14. :return: True if it is a request, False otherwise
  15. """
  16. return self in [MessageType.SALITY_HELLO, MessageType.SALITY_NL_REQUEST]
  17. def is_response(self):
  18. """
  19. Checks whether the given message type is a response or not.
  20. :return: True if it is a response, False otherwise
  21. """
  22. return self in [MessageType.SALITY_HELLO_REPLY, MessageType.SALITY_NL_REPLY]
  23. class Message:
  24. INVALID_LINENO = -1
  25. """
  26. Defines a compact message type that contains all necessary information.
  27. """
  28. def __init__(self, msg_id: int, src, dst, type_: MessageType, time: float, refer_msg_id: int = -1, line_no=-1):
  29. """
  30. Constructs a message with the given parameters.
  31. :param msg_id: the ID of the message
  32. :param src: something identifying the source, e.g. ID or configuration
  33. :param dst: something identifying the destination, e.g. ID or configuration
  34. :param type_: the type of the message
  35. :param time: the timestamp of the message
  36. :param refer_msg_id: the ID this message is a request for or reply to. -1 if there is no related message.
  37. :param line_no: The line number this message appeared at in the original CSV file
  38. """
  39. self.msg_id = msg_id
  40. self.src = src
  41. self.dst = dst
  42. self.type = type_
  43. self.time = time
  44. self.csv_time = time
  45. self.refer_msg_id = refer_msg_id
  46. self.line_no = line_no
  47. def __str__(self):
  48. str_ = "{0}. at {1}: {2}-->{3}, {4}, refer:{5} (line {6})".format(self.msg_id, self.time, self.src, self.dst,
  49. self.type, self.refer_msg_id, self.line_no)
  50. return str_