ID2T-toolkit/code/Core/LabelManager.py

import importlib
import datetime as dt
import os.path
import xml.dom.minidom as minidom
import pytz as pytz

import ID2TLib.Label as Label
import ID2TLib.TestLibrary as Lib


class LabelManager:
    TAG_ROOT = 'labels'
    TAG_INPUT = 'input'
    TAG_OUTPUT = 'output'
    TAG_FILE_NAME = 'filename'
    TAG_FILE_HASH = 'sha256'
    TAG_ATTACK = 'attack'
    TAG_ATTACK_NAME = 'name'
    TAG_ATTACK_NOTE = 'note'
    TAG_ATTACK_SEED = 'seed'
    TAG_TIMESTAMP_START = 'timestamp_start'
    TAG_TIMESTAMP_END = 'timestamp_end'
    TAG_TIMESTAMP = 'timestamp'
    TAG_TIMESTAMP_HR = 'timestamp_hr'
    TAG_PARAMETERS = 'parameters'
    ATTR_VERSION = 'version_parser'
    ATTR_PARAM_USERSPECIFIED = 'user_specified'

    # update this attribute if XML scheme was modified
    ATTR_VERSION_VALUE = '0.3'

    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()
        self.filepath_input_pcap = filepath_pcap

        if filepath_pcap is not None:
            self.label_file_path = os.path.splitext(filepath_pcap)[0] + '_labels.xml'
            # only load labels if label file is existing
            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)

        # sorts the labels ascending by their timestamp
        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_fileinfo(xml_tag_root, filename) -> minidom.Element:
            """
            Creates the subtree for pcap file information (filename and hash).

            :return: The root node of the XML subtree
            """

            input_root = doc.createElement(xml_tag_root)

            file = doc.createElement(self.TAG_FILE_NAME)
            file.appendChild(doc.createTextNode(os.path.split(filename)[-1]))
            input_root.appendChild(file)

            hash_node = doc.createElement(self.TAG_FILE_HASH)
            hash_node.appendChild(doc.createTextNode(Lib.get_sha256(filename)))
            input_root.appendChild(hash_node)

            return input_root

        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)

            # add timestamp in unix format
            timestamp = doc.createElement(self.TAG_TIMESTAMP)
            timestamp.appendChild(doc.createTextNode(str(timestamp_entry)))
            timestamp_root.appendChild(timestamp)

            # add timestamp in human-readable format
            timestamp_hr = doc.createElement(self.TAG_TIMESTAMP_HR)
            timestamp_hr_text = dt.datetime.fromtimestamp(timestamp_entry).astimezone(pytz.timezone('UTC')).strftime('%Y-%m-%d %H:%M:%S.%f')
            timestamp_hr.appendChild(doc.createTextNode(timestamp_hr_text))
            timestamp_root.appendChild(timestamp_hr)

            return timestamp_root

        def get_subtree_parameters(parameters):
            """
            Creates a subtree containing all parameters used to construct the attack

            :param parameters: The list of parameters used to run the attack
            :return: The root node of the XML subtree
            """
            parameters_root = doc.createElement(self.TAG_PARAMETERS)

            for param_key, param_value in parameters.items():
                param = doc.createElement(param_key.value)
                param.appendChild(doc.createTextNode(str(param_value.value)))
                param.setAttribute(self.ATTR_PARAM_USERSPECIFIED, str(param_value.user_specified))
                parameters_root.appendChild(param)

            return parameters_root

        if filepath is not None:
            self.label_file_path = os.path.splitext(filepath)[0] + '_labels.xml'

        # Generate XML
        doc = minidom.Document()
        node = doc.createElement(self.TAG_ROOT)
        node.setAttribute(self.ATTR_VERSION, self.ATTR_VERSION_VALUE)
        node.appendChild(get_subtree_fileinfo(self.TAG_INPUT, self.filepath_input_pcap))
        node.appendChild(get_subtree_fileinfo(self.TAG_OUTPUT, filepath))

        for label in self.labels:
            xml_tree = doc.createElement(self.TAG_ATTACK)

            # add attack to XML tree
            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)
            attack_seed = doc.createElement(self.TAG_ATTACK_SEED)
            attack_seed.appendChild(doc.createTextNode(str(label.seed)))
            xml_tree.appendChild(attack_seed)

            # add timestamp_start to XML tree
            xml_tree.appendChild(get_subtree_timestamp(self.TAG_TIMESTAMP_START, label.timestamp_start))

            # add timestamp_end to XML tree
            xml_tree.appendChild(get_subtree_timestamp(self.TAG_TIMESTAMP_END, label.timestamp_end))

            # add parameters to XML tree
            xml_tree.appendChild(get_subtree_parameters(label.parameters))

            node.appendChild(xml_tree)

        doc.appendChild(node)

        # Write XML to file
        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 = minidom.parse(self.label_file_path)
        except Exception:
            # TODO: more specific exception
            print('ERROR: Provided label file could not be parsed. Ignoring label file')
            return

        # Check if version of parser and version of file match
        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.")

        self.input_filename = get_value_from_node(dom, self.TAG_INPUT, 1, 0)
        self.input_hash = get_value_from_node(dom, self.TAG_INPUT, 3, 0)
        self.output_filename = get_value_from_node(dom, self.TAG_OUTPUT, 1, 0)
        self.output_hash = get_value_from_node(dom, self.TAG_OUTPUT, 3, 0)

        # Parse attacks from XML 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)
            attack_seed = get_value_from_node(a, self.TAG_ATTACK_SEED, 0)

            # Instantiate this attack to create a parameter list with the correct types
            attack_module = importlib.import_module("Attack." + attack_name)
            attack_class = getattr(attack_module, attack_name)
            attack = attack_class()

            # Loop through all parameters listed in the XML file
            param = a.getElementsByTagName(self.TAG_PARAMETERS)[0]
            for param in param.childNodes:
                # Skip empty text nodes returned by minidom
                if not isinstance(param, minidom.Text):
                    import distutils.util
                    param_name = param.tagName
                    param_value = param.childNodes[0].nodeValue
                    param_userspecified = bool(distutils.util.strtobool(param.getAttribute(self.ATTR_PARAM_USERSPECIFIED)))
                    attack.add_param_value(param_name, param_value, param_userspecified)

            # Create the label from the data read
            label = Label.Label(attack_name, float(timestamp_start), float(timestamp_end), attack_seed, attack.params,
                                attack_note)
            self.labels.append(label)
            count_labels += 1

        print("Read " + str(count_labels) + " label(s) successfully.")