3
0

LabelManager.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251
  1. import importlib
  2. import datetime as dt
  3. import os.path
  4. import xml.dom.minidom as minidom
  5. import pytz as pytz
  6. import ID2TLib.Label as Label
  7. import ID2TLib.TestLibrary as Lib
  8. class LabelManager:
  9. TAG_ROOT = 'labels'
  10. TAG_INPUT = 'input'
  11. TAG_OUTPUT = 'output'
  12. TAG_FILE_NAME = 'filename'
  13. TAG_FILE_HASH = 'sha256'
  14. TAG_ATTACK = 'attack'
  15. TAG_ATTACK_NAME = 'name'
  16. TAG_ATTACK_NOTE = 'note'
  17. TAG_ATTACK_SEED = 'seed'
  18. TAG_TIMESTAMP_START = 'timestamp_start'
  19. TAG_TIMESTAMP_END = 'timestamp_end'
  20. TAG_TIMESTAMP = 'timestamp'
  21. TAG_TIMESTAMP_HR = 'timestamp_hr'
  22. TAG_PARAMETERS = 'parameters'
  23. ATTR_VERSION = 'version_parser'
  24. ATTR_PARAM_USERSPECIFIED = 'user_specified'
  25. # update this attribute if XML scheme was modified
  26. ATTR_VERSION_VALUE = '0.3'
  27. def __init__(self, filepath_pcap=None):
  28. """
  29. Creates a new LabelManager for managing the attack's labels.
  30. :param filepath_pcap: The path to the PCAP file associated to the labels.
  31. """
  32. self.labels = list()
  33. self.filepath_input_pcap = filepath_pcap
  34. if filepath_pcap is not None:
  35. self.label_file_path = os.path.splitext(filepath_pcap)[0] + '_labels.xml'
  36. # only load labels if label file is existing
  37. if os.path.exists(self.label_file_path):
  38. self.load_labels()
  39. def add_labels(self, labels):
  40. """
  41. Adds a label to the internal list of labels.
  42. :param labels: The labels to be added
  43. """
  44. if isinstance(labels, list):
  45. self.labels = self.labels + [labels]
  46. elif isinstance(labels, tuple):
  47. for l in labels:
  48. self.labels.append(l)
  49. else:
  50. self.labels.append(labels)
  51. # sorts the labels ascending by their timestamp
  52. self.labels.sort()
  53. def write_label_file(self, filepath=None):
  54. """
  55. Writes previously added/loaded labels to a XML file. Uses the given filepath as destination path, if no path is
  56. given, uses the path in label_file_path.
  57. :param filepath: The path where the label file should be written to.
  58. """
  59. def get_subtree_fileinfo(xml_tag_root, filename) -> minidom.Element:
  60. """
  61. Creates the subtree for pcap file information (filename and hash).
  62. :return: The root node of the XML subtree
  63. """
  64. input_root = doc.createElement(xml_tag_root)
  65. file = doc.createElement(self.TAG_FILE_NAME)
  66. file.appendChild(doc.createTextNode(os.path.split(filename)[-1]))
  67. input_root.appendChild(file)
  68. hash_node = doc.createElement(self.TAG_FILE_HASH)
  69. hash_node.appendChild(doc.createTextNode(Lib.get_sha256(filename)))
  70. input_root.appendChild(hash_node)
  71. return input_root
  72. def get_subtree_timestamp(xml_tag_root, timestamp_entry):
  73. """
  74. Creates the subtree for a given timestamp, consisting of the unix time format (seconds) and a human-readable
  75. output.
  76. :param xml_tag_root: The tag name for the root of the subtree
  77. :param timestamp_entry: The timestamp as unix time
  78. :return: The root node of the XML subtree
  79. """
  80. timestamp_root = doc.createElement(xml_tag_root)
  81. # add timestamp in unix format
  82. timestamp = doc.createElement(self.TAG_TIMESTAMP)
  83. timestamp.appendChild(doc.createTextNode(str(timestamp_entry)))
  84. timestamp_root.appendChild(timestamp)
  85. # add timestamp in human-readable format
  86. timestamp_hr = doc.createElement(self.TAG_TIMESTAMP_HR)
  87. timestamp_hr_text = dt.datetime.fromtimestamp(timestamp_entry).astimezone(pytz.timezone('UTC')).strftime('%Y-%m-%d %H:%M:%S.%f')
  88. timestamp_hr.appendChild(doc.createTextNode(timestamp_hr_text))
  89. timestamp_root.appendChild(timestamp_hr)
  90. return timestamp_root
  91. def get_subtree_parameters(parameters):
  92. """
  93. Creates a subtree containing all parameters used to construct the attack
  94. :param parameters: The list of parameters used to run the attack
  95. :return: The root node of the XML subtree
  96. """
  97. parameters_root = doc.createElement(self.TAG_PARAMETERS)
  98. for param_key, param_value in parameters.items():
  99. param = doc.createElement(param_key.value)
  100. param.appendChild(doc.createTextNode(str(param_value.value)))
  101. param.setAttribute(self.ATTR_PARAM_USERSPECIFIED, str(param_value.user_specified))
  102. parameters_root.appendChild(param)
  103. return parameters_root
  104. if filepath is not None:
  105. self.label_file_path = os.path.splitext(filepath)[0] + '_labels.xml'
  106. # Generate XML
  107. doc = minidom.Document()
  108. node = doc.createElement(self.TAG_ROOT)
  109. node.setAttribute(self.ATTR_VERSION, self.ATTR_VERSION_VALUE)
  110. node.appendChild(get_subtree_fileinfo(self.TAG_INPUT, self.filepath_input_pcap))
  111. node.appendChild(get_subtree_fileinfo(self.TAG_OUTPUT, filepath))
  112. for label in self.labels:
  113. xml_tree = doc.createElement(self.TAG_ATTACK)
  114. # add attack to XML tree
  115. attack_name = doc.createElement(self.TAG_ATTACK_NAME)
  116. attack_name.appendChild(doc.createTextNode(str(label.attack_name)))
  117. xml_tree.appendChild(attack_name)
  118. attack_note = doc.createElement(self.TAG_ATTACK_NOTE)
  119. attack_note.appendChild(doc.createTextNode(str(label.attack_note)))
  120. xml_tree.appendChild(attack_note)
  121. attack_seed = doc.createElement(self.TAG_ATTACK_SEED)
  122. attack_seed.appendChild(doc.createTextNode(str(label.seed)))
  123. xml_tree.appendChild(attack_seed)
  124. # add timestamp_start to XML tree
  125. xml_tree.appendChild(get_subtree_timestamp(self.TAG_TIMESTAMP_START, label.timestamp_start))
  126. # add timestamp_end to XML tree
  127. xml_tree.appendChild(get_subtree_timestamp(self.TAG_TIMESTAMP_END, label.timestamp_end))
  128. # add parameters to XML tree
  129. xml_tree.appendChild(get_subtree_parameters(label.parameters))
  130. node.appendChild(xml_tree)
  131. doc.appendChild(node)
  132. # Write XML to file
  133. file = open(self.label_file_path, 'w')
  134. file.write(doc.toprettyxml())
  135. file.close()
  136. def load_labels(self):
  137. """
  138. Loads the labels from an already existing label XML file located at label_file_path (set by constructor).
  139. """
  140. def get_value_from_node(node, tag_name, *child_number):
  141. """
  142. Returns the value located in the tag specified by tag_name from a given node. Walks therefor the
  143. node's children along as indicated by child_number, e.g., childNumber = (1,2,) first goes to the 1st child,
  144. and then to the 2nd child of the first child -> elem.childNodes[1].childNodes[2].
  145. """
  146. elem = node.getElementsByTagName(tag_name)
  147. if len(elem) == 1:
  148. elem = elem[0]
  149. for c in child_number:
  150. if len(elem.childNodes) > 0:
  151. elem = elem.childNodes[c]
  152. else:
  153. return ""
  154. return elem.data
  155. else:
  156. return ""
  157. print("Label file found. Loading labels...")
  158. try:
  159. dom = minidom.parse(self.label_file_path)
  160. except Exception:
  161. # TODO: more specific exception
  162. print('ERROR: Provided label file could not be parsed. Ignoring label file')
  163. return
  164. # Check if version of parser and version of file match
  165. version = dom.getElementsByTagName(self.TAG_ROOT)
  166. if len(version) > 0:
  167. version = version[0].getAttribute(self.ATTR_VERSION)
  168. if version == [] or not version == self.ATTR_VERSION_VALUE:
  169. print(
  170. "The file " + self.label_file_path + " was created by another version of ID2TLib.LabelManager. "
  171. "Ignoring label file.")
  172. self.input_filename = get_value_from_node(dom, self.TAG_INPUT, 1, 0)
  173. self.input_hash = get_value_from_node(dom, self.TAG_INPUT, 3, 0)
  174. self.output_filename = get_value_from_node(dom, self.TAG_OUTPUT, 1, 0)
  175. self.output_hash = get_value_from_node(dom, self.TAG_OUTPUT, 3, 0)
  176. # Parse attacks from XML file
  177. attacks = dom.getElementsByTagName(self.TAG_ATTACK)
  178. count_labels = 0
  179. for a in attacks:
  180. attack_name = get_value_from_node(a, self.TAG_ATTACK_NAME, 0)
  181. attack_note = get_value_from_node(a, self.TAG_ATTACK_NOTE, 0)
  182. timestamp_start = get_value_from_node(a, self.TAG_TIMESTAMP_START, 1, 0)
  183. timestamp_end = get_value_from_node(a, self.TAG_TIMESTAMP_END, 1, 0)
  184. attack_seed = get_value_from_node(a, self.TAG_ATTACK_SEED, 0)
  185. # Instantiate this attack to create a parameter list with the correct types
  186. attack_module = importlib.import_module("Attack." + attack_name)
  187. attack_class = getattr(attack_module, attack_name)
  188. attack = attack_class()
  189. # Loop through all parameters listed in the XML file
  190. param = a.getElementsByTagName(self.TAG_PARAMETERS)[0]
  191. for param in param.childNodes:
  192. # Skip empty text nodes returned by minidom
  193. if not isinstance(param, minidom.Text):
  194. import distutils.util
  195. param_name = param.tagName
  196. param_value = param.childNodes[0].nodeValue
  197. param_userspecified = bool(distutils.util.strtobool(param.getAttribute(self.ATTR_PARAM_USERSPECIFIED)))
  198. attack.add_param_value(param_name, param_value, param_userspecified)
  199. # Create the label from the data read
  200. label = Label.Label(attack_name, float(timestamp_start), float(timestamp_end), attack_seed, attack.params,
  201. attack_note)
  202. self.labels.append(label)
  203. count_labels += 1
  204. print("Read " + str(count_labels) + " label(s) successfully.")