350 lines
12 KiB
Python
350 lines
12 KiB
Python
"""受理材料的校验、恶意文件扫描、解析与对象存储编排。"""
|
|
|
|
import hashlib
|
|
import io
|
|
import socket
|
|
import struct
|
|
import zipfile
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
from django.conf import settings
|
|
|
|
from apps.tickets.models import TicketDraftMaterial
|
|
from apps.tickets.storage import DraftImageError, get_object, prepare_image
|
|
|
|
SUPPORTED_EXTENSIONS = (
|
|
"jpg",
|
|
"jpeg",
|
|
"png",
|
|
"webp",
|
|
"pdf",
|
|
"txt",
|
|
"log",
|
|
"docx",
|
|
"xlsx",
|
|
"pptx",
|
|
)
|
|
IMAGE_EXTENSIONS = {"jpg", "jpeg", "png", "webp"}
|
|
GENERIC_MIME_TYPES = {"", "application/octet-stream", "binary/octet-stream"}
|
|
MIME_TYPES = {
|
|
"jpg": {"image/jpeg"},
|
|
"jpeg": {"image/jpeg"},
|
|
"png": {"image/png"},
|
|
"webp": {"image/webp"},
|
|
"pdf": {"application/pdf"},
|
|
"txt": {"text/plain"},
|
|
"log": {"text/plain", "text/x-log"},
|
|
"docx": {"application/vnd.openxmlformats-officedocument.wordprocessingml.document"},
|
|
"xlsx": {"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"},
|
|
"pptx": {"application/vnd.openxmlformats-officedocument.presentationml.presentation"},
|
|
}
|
|
CANONICAL_MIME_TYPES = {
|
|
"jpg": "image/jpeg",
|
|
"jpeg": "image/jpeg",
|
|
"png": "image/png",
|
|
"webp": "image/webp",
|
|
"pdf": "application/pdf",
|
|
"txt": "text/plain",
|
|
"log": "text/plain",
|
|
"docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
|
"xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
|
"pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
|
}
|
|
MATERIAL_TYPES = {
|
|
**{extension: "image" for extension in IMAGE_EXTENSIONS},
|
|
"pdf": "document",
|
|
"txt": "text",
|
|
"log": "log",
|
|
"docx": "document",
|
|
"xlsx": "spreadsheet",
|
|
"pptx": "presentation",
|
|
}
|
|
OFFICE_SIGNATURES = {
|
|
"docx": "word/document.xml",
|
|
"xlsx": "xl/workbook.xml",
|
|
"pptx": "ppt/presentation.xml",
|
|
}
|
|
|
|
|
|
class MaterialValidationError(ValueError):
|
|
def __init__(self, status_code: int, message: str):
|
|
super().__init__(message)
|
|
self.status_code = status_code
|
|
|
|
|
|
class MaterialScanError(RuntimeError):
|
|
pass
|
|
|
|
|
|
class MaterialParseError(RuntimeError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PreparedMaterial:
|
|
data: bytes
|
|
file_name: str
|
|
material_type: str
|
|
extension: str
|
|
claimed_mime_type: str
|
|
mime_type: str
|
|
sha256: str
|
|
|
|
|
|
def _decode_text(raw: bytes) -> str:
|
|
if b"\x00" in raw:
|
|
raise MaterialValidationError(415, "文本或日志材料包含二进制内容")
|
|
for encoding in ("utf-8-sig", "gb18030"):
|
|
try:
|
|
return raw.decode(encoding)
|
|
except UnicodeDecodeError:
|
|
continue
|
|
raise MaterialValidationError(415, "文本或日志材料编码无法识别")
|
|
|
|
|
|
def _validate_office_signature(raw: bytes, extension: str) -> None:
|
|
try:
|
|
with zipfile.ZipFile(io.BytesIO(raw)) as archive:
|
|
names = set(archive.namelist())
|
|
except (OSError, zipfile.BadZipFile) as error:
|
|
raise MaterialValidationError(415, "Office 文件签名无效") from error
|
|
if OFFICE_SIGNATURES[extension] not in names:
|
|
raise MaterialValidationError(415, f"文件内容与 .{extension} 扩展名不匹配")
|
|
|
|
|
|
def prepare_material(
|
|
raw: bytes,
|
|
original_name: str,
|
|
claimed_mime_type: str = "",
|
|
) -> PreparedMaterial:
|
|
"""校验大小、扩展名、MIME 和文件签名，返回可安全落库的材料元数据。"""
|
|
|
|
if not raw:
|
|
raise MaterialValidationError(400, "不能上传空文件")
|
|
max_bytes = settings.TICKET_MATERIAL_MAX_FILE_BYTES
|
|
if len(raw) > max_bytes:
|
|
raise MaterialValidationError(413, f"单个材料不能超过 {max_bytes // 1024 // 1024} MB")
|
|
|
|
safe_name = (original_name or "material").replace("\\", "/").rsplit("/", 1)[-1]
|
|
extension = Path(safe_name).suffix.lower().lstrip(".")
|
|
if extension not in SUPPORTED_EXTENSIONS:
|
|
supported = "、".join(item.upper() for item in SUPPORTED_EXTENSIONS)
|
|
raise MaterialValidationError(415, f"仅支持 {supported}")
|
|
|
|
claimed = (claimed_mime_type or "").split(";", 1)[0].strip().lower()
|
|
if claimed not in GENERIC_MIME_TYPES and claimed not in MIME_TYPES[extension]:
|
|
raise MaterialValidationError(415, "文件 MIME 类型与扩展名不匹配")
|
|
|
|
if extension in IMAGE_EXTENSIONS:
|
|
try:
|
|
image = prepare_image(raw, safe_name)
|
|
except DraftImageError as error:
|
|
raise MaterialValidationError(error.status_code, str(error)) from error
|
|
detected_group = {"jpg", "jpeg"} if image.extension == "jpg" else {image.extension}
|
|
if extension not in detected_group:
|
|
raise MaterialValidationError(415, "图片内容与扩展名不匹配")
|
|
return PreparedMaterial(
|
|
data=raw,
|
|
file_name=safe_name,
|
|
material_type="image",
|
|
extension=extension,
|
|
claimed_mime_type=claimed,
|
|
mime_type=image.mime_type,
|
|
sha256=hashlib.sha256(raw).hexdigest(),
|
|
)
|
|
|
|
if extension == "pdf":
|
|
if not raw.startswith(b"%PDF-"):
|
|
raise MaterialValidationError(415, "PDF 文件签名无效")
|
|
elif extension in {"txt", "log"}:
|
|
_decode_text(raw)
|
|
else:
|
|
_validate_office_signature(raw, extension)
|
|
|
|
return PreparedMaterial(
|
|
data=raw,
|
|
file_name=safe_name,
|
|
material_type=MATERIAL_TYPES[extension],
|
|
extension=extension,
|
|
claimed_mime_type=claimed,
|
|
mime_type=CANONICAL_MIME_TYPES[extension],
|
|
sha256=hashlib.sha256(raw).hexdigest(),
|
|
)
|
|
|
|
|
|
def scan_bytes(raw: bytes) -> dict:
|
|
"""调用配置的扫描器；扫描不可用或响应异常时一律按失败处理。"""
|
|
|
|
backend = settings.TICKET_MATERIAL_SCANNER_BACKEND
|
|
if backend == "test":
|
|
return {"status": "clean", "engine": "test"}
|
|
if backend != "clamav":
|
|
raise MaterialScanError("未配置可用的恶意文件扫描器")
|
|
|
|
try:
|
|
with socket.create_connection(
|
|
(settings.CLAMAV_HOST, settings.CLAMAV_PORT),
|
|
timeout=settings.CLAMAV_TIMEOUT_SECONDS,
|
|
) as connection:
|
|
connection.sendall(b"zINSTREAM\0")
|
|
for offset in range(0, len(raw), 1024 * 1024):
|
|
chunk = raw[offset : offset + 1024 * 1024]
|
|
connection.sendall(struct.pack(">I", len(chunk)))
|
|
connection.sendall(chunk)
|
|
connection.sendall(struct.pack(">I", 0))
|
|
response = b""
|
|
while not response.endswith(b"\0"):
|
|
block = connection.recv(4096)
|
|
if not block:
|
|
break
|
|
response += block
|
|
except OSError as error:
|
|
raise MaterialScanError("恶意文件扫描服务暂时不可用") from error
|
|
|
|
message = response.rstrip(b"\0").decode("utf-8", errors="replace")
|
|
if message.endswith(" OK"):
|
|
return {"status": "clean", "engine": "clamav", "message": message}
|
|
if message.endswith(" FOUND"):
|
|
signature = message.rsplit(": ", 1)[-1].removesuffix(" FOUND")
|
|
return {
|
|
"status": "infected",
|
|
"engine": "clamav",
|
|
"signature": signature,
|
|
"message": message,
|
|
}
|
|
raise MaterialScanError(f"恶意文件扫描未完成：{message or '无响应'}")
|
|
|
|
|
|
def _parse_pdf(raw: bytes) -> str:
|
|
from pypdf import PdfReader
|
|
|
|
reader = PdfReader(io.BytesIO(raw))
|
|
return "\n\n".join((page.extract_text() or "").strip() for page in reader.pages).strip()
|
|
|
|
|
|
def _parse_docx(raw: bytes) -> str:
|
|
from docx import Document
|
|
|
|
document = Document(io.BytesIO(raw))
|
|
return "\n".join(paragraph.text.strip() for paragraph in document.paragraphs if paragraph.text)
|
|
|
|
|
|
def _parse_xlsx(raw: bytes) -> str:
|
|
from openpyxl import load_workbook
|
|
|
|
workbook = load_workbook(io.BytesIO(raw), read_only=True, data_only=True)
|
|
sections = []
|
|
try:
|
|
for sheet in workbook.worksheets:
|
|
rows = []
|
|
for row in sheet.iter_rows(values_only=True):
|
|
values = [str(value).strip() for value in row if value not in (None, "")]
|
|
if values:
|
|
rows.append("\t".join(values))
|
|
if rows:
|
|
sections.append(f"[工作表：{sheet.title}]\n" + "\n".join(rows))
|
|
finally:
|
|
workbook.close()
|
|
return "\n\n".join(sections)
|
|
|
|
|
|
def _parse_pptx(raw: bytes) -> str:
|
|
from pptx import Presentation
|
|
|
|
presentation = Presentation(io.BytesIO(raw))
|
|
slides = []
|
|
for index, slide in enumerate(presentation.slides, start=1):
|
|
texts = [
|
|
shape.text.strip()
|
|
for shape in slide.shapes
|
|
if hasattr(shape, "text") and shape.text.strip()
|
|
]
|
|
if texts:
|
|
slides.append(f"[幻灯片 {index}]\n" + "\n".join(texts))
|
|
return "\n\n".join(slides)
|
|
|
|
|
|
def parse_material(raw: bytes, extension: str) -> str:
|
|
"""按已校验格式提取文本，解析失败时不伪造空成功结果。"""
|
|
|
|
try:
|
|
if extension in {"txt", "log"}:
|
|
text = _decode_text(raw)
|
|
elif extension == "pdf":
|
|
text = _parse_pdf(raw)
|
|
elif extension == "docx":
|
|
text = _parse_docx(raw)
|
|
elif extension == "xlsx":
|
|
text = _parse_xlsx(raw)
|
|
elif extension == "pptx":
|
|
text = _parse_pptx(raw)
|
|
else:
|
|
return ""
|
|
except MaterialValidationError:
|
|
raise
|
|
except Exception as error:
|
|
raise MaterialParseError("材料内容解析失败") from error
|
|
return text[: settings.TICKET_MATERIAL_MAX_PARSED_CHARS].strip()
|
|
|
|
|
|
def process_material(material: TicketDraftMaterial, raw: bytes | None = None) -> None:
|
|
"""扫描并解析单个材料，分别持久化扫描与解析终态。"""
|
|
|
|
raw = raw if raw is not None else get_object(material.storage_path)
|
|
material.scan_status = "scanning"
|
|
material.scan_result = {}
|
|
material.parse_status = "pending"
|
|
material.parse_error = ""
|
|
material.save(
|
|
update_fields=["scan_status", "scan_result", "parse_status", "parse_error"]
|
|
)
|
|
try:
|
|
scan_result = scan_bytes(raw)
|
|
except Exception as error:
|
|
message = str(error) if isinstance(error, MaterialScanError) else "恶意文件扫描执行失败"
|
|
material.scan_status = "failed"
|
|
material.scan_result = {"status": "failed", "message": message}
|
|
material.parse_status = "blocked"
|
|
material.parse_error = "恶意文件扫描未通过，未执行解析"
|
|
material.save(
|
|
update_fields=["scan_status", "scan_result", "parse_status", "parse_error"]
|
|
)
|
|
return
|
|
|
|
material.scan_status = scan_result["status"]
|
|
material.scan_result = scan_result
|
|
if material.scan_status != "clean":
|
|
material.parse_status = "blocked"
|
|
material.parse_error = "材料被恶意文件扫描拦截，未执行解析"
|
|
material.save(
|
|
update_fields=["scan_status", "scan_result", "parse_status", "parse_error"]
|
|
)
|
|
return
|
|
|
|
if material.material_type == "image":
|
|
material.parse_status = "not_applicable"
|
|
material.parsed_text = ""
|
|
material.save(
|
|
update_fields=["scan_status", "scan_result", "parse_status", "parsed_text"]
|
|
)
|
|
return
|
|
|
|
try:
|
|
material.parsed_text = parse_material(raw, material.extension)
|
|
material.parse_status = "parsed"
|
|
material.parse_error = ""
|
|
except (MaterialParseError, MaterialValidationError) as error:
|
|
material.parsed_text = ""
|
|
material.parse_status = "failed"
|
|
material.parse_error = str(error)
|
|
material.save(
|
|
update_fields=[
|
|
"scan_status",
|
|
"scan_result",
|
|
"parse_status",
|
|
"parse_error",
|
|
"parsed_text",
|
|
]
|
|
)
|