"""电子书存储：文件放 data/ebooks/{textbook_id}/ 下，数据库只存文件名与元信息。 上传按分块写临时文件（不整体进内存），先校验扩展名与文件头魔数， 再原子改名，避免半截文件污染在线阅读。 """ from __future__ import annotations import hashlib import os import re import shutil import uuid import zipfile from datetime import datetime from pathlib import Path from xml.etree import ElementTree from fastapi import HTTPException, UploadFile from sqlalchemy.orm import Session from pypdf import PdfReader from database import REPO_ROOT from models import EbookProgress, Textbook EBOOK_ROOT = REPO_ROOT / "data" / "ebooks" MAX_EBOOK_BYTES = 500 * 1024 * 1024 # 500MB _CHUNK = 1024 * 1024 # 扩展名 -> (格式标识, 魔数, MIME) KINDS: dict[str, tuple[str, bytes]] = { ".pdf": ("pdf", b"%PDF-"), ".epub": ("epub", b"PK\x03\x04"), } MIME_BY_FORMAT = {"pdf": "application/pdf", "epub": "application/epub+zip"} def book_dir(textbook_id: int) -> Path: return EBOOK_ROOT / str(textbook_id) def resolve_path(textbook_id: int, filename: str) -> Path: """只允许该书目录下的裸文件名，杜绝 ../ 穿越。""" safe = os.path.basename(filename) if not safe or safe.startswith("."): raise HTTPException(status_code=404, detail="电子书文件不存在") path = (book_dir(textbook_id) / safe).resolve() root = book_dir(textbook_id).resolve() if path.parent != root or not path.is_file(): raise HTTPException(status_code=404, detail="电子书文件不存在") return path def media_type(fmt: str) -> str: return MIME_BY_FORMAT.get(fmt, "application/octet-stream") def file_kind(filename: str) -> tuple[str, str]: """按扩展名给出 (格式, 魔数)；不支持的扩展名直接拒绝。""" ext = Path(filename or "").suffix.lower() kind = KINDS.get(ext) if kind is None: raise HTTPException(status_code=400, detail="只支持 PDF 与 EPUB 两种格式") return kind[0], kind[1] def _safe_name(filename: str) -> str: """原始文件名只用于展示：去掉路径与可能破坏排版的控制字符。""" name = os.path.basename(filename or "").strip() name = re.sub(r"[\r\n\t]+", " ", name) return name[:180] async def store_upload(db: Session, textbook: Textbook, upload: UploadFile) -> Textbook: fmt, magic = file_kind(upload.filename or "") target_dir = book_dir(textbook.id) target_dir.mkdir(parents=True, exist_ok=True) temp = target_dir / f".upload-{uuid.uuid4().hex}.part" size = 0 try: with temp.open("wb") as handle: while True: chunk = await upload.read(_CHUNK) if not chunk: break if size == 0 and not chunk.startswith(magic): raise HTTPException( status_code=400, detail="文件内容与扩展名不符，请上传真实的 PDF / EPUB 文件" ) size += len(chunk) if size > MAX_EBOOK_BYTES: raise HTTPException(status_code=400, detail="电子书超过 500MB 上限") handle.write(chunk) if size == 0: raise HTTPException(status_code=400, detail="上传的文件是空的") final = target_dir / f"{uuid.uuid4().hex}.{fmt}" os.replace(temp, final) finally: if temp.exists(): temp.unlink(missing_ok=True) await upload.close() # 覆盖式上传：旧文件随后清理，磁盘上不留重复副本 if textbook.ebook_file: try: resolve_path(textbook.id, textbook.ebook_file).unlink(missing_ok=True) except HTTPException: pass textbook.ebook_file = final.name textbook.ebook_format = fmt textbook.ebook_name = _safe_name(upload.filename or final.name) textbook.ebook_size = size textbook.ebook_pages = count_pages(final, fmt) textbook.ebook_uploaded_at = datetime.utcnow() db.commit() return textbook def clear_ebook(db: Session, textbook: Textbook) -> Textbook: if textbook.ebook_file: try: resolve_path(textbook.id, textbook.ebook_file).unlink(missing_ok=True) except HTTPException: pass textbook.ebook_file = "" textbook.ebook_name = "" textbook.ebook_format = "" textbook.ebook_size = 0 textbook.ebook_pages = 0 textbook.ebook_uploaded_at = None db.commit() return textbook def sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as reader: for chunk in iter(lambda: reader.read(_CHUNK), b""): digest.update(chunk) return digest.hexdigest() def find_identical(textbook_id: int, source: Path) -> Path | None: """该书目录下是否已经有一模一样的文件（按大小 + SHA-256 判定）。""" target_dir = book_dir(textbook_id) if not target_dir.is_dir() or not source.is_file(): return None size = source.stat().st_size digest = sha256_file(source) for candidate in sorted(target_dir.iterdir()): if not candidate.is_file() or candidate.stat().st_size != size: continue if sha256_file(candidate) == digest: return candidate return None def adopt_file( db: Session, textbook: Textbook, filename: str, display_name: str = "", ) -> Textbook: """登记该书目录下已经在磁盘上的文件（认领孤儿电子书，不重新复制）。""" path = resolve_path(textbook.id, filename) fmt, _ = file_kind(path.name) textbook.ebook_file = path.name textbook.ebook_format = fmt textbook.ebook_name = _safe_name(display_name or path.name) textbook.ebook_size = path.stat().st_size textbook.ebook_pages = count_pages(path, fmt) textbook.ebook_uploaded_at = datetime.utcnow() db.commit() return textbook def attach_local_file( db: Session, textbook: Textbook, source: Path, display_name: str = "", ) -> Textbook: """给种子脚本用：把仓库里的电子书复制进 data/ebooks/{id}/ 并登记。""" fmt, magic = file_kind(source.name) if not source.is_file(): raise FileNotFoundError(str(source)) target_dir = book_dir(textbook.id) target_dir.mkdir(parents=True, exist_ok=True) # 同一本书重复导入（例如换库重跑种子）时复用已有文件，避免多占几十 MB identical = find_identical(textbook.id, source) if identical is not None: return adopt_file(db, textbook, identical.name, display_name=display_name) final = target_dir / f"{uuid.uuid4().hex}.{fmt}" with source.open("rb") as reader, final.open("wb") as writer: head = reader.read(len(magic)) if not head.startswith(magic): final.unlink(missing_ok=True) raise ValueError(f"{source.name} 不是有效的 {fmt.upper()} 文件") writer.write(head) shutil.copyfileobj(reader, writer, _CHUNK) if textbook.ebook_file: try: resolve_path(textbook.id, textbook.ebook_file).unlink(missing_ok=True) except HTTPException: pass textbook.ebook_file = final.name textbook.ebook_format = fmt textbook.ebook_name = _safe_name(display_name or source.name) textbook.ebook_size = final.stat().st_size textbook.ebook_pages = count_pages(final, fmt) textbook.ebook_uploaded_at = datetime.utcnow() db.commit() return textbook def _pdf_pages(path: Path) -> int: return len(PdfReader(path).pages) def _epub_pages(path: Path) -> int: """EPUB 没有“页”，用正文文档数（spine 条目）作为篇幅。""" with zipfile.ZipFile(path) as archive: container = ElementTree.fromstring(archive.read("META-INF/container.xml")) opf = next( ( node.get("full-path") for node in container.iter() if node.tag.endswith("rootfile") ), "", ) if not opf: return 0 package = ElementTree.fromstring(archive.read(opf)) spine = next((node for node in package if node.tag.endswith("spine")), None) if spine is None: return 0 return len( [ item for item in spine if item.tag.endswith("itemref") and item.get("idref") ] ) def count_pages(path: Path, fmt: str) -> int: """解析篇幅；文件损坏时返回 0，交给阅读器再回填，不阻断上传。""" try: return _pdf_pages(path) if fmt == "pdf" else _epub_pages(path) except Exception: # noqa: BLE001 - 解析器版本/畸形文件都按“未知篇幅”处理 return 0 def backfill_pages(db: Session) -> None: """给存量数据补篇幅：上传时没解析过（或当时解析失败）的书。""" rows = ( db.query(Textbook) .filter(Textbook.ebook_file != "", Textbook.ebook_pages == 0) .all() ) changed = False for textbook in rows: try: path = resolve_path(textbook.id, textbook.ebook_file) except HTTPException: continue pages = count_pages(path, textbook.ebook_format) if pages and textbook.ebook_pages != pages: textbook.ebook_pages = pages changed = True if textbook.ebook_uploaded_at is None: textbook.ebook_uploaded_at = datetime.utcfromtimestamp( path.stat().st_mtime ) changed = True if changed: db.commit() def drop_book_files(textbook_id: int) -> None: shutil.rmtree(book_dir(textbook_id), ignore_errors=True) def drop_progress(db: Session, textbook_id: int) -> None: db.query(EbookProgress).filter(EbookProgress.textbook_id == textbook_id).delete()