290 lines
9.7 KiB
Python
290 lines
9.7 KiB
Python
"""电子书存储：文件放 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()
|