refactor: improve code formatting and enhance XLSX validation logic

apps/chat/serializers/chat_authentication.py

@@ -15,7 +15,6 @@
 from rest_framework import serializers
 
 from application.models import ApplicationAccessToken, Application, ApplicationVersion
-from application.serializers.application import ApplicationSerializerModel
 from common.auth.common import FileToken, ChatToken
 from common.auth.constants.operate_constants import Operate
 from common.constants.authentication_type import AuthenticationType
@@ -38,7 +37,7 @@ def auth(self, request):
             # 校验token
             if token is not None:
                 token_details = signing.loads(token[7:])
-        except Exception as e:
+        except Exception:
             pass
         chat_user_id = token_details.get("id") or str(uuid.uuid7())
         _type = AuthenticationType.CHAT_USER
@@ -72,7 +71,7 @@ def auth(self, request, with_valid=True):
             # 校验token
             if token is not None:
                 token_details = signing.loads(token[7:])
-        except Exception as e:
+        except Exception:
             pass
         if with_valid:
             self.is_valid(raise_exception=True)
@@ -222,7 +221,14 @@ def profile(self, with_valid=True):
             node for node in ((application.work_flow or {}).get("nodes", []) or []) if node.get("id") == "base-node"
         ]
         return {
-            **ApplicationSerializerModel(application).data,
+            "id": application.id,
+            "name": application.name,
+            "desc": application.desc,
+            "prologue": application.prologue,
+            "icon": application.icon,
+            "type": application.type,
+            "dialogue_number": application.dialogue_number,
+            "problem_optimization": application.problem_optimization,
             "stt_model_id": application.stt_model_id,
             "tts_model_id": application.tts_model_id,
             "stt_model_enable": application.stt_model_enable,

apps/common/handle/impl/common_handle.py

@@ -1,14 +1,14 @@
 # coding=utf-8
 """
-    @project: MaxKB
-    @Author：虎
-    @file： tools.py
-    @date：2024/9/11 16:41
-    @desc:
+@project: MaxKB
+@Author：虎
+@file： tools.py
+@date：2024/9/11 16:41
+@desc:
 """
+
 import io
 import traceback
-from functools import reduce
 from io import BytesIO
 from xml.etree.ElementTree import fromstring
 from zipfile import ZipFile
@@ -23,8 +23,42 @@
 from knowledge.models import File
 
 from PIL import ImageFile
+
 ImageFile.LOAD_TRUNCATED_IMAGES = True
-PILImage.MAX_IMAGE_PIXELS = None
+
+# 全局图片解码像素上限（不再禁用 Pillow 的解压炸弹保护）。
+# 超过该上限 Pillow 会告警，超过 2 倍会直接抛错，避免超大图片耗尽 worker 内存。
+PILImage.MAX_IMAGE_PIXELS = 50_000_000
+
+# 内嵌图片解码保护（防解压炸弹 / 超大尺寸图片导致共享 worker OOM）。
+MAX_EMBED_IMAGE_PIXELS = 16_000_000
+MAX_EMBED_IMAGE_AGGREGATE_PIXELS = 64_000_000
+
+# XLSX(zip) 压缩包防护，限制成员数 / 解压后总大小 / 解压膨胀比。
+MAX_EMBED_ARCHIVE_MEMBERS = 10_000
+MAX_EMBED_ARCHIVE_UNCOMPRESSED_BYTES = 1024 * 1024 * 1024
+MAX_EMBED_ARCHIVE_EXPANSION_RATIO = 50
+
+
+def validate_xlsx_archive(archive: ZipFile):
+    infolist = archive.infolist()
+    if len(infolist) > MAX_EMBED_ARCHIVE_MEMBERS:
+        raise ValueError(f"XLSX archive member count exceeds limit: {len(infolist)}")
+    total_uncompressed = sum(info.file_size for info in infolist)
+    total_compressed = sum(info.compress_size for info in infolist)
+    if total_uncompressed > MAX_EMBED_ARCHIVE_UNCOMPRESSED_BYTES:
+        raise ValueError("XLSX archive uncompressed size exceeds limit")
+    if total_compressed > 0 and total_uncompressed > total_compressed * MAX_EMBED_ARCHIVE_EXPANSION_RATIO:
+        raise ValueError("XLSX archive expansion ratio exceeds limit")
+
+
+def validate_xlsx_buffer(buffer):
+    archive = ZipFile(buffer)
+    try:
+        validate_xlsx_archive(archive)
+    finally:
+        archive.close()
+
 
 def parse_element(element) -> {}:
     data = {}
@@ -87,15 +121,16 @@ def handle_images(deps, archive: ZipFile) -> []:
 
 def xlsx_embed_cells_images(buffer) -> {}:
     archive = ZipFile(buffer)
+    validate_xlsx_archive(archive)
     # 解析cellImage.xml文件
     deps = get_dependents(archive, get_rels_path("xl/cellimages.xml"))
     image_rel = handle_images(deps=deps, archive=archive)
     # 工作表及其中图片ID
     sheet_list = {}
     for item in archive.namelist():
-        if not item.startswith('xl/worksheets/sheet'):
+        if not item.startswith("xl/worksheets/sheet"):
             continue
-        key = item.split('/')[-1].split('.')[0].split('sheet')[-1]
+        key = item.split("/")[-1].split(".")[0].split("sheet")[-1]
         sheet_list[key] = parse_element_sheet_xml(fromstring(archive.read(item)))
     cell_images_xml = parse_element(fromstring(archive.read("xl/cellimages.xml")))
     cell_images_rel = {}
@@ -104,28 +139,36 @@ def xlsx_embed_cells_images(buffer) -> {}:
     for cnv, embed in cell_images_xml.items():
         cell_images_xml[cnv] = cell_images_rel.get(embed)
     result = {}
+    total_pixels = 0
     for key, img in cell_images_xml.items():
-        all_cells = [
-            cell
-            for _sheet_id, sheet in sheet_list.items()
-            if sheet is not None
-            for cell in sheet or []
-        ]
-
-        image_excel_id_list = [
-            cell for cell in all_cells
-            if isinstance(cell, str) and key in cell
-        ]
+        all_cells = [cell for _sheet_id, sheet in sheet_list.items() if sheet is not None for cell in sheet or []]
+
+        image_excel_id_list = [cell for cell in all_cells if isinstance(cell, str) and key in cell]
         # print(key, img)
         if img is None:
             continue
         if len(image_excel_id_list) > 0:
             image_excel_id = image_excel_id_list[-1]
             f = archive.open(img.target)
             img_byte = io.BytesIO()
-            im = PILImage.open(f).convert('RGB')
-            im.save(img_byte, format='JPEG')
-            image = File(id=uuid.uuid7(), file_name=img.path, meta={'debug': False, 'content': img_byte.getvalue()})
-            result['=' + image_excel_id] = image
+            try:
+                with PILImage.open(f) as im:
+                    width, height = im.size
+                    pixels = width * height
+                    if pixels > MAX_EMBED_IMAGE_PIXELS:
+                        maxkb_logger.warning(
+                            f"Skip oversized embedded image {img.path}: {width}x{height} pixels exceeds limit"
+                        )
+                        continue
+                    total_pixels += pixels
+                    if total_pixels > MAX_EMBED_IMAGE_AGGREGATE_PIXELS:
+                        maxkb_logger.warning("Skip embedded images in archive: aggregate pixels exceed limit")
+                        break
+                    im.convert("RGB").save(img_byte, format="JPEG")
+            except Exception as e:
+                maxkb_logger.error(f"Error decoding image {img.target}: {e}, {traceback.format_exc()}")
+                continue
+            image = File(id=uuid.uuid7(), file_name=img.path, meta={"debug": False, "content": img_byte.getvalue()})
+            result["=" + image_excel_id] = image
     archive.close()
     return result

apps/common/handle/impl/qa/xlsx_parse_qa_handle.py

@@ -1,18 +1,19 @@
 # coding=utf-8
 """
-    @project: maxkb
-    @Author：虎
-    @file： xlsx_parse_qa_handle.py
-    @date：2024/5/21 14:59
-    @desc:
+@project: maxkb
+@Author：虎
+@file： xlsx_parse_qa_handle.py
+@date：2024/5/21 14:59
+@desc:
 """
+
 import io
 import traceback
 
 import openpyxl
 
 from common.handle.base_parse_qa_handle import BaseParseQAHandle, get_title_row_index_dict, get_row_value
-from common.handle.impl.common_handle import xlsx_embed_cells_images
+from common.handle.impl.common_handle import xlsx_embed_cells_images, validate_xlsx_buffer
 from common.utils.logger import maxkb_logger
 
 
@@ -22,28 +23,26 @@ def handle_sheet(file_name, sheet, image_dict):
         title_row_list = next(rows)
         title_row_list = [row.value for row in title_row_list]
     except Exception as e:
-        return {'name': file_name, 'paragraphs': []}
+        return {"name": file_name, "paragraphs": []}
     if len(title_row_list) == 0:
-        return {'name': file_name, 'paragraphs': []}
+        return {"name": file_name, "paragraphs": []}
     title_row_index_dict = get_title_row_index_dict(title_row_list)
     paragraph_list = []
     for row in rows:
-        content = get_row_value(row, title_row_index_dict, 'content')
+        content = get_row_value(row, title_row_index_dict, "content")
         if content is None or content.value is None:
             continue
-        problem = get_row_value(row, title_row_index_dict, 'problem_list')
-        problem = str(problem.value) if problem is not None and problem.value is not None else ''
-        problem_list = [{'content': p[0:255]} for p in problem.split('\n') if len(p.strip()) > 0]
-        title = get_row_value(row, title_row_index_dict, 'title')
-        title = str(title.value) if title is not None and title.value is not None else ''
+        problem = get_row_value(row, title_row_index_dict, "problem_list")
+        problem = str(problem.value) if problem is not None and problem.value is not None else ""
+        problem_list = [{"content": p[0:255]} for p in problem.split("\n") if len(p.strip()) > 0]
+        title = get_row_value(row, title_row_index_dict, "title")
+        title = str(title.value) if title is not None and title.value is not None else ""
         content = str(content.value)
         image = image_dict.get(content, None)
         if image is not None:
-            content = f'![](./oss/file/{image.id})'
-        paragraph_list.append({'title': title[0:255],
-                               'content': content[0:102400],
-                               'problem_list': problem_list})
-    return {'name': file_name, 'paragraphs': paragraph_list}
+            content = f"![](./oss/file/{image.id})"
+        paragraph_list.append({"title": title[0:255], "content": content[0:102400], "problem_list": problem_list})
+    return {"name": file_name, "paragraphs": paragraph_list}
 
 
 class XlsxParseQAHandle(BaseParseQAHandle):
@@ -56,6 +55,7 @@ def support(self, file, get_buffer):
     def handle(self, file, get_buffer, save_image):
         buffer = get_buffer(file)
         try:
+            validate_xlsx_buffer(io.BytesIO(buffer))
             workbook = openpyxl.load_workbook(io.BytesIO(buffer))
             try:
                 image_dict: dict = xlsx_embed_cells_images(io.BytesIO(buffer))
@@ -64,12 +64,16 @@ def handle(self, file, get_buffer, save_image):
                 image_dict = {}
             worksheets = workbook.worksheets
             worksheets_size = len(worksheets)
-            return [row for row in
-                    [handle_sheet(file.name,
-                                  sheet,
-                                  image_dict) if worksheets_size == 1 and sheet.title == 'Sheet1' else handle_sheet(
-                        sheet.title, sheet, image_dict) for sheet
-                     in worksheets] if row is not None]
+            return [
+                row
+                for row in [
+                    handle_sheet(file.name, sheet, image_dict)
+                    if worksheets_size == 1 and sheet.title == "Sheet1"
+                    else handle_sheet(sheet.title, sheet, image_dict)
+                    for sheet in worksheets
+                ]
+                if row is not None
+            ]
         except Exception as e:
             maxkb_logger.error(f"Error processing XLSX file {file.name}: {e}, {traceback.format_exc()}")
-            return [{'name': file.name, 'paragraphs': []}]
+            return [{"name": file.name, "paragraphs": []}]

apps/common/handle/impl/table/xlsx_parse_table_handle.py

@@ -6,15 +6,15 @@
 from openpyxl import load_workbook
 
 from common.handle.base_parse_table_handle import BaseParseTableHandle
-from common.handle.impl.common_handle import xlsx_embed_cells_images
+from common.handle.impl.common_handle import xlsx_embed_cells_images, validate_xlsx_buffer
 from common.handle.impl.xlsx_utils import iter_sheet_content_rows
 from common.utils.logger import maxkb_logger
 
 
 class XlsxParseTableHandle(BaseParseTableHandle):
     def support(self, file, get_buffer):
         file_name: str = file.name.lower()
-        if file_name.endswith('.xlsx'):
+        if file_name.endswith(".xlsx"):
             return True
         return False
 
@@ -30,7 +30,7 @@ def fill_merged_cells(self, sheet, image_dict):
             return data
         for idx, cell in enumerate(title_row):
             if cell.value is None:
-                headers.append(' ' * (idx + 1))
+                headers.append(" " * (idx + 1))
             else:
                 headers.append(cell.value)
 
@@ -47,10 +47,10 @@ def fill_merged_cells(self, sheet, image_dict):
                             cell_value = sheet[merged_range.min_row][merged_range.min_col - 1].value
                             break
                 if cell_value is None:
-                    cell_value = ''
+                    cell_value = ""
                 image = image_dict.get(cell_value, None)
                 if image is not None:
-                    cell_value = f'![](./oss/file/{image.id})'
+                    cell_value = f"![](./oss/file/{image.id})"
 
                 # 使用标题作为键，单元格的值作为值存入字典
                 row_data[headers[col_idx]] = cell_value
@@ -61,6 +61,7 @@ def fill_merged_cells(self, sheet, image_dict):
     def handle(self, file, get_buffer, save_image):
         buffer = get_buffer(file)
         try:
+            validate_xlsx_buffer(io.BytesIO(buffer))
             wb = load_workbook(io.BytesIO(buffer))
             try:
                 image_dict: dict = xlsx_embed_cells_images(io.BytesIO(buffer))
@@ -76,13 +77,13 @@ def handle(self, file, get_buffer, save_image):
                 for row in data:
                     row_output = "; ".join([f"{key}: {value}" for key, value in row.items()])
                     # print(row_output)
-                    paragraphs.append({'title': '', 'content': row_output})
+                    paragraphs.append({"title": "", "content": row_output})
 
-                result.append({'name': sheetname, 'paragraphs': paragraphs})
+                result.append({"name": sheetname, "paragraphs": paragraphs})
 
         except BaseException as e:
             maxkb_logger.error(f"Error processing XLSX file {file.name}: {e}, {traceback.format_exc()}")
-            return [{'name': file.name, 'paragraphs': []}]
+            return [{"name": file.name, "paragraphs": []}]
         return result
 
     def get_content(self, file, save_image):
@@ -94,9 +95,9 @@ def get_content(self, file, save_image):
                 if len(image_dict) > 0:
                     save_image(image_dict.values())
             except Exception as e:
-                maxkb_logger.error(f'Exception: {e}')
+                maxkb_logger.error(f"Exception: {e}")
                 image_dict = {}
-            md_tables = ''
+            md_tables = ""
             # 遍历所有工作表
             for sheetname in workbook.sheetnames:
                 sheet = workbook[sheetname]
@@ -105,22 +106,25 @@ def get_content(self, file, save_image):
                     continue
 
                 # 添加 sheet 名称作为标题
-                md_tables += f'## {sheetname}\n\n'
+                md_tables += f"## {sheetname}\n\n"
 
                 # 提取表头和内容
                 headers = [f"{key}" for key, value in rows[0].items()]
 
                 # 构建 Markdown 表格
-                md_table = '| ' + ' | '.join(headers) + ' |\n'
-                md_table += '| ' + ' | '.join(['---'] * len(headers)) + ' |\n'
+                md_table = "| " + " | ".join(headers) + " |\n"
+                md_table += "| " + " | ".join(["---"] * len(headers)) + " |\n"
                 for row in rows:
-                    r = [f'{value}' for key, value in row.items()]
-                    md_table += '| ' + ' | '.join(
-                        [str(cell).replace('\n', '<br>') if cell is not None else '' for cell in r]) + ' |\n'
+                    r = [f"{value}" for key, value in row.items()]
+                    md_table += (
+                        "| "
+                        + " | ".join([str(cell).replace("\n", "<br>") if cell is not None else "" for cell in r])
+                        + " |\n"
+                    )
 
-                md_tables += md_table + '\n\n'
+                md_tables += md_table + "\n\n"
 
             return md_tables
         except Exception as e:
-            maxkb_logger.error(f'excel split handle error: {e}')
-            return f'error: {e}'
+            maxkb_logger.error(f"excel split handle error: {e}")
+            return f"error: {e}"