diff --git a/finance/exporting.py b/finance/exporting.py index 28e5e18..b08ecf2 100644 --- a/finance/exporting.py +++ b/finance/exporting.py @@ -1,7 +1,9 @@ import hashlib import hmac +import logging import os import secrets +import stat import warnings from datetime import timedelta from decimal import Decimal, InvalidOperation @@ -28,6 +30,58 @@ FORMULA_VERSION = "3" MAX_ROWS = 100_000 MAX_GROUPS = 200 +logger = logging.getLogger(__name__) + + +def _validated_export_directory(path, description): + """导出目录必须不存在或为真实目录；符号链接一律拒绝并保持文件原样。""" + try: + info = os.lstat(path) + except FileNotFoundError: + path.mkdir(parents=True, mode=0o700) + info = os.lstat(path) + if stat.S_ISLNK(info.st_mode) or not stat.S_ISDIR(info.st_mode): + raise ValueError(f"{description}必须是真实目录，不能是符号链接") + os.chmod(path, 0o700) + return path + + +def _validated_export_target(user_dir, job): + """目标文件若已存在必须是普通文件；符号链接或其他类型一律拒绝覆盖。""" + target = user_dir / f"{job.pk}.xlsx" + try: + info = os.lstat(target) + except FileNotFoundError: + return target + if stat.S_ISLNK(info.st_mode) or not stat.S_ISREG(info.st_mode): + raise ValueError("公式导出目标路径已存在且不是普通文件，已拒绝生成") + return target + + +def _unlink_created_file(path): + """只通过不跟随链接的目录句柄删除受控目录内的普通文件。 + + 返回 "removed"（已删除）、"missing"（本就不存在）或 "blocked"（无法安全删除， + 例如目录已被换成符号链接或权限不足），调用方据此保留失败任务与磁盘的一致性。 + """ + try: + dir_fd = os.open(path.parent, os.O_RDONLY | os.O_NOFOLLOW) + except OSError: + return "blocked" + try: + try: + info = os.lstat(path.name, dir_fd=dir_fd) + except FileNotFoundError: + return "missing" + if not stat.S_ISREG(info.st_mode): + return "blocked" + os.unlink(path.name, dir_fd=dir_fd) + return "removed" + except OSError: + return "blocked" + finally: + os.close(dir_fd) + def _excel_number(value): if value is None: @@ -594,41 +648,34 @@ def build_formula_workbook(ds, member, export_context, output): def process_formula_export(export_id): job = FormulaExport.objects.select_related("dataset", "user").get(pk=export_id) - temporary = None target = None + created = [] try: - root = Path(settings.MEDIA_ROOT) / "formula-exports" - if root.is_symlink(): - raise ValueError("公式导出目录不能是符号链接") - root.mkdir(parents=True, exist_ok=True, mode=0o700) - root.chmod(0o700) - target = root / str(job.user_id) / f"{job.pk}.xlsx" - target.parent.mkdir(parents=True, exist_ok=True, mode=0o700) - if target.parent.is_symlink() or target.is_symlink(): - raise ValueError("公式导出路径不能是符号链接") - target.parent.chmod(0o700) - temporary = target.with_name(f"{job.pk}.{secrets.token_hex(8)}.tmp") member = Membership.objects.get(user=job.user, organization=job.dataset.organization) if set(member.departments) != set(job.departments): raise PermissionError("科室授权已变更，请重新提交导出") + root = _validated_export_directory(Path(settings.MEDIA_ROOT) / "formula-exports", "公式导出目录") + user_dir = _validated_export_directory(root / str(job.user_id), "公式导出用户目录") + target = _validated_export_target(user_dir, job) + # 路径完整验证通过后才把目标登记为可清理文件；此前任何失败都不会删除磁盘文件。 + created.append(target) + temporary = target.with_name(f"{job.pk}.{secrets.token_hex(8)}.tmp") descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + created.append(temporary) with os.fdopen(descriptor, "wb") as file: stats = build_formula_workbook(job.dataset, member, job.context, file) - temporary.replace(target) + os.replace(temporary, target) + created = [target] job.path, job.status, job.error = str(target), "succeeded", "" job.finished_at = timezone.now() job.expires_at = timezone.now() + timedelta(hours=24) job.save(update_fields=["path", "status", "error", "finished_at", "expires_at"]) return stats except Exception as exc: - for path in (temporary, target): - try: - if path is not None and not path.is_symlink(): - path.unlink(missing_ok=True) - except OSError: - pass - job.status, job.error, job.path = "failed", str(exc)[:500], "" - job.finished_at, job.expires_at = timezone.now(), None + job.status, job.error, job.finished_at, job.expires_at = "failed", str(exc)[:500], timezone.now(), None + # 只清理本次登记且确认位于受控目录内的文件；目录被换成符号链接时拒绝删除。 + leftover = [path for path in created if _unlink_created_file(path) == "blocked"] + job.path = str(target) if target in leftover else "" job.save(update_fields=["status", "error", "path", "finished_at", "expires_at"]) raise diff --git a/tests/test_formula_export_reliability.py b/tests/test_formula_export_reliability.py index 0530117..9e15b8d 100644 --- a/tests/test_formula_export_reliability.py +++ b/tests/test_formula_export_reliability.py @@ -198,6 +198,114 @@ def test_formula_export_worker_records_path_preparation_failure(workspace, setti assert job.status == "failed" +def _queued_job(ds): + return FormulaExport.objects.create( + dataset=ds, + user=ds.created_by, + context={"start": "2026-08", "end": "2026-08", "comparison": "yoy", "group": "department"}, + ) + + +def test_failed_export_rejects_symlinked_user_directory_without_touching_files(workspace, settings): + from finance.exporting import process_formula_export + + ds, *_ = workspace + job = _queued_job(ds) + root = Path(settings.MEDIA_ROOT) / "formula-exports" + root.mkdir(parents=True) + protected = Path(settings.MEDIA_ROOT) / "protected" + protected.mkdir() + victim = protected / f"{job.pk}.xlsx" + victim.write_bytes(b"protected workbook") + (root / str(job.user_id)).symlink_to(protected, target_is_directory=True) + + with pytest.raises(ValueError, match="符号链接"): + process_formula_export(job.pk) + + job.refresh_from_db() + assert job.status == "failed" + assert not job.path + assert victim.read_bytes() == b"protected workbook" + assert list(protected.glob("*.tmp")) == [] + + +def test_failed_export_rejects_symlinked_target_and_keeps_link(workspace, settings): + from finance.exporting import process_formula_export + + ds, *_ = workspace + job = _queued_job(ds) + user_dir = Path(settings.MEDIA_ROOT) / "formula-exports" / str(job.user_id) + user_dir.mkdir(parents=True) + protected_file = Path(settings.MEDIA_ROOT) / "protected.xlsx" + protected_file.write_bytes(b"protected workbook") + (user_dir / f"{job.pk}.xlsx").symlink_to(protected_file) + + with pytest.raises(ValueError, match="普通文件"): + process_formula_export(job.pk) + + job.refresh_from_db() + assert job.status == "failed" + assert protected_file.read_bytes() == b"protected workbook" + assert (user_dir / f"{job.pk}.xlsx").is_symlink() + assert list(user_dir.glob("*.tmp")) == [] + + +def test_failed_export_refuses_cleanup_when_parent_directory_swapped_to_symlink( + workspace, settings, monkeypatch +): + from finance.exporting import process_formula_export + + ds, *_ = workspace + job = _queued_job(ds) + root = Path(settings.MEDIA_ROOT) / "formula-exports" + user_dir = root / str(job.user_id) + user_dir.mkdir(parents=True) + protected = Path(settings.MEDIA_ROOT) / "protected" + protected.mkdir() + victim = protected / f"{job.pk}.xlsx" + victim.write_bytes(b"protected workbook") + + def swap_and_fail(*args, **kwargs): + args[-1].write(b"partial") + user_dir.rename(user_dir.with_name("swapped")) + user_dir.symlink_to(protected, target_is_directory=True) + raise ValueError("controlled failure after path validation") + + monkeypatch.setattr("finance.exporting.build_formula_workbook", swap_and_fail) + with pytest.raises(ValueError, match="controlled failure"): + process_formula_export(job.pk) + + job.refresh_from_db() + assert job.status == "failed" + assert victim.read_bytes() == b"protected workbook" + assert user_dir.is_symlink() + # 本次创建的临时文件残留在被换名的原目录中，由清理命令按保留期回收。 + assert list((root / "swapped").glob(f"{job.pk}.*.tmp")) + + +def test_formula_export_retry_after_failure_succeeds(workspace, monkeypatch): + from finance.exporting import process_formula_export + + ds, *_ = workspace + job = _queued_job(ds) + + def fail(*args, **kwargs): + args[-1].write(b"partial") + raise ValueError("controlled export failure") + + monkeypatch.setattr("finance.exporting.build_formula_workbook", fail) + with pytest.raises(ValueError, match="controlled export failure"): + process_formula_export(job.pk) + monkeypatch.undo() + + stats = process_formula_export(job.pk) + job.refresh_from_db() + assert stats["rows"] == 5 + assert job.status == "succeeded" + assert Path(job.path).read_bytes()[:2] == b"PK" + assert not list(Path(job.path).parent.glob(f"{job.pk}.*.tmp")) + + def test_formula_workbook_rejects_unclassified_label_collision(workspace): from decimal import Decimal