import pickle
import threading
from myutils.ConfigManager import myCongif

class PickleManager:
    def __init__(self):
        self.lock = threading.Lock()  # 线程锁
        self.tree_file = myCongif.get_data("TreeFile")

    def WriteData(self,attack_tree,filename=""):
        if filename:
            filepath = "tree_data/"+filename
        else:
            filepath = self.tree_file

        with self.lock:
            with open(filepath, 'wb') as f:
                pickle.dump(attack_tree, f)

    def ReadData(self,filename=""):
        attack_tree = None
        if filename:
            filepath = "tree_data/"+filename
        else:
            filepath = self.tree_file

        with self.lock:
            with open(filepath, "rb") as f:
                attack_tree = pickle.load(f)
        return attack_tree

g_PKM = PickleManager()