zfsafe/myutils/PickleManager.py

import pickle
import threading
import os
from myutils.ConfigManager import myCongif

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

    def getfile_path(self,filename=""):
        filepath = self.tree_file
        if filename:
            filepath = "tree_data/"+filename
        return filepath

    def WriteData(self,attack_tree,filename=""):
        filepath = self.getfile_path(filename)
        with self.lock:
            with open(filepath, 'wb') as f:
                pickle.dump(attack_tree, f)

    def ReadData(self,filename=""):
        attack_tree = None
        filepath = self.getfile_path(filename)
        with self.lock:
            with open(filepath, "rb") as f:
                attack_tree = pickle.load(f)
        return attack_tree

    def DelData(self,filename=""):
        filepath = self.getfile_path(filename)
        #删除文件
        try:
            os.remove(filepath)
            return True
        except FileNotFoundError:
            # 文件不存在
            return True
        except PermissionError:
            # 没有删除权限/文件被占用
            return False
        except IsADirectoryError:
            # 路径指向的是目录而非文件
            return False
        except Exception as e:
            # 其他未知错误（如路径非法、存储介质故障等）
            print(f"删除文件时发生意外错误: {str(e)}")
            return False



g_PKM = PickleManager()