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

import tempfile

class AttackContext:
    def __init__(self):
        # 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 = {}

    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 reset(self):
        self.allocated_files.clear()

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

        self.allocated_files.clear()
        return return_