LabelManager.py 5.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140
  1. import os.path
  2. from datetime import datetime
  3. from xml.dom.minidom import *
  4. import ID2TLib.Label as Label
  5. class LabelManager:
  6. TAG_ROOT = 'LABELS'
  7. TAG_ATTACK = 'attack'
  8. TAG_ATTACK_NAME = 'attack_name'
  9. TAG_ATTACK_NOTE = 'attack_note'
  10. TAG_TIMESTAMP_START = 'timestamp_start'
  11. TAG_TIMESTAMP_END = 'timestamp_end'
  12. TAG_TIMESTAMP = 'timestamp'
  13. TAG_TIMESTAMP_HR = 'timestamp_hr'
  14. ATTR_VERSION = 'version_parser'
  15. # update this attribute if XML scheme was modified
  16. ATTR_VERSION_VALUE = '0.2'
  17. def __init__(self, filepath_pcap=None):
  18. """
  19. Creates a new LabelManager for managing the attack's labels.
  20. :param filepath_pcap: The path to the PCAP file associated to the labels.
  21. """
  22. self.labels = list()
  23. if filepath_pcap is not None:
  24. self.labelFilePath = filepath_pcap.strip('.pcap') + '_labels.xml'
  25. # only load labels if label file is existing
  26. if os.path.exists(self.labelFilePath):
  27. self._load_labels()
  28. def add_labels(self, labels):
  29. """
  30. Adds a label to the internal list of labels.
  31. :param labels: The labels to be added
  32. """
  33. if isinstance(labels, list):
  34. self.labels = self.labels + [labels]
  35. elif isinstance(labels, tuple):
  36. for l in labels:
  37. self.labels.append(l)
  38. else:
  39. self.labels.append(labels)
  40. # sorts the labels ascending by their timestamp
  41. self.labels.sort()
  42. def write_label_file(self, filepath=None):
  43. """
  44. Writes previously added/loaded labels to a XML file. Uses the given filepath as destination path, if no path is
  45. given, uses the path in labelFilePath.
  46. :param filepath: The path where the label file should be written to.
  47. """
  48. def get_subtree_timestamp(xml_tag_root, timestamp_entry) -> Element:
  49. """
  50. Creates the subtree for a given timestamp, consisting of the unix time format (seconds) and a human-readable
  51. output.
  52. :param xml_tag_root: The tag name for the root of the subtree
  53. :param timestamp_entry: The timestamp as unix time
  54. :return: The root node of the XML subtree
  55. """
  56. timestamp_root = doc.createElement(xml_tag_root)
  57. # add timestamp in unix format
  58. timestamp = doc.createElement(self.TAG_TIMESTAMP)
  59. timestamp.appendChild(doc.createTextNode(str(timestamp_entry)))
  60. timestamp_root.appendChild(timestamp)
  61. # add timestamp in human-readable format
  62. timestamp_hr = doc.createElement(self.TAG_TIMESTAMP_HR)
  63. timestamp_hr_text = datetime.fromtimestamp(int(timestamp_entry)).strftime('%Y-%m-%d %H:%M:%S')
  64. timestamp_hr.appendChild(doc.createTextNode(timestamp_hr_text))
  65. timestamp_root.appendChild(timestamp_hr)
  66. return timestamp_root
  67. if filepath is not None:
  68. self.labelFilePath = filepath.strip('.pcap') + '_labels.xml'
  69. # Generate XML
  70. doc = Document()
  71. node = doc.createElement(self.TAG_ROOT)
  72. node.setAttribute(self.ATTR_VERSION, self.ATTR_VERSION_VALUE)
  73. for label in self.labels:
  74. xml_tree = doc.createElement(self.TAG_ATTACK)
  75. # add attack to XML tree
  76. attack_name = doc.createElement(self.TAG_ATTACK_NAME)
  77. attack_name.appendChild(doc.createTextNode(str(label.attack_name)))
  78. xml_tree.appendChild(attack_name)
  79. attack_note = doc.createElement(self.TAG_ATTACK_NOTE)
  80. attack_note.appendChild(doc.createTextNode(str(label.attack_note)))
  81. xml_tree.appendChild(attack_note)
  82. # add timestamp_start to XML tree
  83. xml_tree.appendChild(get_subtree_timestamp(self.TAG_TIMESTAMP_START, label.timestamp_start))
  84. # add timestamp_end to XML tree
  85. xml_tree.appendChild(get_subtree_timestamp(self.TAG_TIMESTAMP_END, label.timestamp_end))
  86. node.appendChild(xml_tree)
  87. doc.appendChild(node)
  88. # Write XML to file
  89. file = open(self.labelFilePath, 'w')
  90. file.write(doc.toprettyxml())
  91. file.close()
  92. def _load_labels(self):
  93. """
  94. Loads the labels from an already existing label XML file located at labelFilePath (set by constructor).
  95. """
  96. print("Label file found. Loading labels...")
  97. dom = parse(self.labelFilePath)
  98. # Check if version of parser and version of file match
  99. version = dom.getElementsByTagName(self.TAG_ROOT)[0].getAttribute(self.ATTR_VERSION)
  100. if not version == self.ATTR_VERSION_VALUE:
  101. raise ValueError(
  102. "The file " + self.labelFilePath + " was created by another version of ID2TLib.LabelManager")
  103. # Parse attacks from XML file
  104. attacks = dom.getElementsByTagName(self.TAG_ATTACK)
  105. for a in attacks:
  106. attack_name = a.childNodes[1].firstChild.data
  107. attack_note = a.childNodes[3].firstChild.data
  108. timestamp_start = a.childNodes[5].childNodes[1].firstChild.data
  109. timestamp_end = a.childNodes[7].childNodes[1].firstChild.data
  110. label = Label.Label(attack_name, float(timestamp_start), float(timestamp_end), attack_note)
  111. self.labels.append(label)