123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201 |
- import os.path
- from datetime import datetime
- from xml.dom.minidom import *
- from functools import total_ordering
- @total_ordering
- class Label:
- def __init__(self, attack_name, timestamp_start, timestamp_end, attack_note=""):
- """
- Creates a new attack label
- :param attack_name: The name of the associated attack
- :param timestamp_start: The timestamp as unix time of the first attack packet
- :param timestamp_end: The timestamp as unix time of the last attack packet
- :param attack_note: A note associated to the attack (optional)
- """
- self.attack_name = attack_name
- self.timestamp_start = timestamp_start
- self.timestamp_end = timestamp_end
- self.attack_note = attack_note
- def __eq__(self, other):
- return self.timestamp == other.timestamp
- def __lt__(self, other):
- return self.timestamp_start < other.timestamp_start
- def __gt__(self, other):
- return self.timestamp_start > other.timestamp_start
- def __str__(self):
- return ''.join(
- ['(', self.attack_name, ',', self.attack_note, ',', str(self.timestamp_start), ',', str(self.timestamp_end),
- ')'])
- class LabelManager:
- TAG_ROOT = 'LABELS'
- TAG_ATTACK = 'attack'
- TAG_ATTACK_NAME = 'attack_name'
- TAG_ATTACK_NOTE = 'attack_note'
- TAG_TIMESTAMP_START = 'timestamp_start'
- TAG_TIMESTAMP_END = 'timestamp_end'
- TAG_TIMESTAMP = 'timestamp'
- TAG_TIMESTAMP_HR = 'timestamp_hr'
- ATTR_VERSION = 'version_parser'
-
- ATTR_VERSION_VALUE = '0.2'
- def __init__(self, filepath_pcap=None):
- """
- Creates a new LabelManager for managing the attack's labels.
- :param filepath_pcap: The path to the PCAP file associated to the labels.
- """
- self.labels = list()
- if filepath_pcap is not None:
-
- self.label_file_path = os.path.splitext(filepath_pcap)[0] + '_labels.xml'
-
- if os.path.exists(self.label_file_path):
- self.load_labels()
- def add_labels(self, labels):
- """
- Adds a label to the internal list of labels.
- :param labels: The labels to be added
- """
- if isinstance(labels, list):
- self.labels = self.labels + [labels]
- elif isinstance(labels, tuple):
- for l in labels:
- self.labels.append(l)
- else:
- self.labels.append(labels)
-
- self.labels.sort()
- def write_label_file(self, filepath=None):
- """
- Writes previously added/loaded labels to a XML file. Uses the given filepath as destination path, if no path is
- given, uses the path in label_file_path.
- :param filepath: The path where the label file should be written to.
- """
- def get_subtree_timestamp(xml_tag_root, timestamp_entry):
- """
- Creates the subtree for a given timestamp, consisting of the unix time format (seconds) and a human-readable
- output.
- :param xml_tag_root: The tag name for the root of the subtree
- :param timestamp_entry: The timestamp as unix time
- :return: The root node of the XML subtree
- """
- timestamp_root = doc.createElement(xml_tag_root)
-
- timestamp = doc.createElement(self.TAG_TIMESTAMP)
- timestamp.appendChild(doc.createTextNode(str(timestamp_entry)))
- timestamp_root.appendChild(timestamp)
-
- timestamp_hr = doc.createElement(self.TAG_TIMESTAMP_HR)
- timestamp_hr_text = datetime.fromtimestamp(timestamp_entry).strftime('%Y-%m-%d %H:%M:%S.%f')
- timestamp_hr.appendChild(doc.createTextNode(timestamp_hr_text))
- timestamp_root.appendChild(timestamp_hr)
- return timestamp_root
- if filepath is not None:
- self.label_file_path = os.path.splitext(filepath)[0] + '_labels.xml'
-
- doc = Document()
- node = doc.createElement(self.TAG_ROOT)
- node.setAttribute(self.ATTR_VERSION, self.ATTR_VERSION_VALUE)
- for label in self.labels:
- xml_tree = doc.createElement(self.TAG_ATTACK)
-
- attack_name = doc.createElement(self.TAG_ATTACK_NAME)
- attack_name.appendChild(doc.createTextNode(str(label.attack_name)))
- xml_tree.appendChild(attack_name)
- attack_note = doc.createElement(self.TAG_ATTACK_NOTE)
- attack_note.appendChild(doc.createTextNode(str(label.attack_note)))
- xml_tree.appendChild(attack_note)
-
- xml_tree.appendChild(get_subtree_timestamp(self.TAG_TIMESTAMP_START, label.timestamp_start))
-
- xml_tree.appendChild(get_subtree_timestamp(self.TAG_TIMESTAMP_END, label.timestamp_end))
- node.appendChild(xml_tree)
- doc.appendChild(node)
-
- file = open(self.label_file_path, 'w')
- file.write(doc.toprettyxml())
- file.close()
- def load_labels(self):
- """
- Loads the labels from an already existing label XML file located at label_file_path (set by constructor).
- """
- def get_value_from_node(node, tag_name, *child_number):
- """
- Returns the value located in the tag specified by tag_name from a given node. Walks therefor the
- node's children along as indicated by child_number, e.g., childNumber = (1,2,) first goes to the 1st child, and
- then to the 2nd child of the first child -> elem.childNodes[1].childNodes[2].
- """
- elem = node.getElementsByTagName(tag_name)
- if len(elem) == 1:
- elem = elem[0]
- for c in child_number:
- if len(elem.childNodes) > 0:
- elem = elem.childNodes[c]
- else:
- return ""
- return elem.data
- else:
- return ""
- print("Label file found. Loading labels...")
- try:
- dom = parse(self.label_file_path)
- except Exception:
- print('ERROR: Provided label file could not be parsed. Ignoring label file')
- return
-
- version = dom.getElementsByTagName(self.TAG_ROOT)
- if len(version) > 0:
- version = version[0].getAttribute(self.ATTR_VERSION)
- if version == [] or not version == self.ATTR_VERSION_VALUE:
- print(
- "The file " + self.label_file_path + " was created by another version of ID2TLib.LabelManager. Ignoring label file.")
-
- attacks = dom.getElementsByTagName(self.TAG_ATTACK)
- count_labels = 0
- for a in attacks:
- attack_name = get_value_from_node(a, self.TAG_ATTACK_NAME, 0)
- attack_note = get_value_from_node(a, self.TAG_ATTACK_NOTE, 0)
- timestamp_start = get_value_from_node(a, self.TAG_TIMESTAMP_START, 1, 0)
- timestamp_end = get_value_from_node(a, self.TAG_TIMESTAMP_END, 1, 0)
- label = Label(attack_name, float(timestamp_start), float(timestamp_end), attack_note)
- self.labels.append(label)
- count_labels += 1
- print("Read " + str(count_labels) + " label(s) successfully.")
|