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_