3
0

LabelManager.py 10 KB

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