CommunicationProcessor.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224
  1. from lea import Lea
  2. from random import randrange
  3. from ID2TLib.Botnet.Message import Message
  4. from ID2TLib.Botnet.Message import MessageType
  5. class CommunicationProcessor:
  6. """
  7. Class to process parsed input CSV/XML data and retrieve a mapping or other information.
  8. """
  9. def __init__(self, mtypes: dict, nat: bool):
  10. """
  11. Creates an instance of CommunicationProcessor.
  12. :param mtypes: a dict containing an int to EnumType mapping of MessageTypes
  13. :param nat: whether NAT is present in this network
  14. """
  15. self.packets = []
  16. self.mtypes = mtypes
  17. self.nat = nat
  18. self.messages = []
  19. self.respnd_ids = set()
  20. self.external_init_ids = set()
  21. self.local_init_ids = dict()
  22. self.local_ids = dict()
  23. self.external_ids = set()
  24. def set_mapping(self, packets: list, mapped_ids: dict):
  25. """
  26. Set the selected mapping for this communication processor.
  27. :param packets: all packets contained in the mapped time frame
  28. :param mapped_ids: the chosen IDs
  29. """
  30. self.packets = packets
  31. self.local_init_ids = set(mapped_ids)
  32. @staticmethod
  33. def get_comm_interval(cpp_comm_proc, strategy: str, number_ids: int, max_int_time: int, start_idx: int,
  34. end_idx: int):
  35. """
  36. Finds a communication interval with respect to the given strategy. The interval is maximum of the given seconds
  37. and has at least number_ids communicating initiators in it.
  38. :param cpp_comm_proc: An instance of the C++ communication processor that stores all the input messages and
  39. is responsible for retrieving the interval(s)
  40. :param strategy: The selection strategy (i.e. random, optimal, custom)
  41. :param number_ids: The number of initiator IDs that have to exist in the interval(s)
  42. :param max_int_time: The maximum time period of the interval
  43. :param start_idx: The message index the interval should start at (None if not specified)
  44. :param end_idx: The message index the interval should stop at (inclusive) (None if not specified)
  45. :return: A dict representing the communication interval. It contains the initiator IDs,
  46. the start index and end index of the respective interval. The respective keys
  47. are {IDs, Start, End}. If no interval is found, an empty dict is returned.
  48. """
  49. if strategy == "random":
  50. # try finding not-empty interval 5 times
  51. for i in range(5):
  52. start_idx = randrange(0, cpp_comm_proc.get_message_count())
  53. interval = cpp_comm_proc.find_interval_from_startidx(start_idx, number_ids, max_int_time)
  54. if interval and interval["IDs"]:
  55. return interval
  56. return {}
  57. elif strategy == "optimal":
  58. intervals = cpp_comm_proc.find_optimal_interval(number_ids, max_int_time)
  59. if not intervals:
  60. return {}
  61. else:
  62. for i in range(5):
  63. interval = intervals[randrange(0, len(intervals))]
  64. if interval and interval["IDs"]:
  65. return interval
  66. return {}
  67. elif strategy == "custom":
  68. if (not start_idx) and (not end_idx):
  69. print("Custom strategy was selected, but no (valid) start or end index was specified.")
  70. print("Because of this, a random interval is selected.")
  71. start_idx = randrange(0, cpp_comm_proc.get_message_count())
  72. interval = cpp_comm_proc.find_interval_from_startidx(start_idx, number_ids, max_int_time)
  73. elif (not start_idx) and end_idx:
  74. end_idx -= 1 # because message indices start with 1 (for the user)
  75. interval = cpp_comm_proc.find_interval_from_endidx(end_idx, number_ids, max_int_time)
  76. elif start_idx and (not end_idx):
  77. start_idx -= 1 # because message indices start with 1 (for the user)
  78. interval = cpp_comm_proc.find_interval_from_startidx(start_idx, number_ids, max_int_time)
  79. elif start_idx and end_idx:
  80. start_idx -= 1
  81. end_idx -= 1
  82. ids = cpp_comm_proc.get_interval_init_ids(start_idx, end_idx)
  83. if not ids:
  84. return {}
  85. return {"IDs": ids, "Start": start_idx, "End": end_idx}
  86. if not interval or not interval["IDs"]:
  87. return {}
  88. return interval
  89. def det_id_roles_and_msgs(self):
  90. """
  91. Determine the role of every mapped ID. The role can be initiator, responder or both.
  92. On the side also connect corresponding messages together to quickly find out
  93. which reply belongs to which request and vice versa.
  94. :return: the selected messages
  95. """
  96. mtypes = self.mtypes
  97. # setup initial variables and their values
  98. respnd_ids = set()
  99. # msgs --> the filtered messages, msg_id --> an increasing ID to give every message an artificial primary key
  100. msgs, msg_id = [], 0
  101. # keep track of previous request to find connections
  102. prev_reqs = {}
  103. # used to determine whether a request has been seen yet, so that replies before the first request are skipped
  104. # and do not throw an error by accessing the empty dict prev_reqs (this is not a perfect solution, but it works
  105. # most of the time)
  106. req_seen = False
  107. local_init_ids = self.local_init_ids
  108. external_init_ids = set()
  109. # process every packet individually
  110. for packet in self.packets:
  111. id_src, id_dst, msg_type, time = packet["Src"], packet["Dst"], int(packet["Type"]), float(packet["Time"])
  112. lineno = packet.get("LineNumber", -1)
  113. # if if either one of the IDs is not mapped, continue
  114. if (id_src not in local_init_ids) and (id_dst not in local_init_ids):
  115. continue
  116. # convert message type number to enum type
  117. msg_type = mtypes[msg_type]
  118. # process a request
  119. if msg_type in {MessageType.SALITY_HELLO, MessageType.SALITY_NL_REQUEST}:
  120. if not self.nat and id_dst in local_init_ids and id_src not in local_init_ids:
  121. external_init_ids.add(id_src)
  122. elif id_src not in local_init_ids:
  123. continue
  124. else:
  125. # process ID's role
  126. respnd_ids.add(id_dst)
  127. # convert the abstract message into a message object to handle it better
  128. msg_str = "{0}-{1}".format(id_src, id_dst)
  129. msg = Message(msg_id, id_src, id_dst, msg_type, time, line_no=lineno)
  130. msgs.append(msg)
  131. prev_reqs[msg_str] = msg_id
  132. msg_id += 1
  133. req_seen = True
  134. # process a reply
  135. elif msg_type in {MessageType.SALITY_HELLO_REPLY, MessageType.SALITY_NL_REPLY} and req_seen:
  136. if not self.nat and id_src in local_init_ids and id_dst not in local_init_ids:
  137. # process ID's role
  138. external_init_ids.add(id_dst)
  139. elif id_dst not in local_init_ids:
  140. continue
  141. else:
  142. # process ID's role
  143. respnd_ids.add(id_src)
  144. # convert the abstract message into a message object to handle it better
  145. msg_str = "{0}-{1}".format(id_dst, id_src)
  146. # find the request message ID for this response and set its reference index
  147. refer_idx = prev_reqs.get(msg_str, -1)
  148. if refer_idx != -1:
  149. msgs[refer_idx].refer_msg_id = msg_id
  150. del(prev_reqs[msg_str])
  151. msg = Message(msg_id, id_src, id_dst, msg_type, time, refer_idx, lineno)
  152. msgs.append(msg)
  153. # remove the request to this response from storage
  154. msg_id += 1
  155. elif msg_type == MessageType.TIMEOUT and id_src in local_init_ids and not self.nat:
  156. # convert the abstract message into a message object to handle it better
  157. msg_str = "{0}-{1}".format(id_dst, id_src)
  158. # find the request message ID for this response and set its reference index
  159. refer_idx = prev_reqs.get(msg_str)
  160. if refer_idx is not None:
  161. msgs[refer_idx].refer_msg_id = msg_id
  162. if msgs[refer_idx].type == MessageType.SALITY_NL_REQUEST:
  163. msg = Message(msg_id, id_src, id_dst, MessageType.SALITY_NL_REPLY, time, refer_idx, lineno)
  164. else:
  165. msg = Message(msg_id, id_src, id_dst, MessageType.SALITY_HELLO_REPLY, time, refer_idx, lineno)
  166. msgs.append(msg)
  167. # remove the request to this response from storage
  168. del(prev_reqs[msg_str])
  169. msg_id += 1
  170. # store the retrieved information in this object for later use
  171. self.respnd_ids = sorted(respnd_ids)
  172. self.external_init_ids = sorted(external_init_ids)
  173. self.messages = msgs
  174. # return the selected messages
  175. return self.messages
  176. def det_ext_and_local_ids(self, prob_rspnd_local: int=0):
  177. """
  178. Map the given IDs to a locality (i.e. local or external} considering the given probabilities.
  179. :param prob_rspnd_local: the probabilty that a responder is local
  180. """
  181. external_ids = set()
  182. local_ids = self.local_init_ids.copy()
  183. # set up probabilistic chooser
  184. rspnd_locality = Lea.fromValFreqsDict({"local": prob_rspnd_local*100, "external": (1-prob_rspnd_local)*100})
  185. for id_ in self.external_init_ids:
  186. external_ids.add(id_)
  187. # determine responder localities
  188. for id_ in self.respnd_ids:
  189. if id_ in local_ids or id_ in external_ids:
  190. continue
  191. pos = rspnd_locality.random()
  192. if pos == "local":
  193. local_ids.add(id_)
  194. elif pos == "external":
  195. external_ids.add(id_)
  196. self.local_ids, self.external_ids = local_ids, external_ids
  197. return self.local_ids, self.external_ids