From 91f5e6b8f754beb47b02f7c1893804c1c9543ccb Mon Sep 17 00:00:00 2001
From: binary-husky <qingxu.fu@outlook.com>
Date: Mon, 4 Nov 2024 13:49:49 +0000
Subject: [PATCH] resolve pickle security issue

---
 crazy_functions/latex_fns/latex_pickle_io.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/crazy_functions/latex_fns/latex_pickle_io.py b/crazy_functions/latex_fns/latex_pickle_io.py
index 7b93ea87..d951bf58 100644
--- a/crazy_functions/latex_fns/latex_pickle_io.py
+++ b/crazy_functions/latex_fns/latex_pickle_io.py
@@ -6,12 +6,16 @@ class SafeUnpickler(pickle.Unpickler):
     def get_safe_classes(self):
         from crazy_functions.latex_fns.latex_actions import LatexPaperFileGroup, LatexPaperSplit
         from crazy_functions.latex_fns.latex_toolbox import LinkedListNode
+        from numpy.core.multiarray import scalar
+        from numpy import dtype
         # 定义允许的安全类
         safe_classes = {
             # 在这里添加其他安全的类
             'LatexPaperFileGroup': LatexPaperFileGroup,
             'LatexPaperSplit': LatexPaperSplit,
             'LinkedListNode': LinkedListNode,
+            'scalar': scalar,
+            'dtype': dtype,
         }
         return safe_classes
 
@@ -22,8 +26,6 @@ class SafeUnpickler(pickle.Unpickler):
         for class_name in self.safe_classes.keys():
             if (class_name in f'{module}.{name}'):
                 match_class_name = class_name
-        if module == 'numpy' or module.startswith('numpy.'):
-            return super().find_class(module, name)
         if match_class_name is not None:
             return self.safe_classes[match_class_name]
         # 如果尝试加载未授权的类，则抛出异常