lou
56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
"""
|
|
汇聚器模块 (assembler)
|
|
模块: 服务端 / 汇聚器
|
|
输入: 已收卷集合
|
|
输出: 完整密文文件
|
|
|
|
两道门: 卷齐才合并 (缺卷返回列表让客户端补); 合并后总哈希必须匹配才进解包。
|
|
"""
|
|
|
|
import hashlib
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
from task_manager import TaskManager
|
|
|
|
|
|
class AssemblerError(RuntimeError):
|
|
"""汇聚失败 (卷未齐/总哈希不一致)"""
|
|
|
|
|
|
class Assembler:
|
|
def __init__(self, tasks: TaskManager, tmp_root: str | Path) -> None:
|
|
self.tasks = tasks
|
|
self.tmp_root = Path(tmp_root)
|
|
|
|
def assemble(self, transfer_id: str) -> Path:
|
|
"""按 idx 顺序合并全部卷 -> 返回合并后密文路径"""
|
|
transfer = self.tasks.get(transfer_id)
|
|
if transfer is None:
|
|
raise AssemblerError(f"任务不存在: {transfer_id}")
|
|
total = transfer["chunk_count"]
|
|
|
|
received = self.tasks.received(transfer_id)
|
|
if len(received) < total:
|
|
missing = sorted(set(range(1, total + 1)) - received)
|
|
raise AssemblerError(f"缺卷 {missing}, 等齐再合并")
|
|
|
|
chunk_dir = self.tmp_root / transfer_id
|
|
merged_path = chunk_dir / "merged.bin"
|
|
with open(merged_path, "wb") as out:
|
|
for idx in range(1, total + 1):
|
|
chunk_path = chunk_dir / f"chunk_{idx:04d}"
|
|
with open(chunk_path, "rb") as f:
|
|
shutil.copyfileobj(f, out)
|
|
|
|
# 总哈希校验
|
|
h = hashlib.sha256()
|
|
with open(merged_path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(1 << 20), b""):
|
|
h.update(chunk)
|
|
if h.hexdigest() != transfer["total_sha256"]:
|
|
merged_path.unlink(missing_ok=True)
|
|
raise AssemblerError("合并后总 SHA-256 不一致, 任务 failed (需整体重传)")
|
|
|
|
return merged_path
|