merge to resolve multiple pickle security issues (#1728)

* 注释调试if分支

* support pdf url for latex translation

* Merge pull request from GHSA-mvrw-h7rc-22r8

* 注释调试if分支

* Improve objload security

* Update README.md

* support pdf url for latex translation

---------

Co-authored-by: binary-husky <96192199+binary-husky@users.noreply.github.com>
Co-authored-by: binary-husky <qingxu.fu@outlook.com>

* fix import

---------

Co-authored-by: Longtaotao <longtaotao@bupt.edu.cn>
Co-authored-by: iluem <57590186+Qhaoduoyu@users.noreply.github.com>

crazy_functions/Latex输出PDF.py

@@ -107,6 +107,10 @@ def arxiv_download(chatbot, history, txt, allow_cache=True):
         except ValueError:
             return False
 
+    if txt.startswith('https://arxiv.org/pdf/'):
+        arxiv_id = txt.split('/')[-1]   # 2402.14207v2.pdf
+        txt = arxiv_id.split('v')[0]  # 2402.14207
+
     if ('.' in txt) and ('/' not in txt) and is_float(txt):  # is arxiv ID
         txt = 'https://arxiv.org/abs/' + txt.strip()
     if ('.' in txt) and ('/' not in txt) and is_float(txt[:10]):  # is arxiv ID
@@ -121,6 +125,7 @@ def arxiv_download(chatbot, history, txt, allow_cache=True):
     time.sleep(1)  # 刷新界面
 
     url_ = txt  # https://arxiv.org/abs/1707.06690
+
     if not txt.startswith('https://arxiv.org/abs/'):
         msg = f"解析arxiv网址失败, 期望格式例如: https://arxiv.org/abs/1707.06690。实际得到格式: {url_}。"
         yield from update_ui_lastest_msg(msg, chatbot=chatbot, history=history)  # 刷新界面
@@ -458,23 +463,23 @@ def PDF翻译中文并重新编译PDF(txt, llm_kwargs, plugin_kwargs, chatbot, h
             promote_file_to_downloadzone(translate_pdf, rename_file=None, chatbot=chatbot)
 
             comparison_pdf = [f for f in glob.glob(f'{project_folder}/**/comparison.pdf', recursive=True)][0]
-            promote_file_to_downloadzone(comparison_pdf, rename_file=None, chatbot=chatbot)        
+            promote_file_to_downloadzone(comparison_pdf, rename_file=None, chatbot=chatbot)
 
             zip_res = zip_result(project_folder)
             promote_file_to_downloadzone(file=zip_res, chatbot=chatbot)
 
             return True
-        
+
         except:
             report_exception(chatbot, history, a=f"解析项目: {txt}", b=f"发现重复上传，但是无法找到相关文件")
             yield from update_ui(chatbot=chatbot, history=history)
-            
+
             chatbot.append([f"没有相关文件", '尝试重新翻译PDF...'])
             yield from update_ui(chatbot=chatbot, history=history)
 
             except_flag = True
-            
-    
+
+
     elif not repeat or except_flag:
         yield from update_ui_lastest_msg(f"未发现重复上传", chatbot=chatbot, history=history)
 

crazy_functions/latex_fns/latex_actions.py

@@ -1,10 +1,11 @@
 from toolbox import update_ui, update_ui_lastest_msg, get_log_folder
-from toolbox import get_conf, objdump, objload, promote_file_to_downloadzone
+from toolbox import get_conf, promote_file_to_downloadzone
 from .latex_toolbox import PRESERVE, TRANSFORM
 from .latex_toolbox import set_forbidden_text, set_forbidden_text_begin_end, set_forbidden_text_careful_brace
 from .latex_toolbox import reverse_forbidden_text_careful_brace, reverse_forbidden_text, convert_to_linklist, post_process
 from .latex_toolbox import fix_content, find_main_tex_file, merge_tex_files, compile_latex_with_timeout
 from .latex_toolbox import find_title_and_abs
+from .latex_pickle_io import objdump, objload
 
 import os, shutil
 import re

crazy_functions/latex_fns/latex_pickle_io.py

@@ -0,0 +1,38 @@
+import pickle
+
+
+class SafeUnpickler(pickle.Unpickler):
+
+    def get_safe_classes(self):
+        from .latex_actions import LatexPaperFileGroup, LatexPaperSplit
+        # 定义允许的安全类
+        safe_classes = {
+            # 在这里添加其他安全的类
+            'LatexPaperFileGroup': LatexPaperFileGroup,
+            'LatexPaperSplit' : LatexPaperSplit,
+        }
+        return safe_classes
+
+    def find_class(self, module, name):
+        # 只允许特定的类进行反序列化
+        self.safe_classes = self.get_safe_classes()
+        if f'{module}.{name}' in self.safe_classes:
+            return self.safe_classes[f'{module}.{name}']
+        # 如果尝试加载未授权的类，则抛出异常
+        raise pickle.UnpicklingError(f"Attempted to deserialize unauthorized class '{name}' from module '{module}'")
+
+def objdump(obj, file="objdump.tmp"):
+
+    with open(file, "wb+") as f:
+        pickle.dump(obj, f)
+    return
+
+
+def objload(file="objdump.tmp"):
+    import os
+
+    if not os.path.exists(file):
+        return
+    with open(file, "rb") as f:
+        unpickler = SafeUnpickler(f)
+        return unpickler.load()

toolbox.py

@@ -868,23 +868,6 @@ class ProxyNetworkActivate:
         return
 
 
-def objdump(obj, file="objdump.tmp"):
-    import pickle
-
-    with open(file, "wb+") as f:
-        pickle.dump(obj, f)
-    return
-
-
-def objload(file="objdump.tmp"):
-    import pickle, os
-
-    if not os.path.exists(file):
-        return
-    with open(file, "rb") as f:
-        return pickle.load(f)
-
-
 def Singleton(cls):
     """
     一个单实例装饰器