ID2T-toolkit-BotnetTraffic/code/ID2TLib/AttackContext.py

import tempfile

class AttackContext:
    def __init__(self, out_dir):
        # files allocated by attacks, intended to get copied next to the pcap
        # the keys are the suffix for later use, e.g. "_extra_info.txt"
        # the values are the respectable file names
        self.allocated_files = {}
        self.other_created_files = set()
        self.out_dir = out_dir

    def allocate_file(self, suffix):
        if suffix in self.allocated_files:
            raise ValueError("File with suffix %s is already allocated" % suffix)

        file = tempfile.NamedTemporaryFile(mode = "w", suffix = suffix, delete = False)
        self.allocated_files[suffix] = file.name

        return file

    def add_other_created_file(self, filepath):
        self.other_created_files.add(filepath)

    def get_other_created_files(self):
        return sorted(self.other_created_files)

    def reset(self):
        self.allocated_files.clear()

    def get_allocated_files(self):
        return_ = list(self.allocated_files.items())

        self.allocated_files.clear()
        return return_

    def get_output_dir(self):
        return self.out_dir