document-management/backend/tests/dms/test_b5_file_mutations.py

from __future__ import annotations

import hashlib
import io
import json
import struct
import zipfile
from pathlib import Path

import pytest
from docx import Document as WordDocument


def auth(token: str) -> dict[str, str]:
    return {"Authorization": f"Bearer {token}"}


def docx_bytes() -> bytes:
    stream = io.BytesIO()
    document = WordDocument()
    document.add_paragraph("B5测试文档")
    document.save(stream)
    return stream.getvalue()


def xlsx_bytes() -> bytes:
    stream = io.BytesIO()
    with zipfile.ZipFile(stream, "w") as archive:
        archive.writestr("[Content_Types].xml", "<Types/>")
        archive.writestr("xl/workbook.xml", "<workbook/>")
    return stream.getvalue()


def fake_zip() -> bytes:
    stream = io.BytesIO()
    with zipfile.ZipFile(stream, "w") as archive:
        archive.writestr("hello.txt", "not office")
    return stream.getvalue()


def cfb_bytes(stream_name: str) -> bytes:
    def entry(name: str, kind: int, child: int = 0xFFFFFFFF) -> bytes:
        value = bytearray(128)
        encoded = (name + "\0").encode("utf-16le")
        value[: len(encoded)] = encoded
        struct.pack_into("<H", value, 64, len(encoded))
        value[66] = kind
        value[67] = 1
        struct.pack_into(
            "<III", value, 68, 0xFFFFFFFF, 0xFFFFFFFF, child
        )
        struct.pack_into("<I", value, 116, 0xFFFFFFFE)
        return bytes(value)

    header = bytearray(512)
    header[:8] = bytes.fromhex("D0CF11E0A1B11AE1")
    struct.pack_into("<HH", header, 24, 0x003E, 3)
    struct.pack_into("<H", header, 28, 0xFFFE)
    struct.pack_into("<HH", header, 30, 9, 6)
    struct.pack_into("<I", header, 44, 1)
    struct.pack_into("<I", header, 48, 0)
    struct.pack_into("<I", header, 56, 4096)
    struct.pack_into("<I", header, 60, 0xFFFFFFFE)
    struct.pack_into("<I", header, 68, 0xFFFFFFFE)
    struct.pack_into("<I", header, 76, 1)
    for index in range(1, 109):
        struct.pack_into("<I", header, 76 + 4 * index, 0xFFFFFFFF)
    directory = bytearray(512)
    directory[:128] = entry("Root Entry", 5, 1)
    directory[128:256] = entry(stream_name, 2)
    fat = bytearray(b"\xff" * 512)
    struct.pack_into("<I", fat, 0, 0xFFFFFFFE)
    struct.pack_into("<I", fat, 4, 0xFFFFFFFD)
    return bytes(header + directory + fat)


def category_id(app) -> str:
    from dms.extensions import db
    from dms.models import Category

    with app.app_context():
        category = db.session.scalar(
            db.select(Category).where(Category.category_code == "ROOT_A")
        )
        return str(category.id)


def main_metadata(app, name: str = "B5_TEST_主案") -> dict:
    return {
        "documentName": name,
        "documentType": "MAIN",
        "summary": "B5上传测试",
        "categoryId": category_id(app),
        "securityLevel": "INTERNAL",
        "visibilityType": "ALL_AUTHENTICATED",
        "status": "PUBLISHED",
        "tags": ["B5", "上传"],
    }


def attachment_metadata(name: str = "B5_TEST_附件") -> dict:
    return {
        "documentName": name,
        "attachmentType": "POLICY",
        "summary": "B5附件",
        "tags": ["B5"],
    }


def upload(client, token, path, content, metadata):
    return client.post(
        path,
        data={
            "file": (io.BytesIO(content), metadata["documentName"] + ".docx"),
            "metadata": json.dumps(metadata, ensure_ascii=False),
        },
        headers=auth(token),
        content_type="multipart/form-data",
    )


def test_main_upload_persists_relative_uuid_file_and_hash(
    b5_app, b5_client, token_for
):
    from dms.extensions import db
    from dms.models import Document

    content = docx_bytes()
    response = upload(
        b5_client,
        token_for("admin"),
        "/api/v1/documents",
        content,
        main_metadata(b5_app),
    )
    assert response.status_code == 201
    item = response.get_json()["data"]
    assert item["id"].isdecimal()
    assert "fileRelativePath" not in item
    with b5_app.app_context():
        document = db.session.get(Document, int(item["id"]))
        relative = Path(document.file_relative_path)
        assert not relative.is_absolute()
        assert relative.parent.as_posix() == "original"
        assert document.original_file_name == "B5_TEST_主案.docx"
        assert relative.name != document.original_file_name
        path = Path(b5_app.config["DMS_STORAGE_ROOT"]) / relative
        assert path.read_bytes() == content
        assert document.file_hash == hashlib.sha256(content).hexdigest()


def test_sub_plan_upload_updates_parent_count(b5_app, b5_client, token_for):
    from dms.extensions import db
    from dms.models import Document

    token = token_for("admin")
    main = upload(
        b5_client,
        token,
        "/api/v1/documents",
        docx_bytes(),
        main_metadata(b5_app, "B5_TEST_父主案"),
    ).get_json()["data"]
    metadata = main_metadata(b5_app, "B5_TEST_子案")
    metadata["documentType"] = "SUB_PLAN"
    metadata["parentDocumentId"] = main["id"]
    child = upload(
        b5_client, token, "/api/v1/documents", docx_bytes(), metadata
    )
    assert child.status_code == 201
    assert child.get_json()["data"]["rootDocumentId"] == main["id"]
    with b5_app.app_context():
        assert db.session.get(Document, int(main["id"])).child_count == 1


def test_attachment_upload_forces_fixed_security(b5_client, token_for):
    response = upload(
        b5_client,
        token_for("admin"),
        "/api/v1/attachments",
        docx_bytes(),
        attachment_metadata(),
    )
    assert response.status_code == 201
    item = response.get_json()["data"]
    assert item["documentType"] == "ATTACHMENT"
    assert item["securityLevel"] == "PUBLIC"
    assert item["visibilityType"] == "ALL_AUTHENTICATED"
    assert item["categoryId"] is None


def test_non_admin_upload_forbidden(b5_app, b5_client, token_for):
    response = upload(
        b5_client,
        token_for("user"),
        "/api/v1/documents",
        docx_bytes(),
        main_metadata(b5_app),
    )
    assert response.status_code == 403


@pytest.mark.parametrize(
    "mutator",
    [
        lambda value: value.update(categoryId=1),
        lambda value: value.update(extra="bad"),
        lambda value: value.update(status="INVALID"),
        lambda value: value.update(categoryId="999999999"),
    ],
)
def test_upload_rejects_invalid_metadata(
    b5_app, b5_client, token_for, mutator
):
    metadata = main_metadata(b5_app)
    mutator(metadata)
    response = upload(
        b5_client,
        token_for("admin"),
        "/api/v1/documents",
        docx_bytes(),
        metadata,
    )
    assert response.status_code in {400, 404}


@pytest.mark.parametrize(
    ("filename", "content"),
    [
        ("bad.exe", b"MZ"),
        ("empty.pdf", b""),
        ("mismatch.pdf", docx_bytes()),
        ("broken.docx", b"PK\\x03\\x04broken"),
        ("fake.docx", fake_zip()),
        ("fake.xlsx", fake_zip()),
    ],
    ids=["unsupported", "empty", "mismatch", "broken-docx", "fake-docx", "fake-xlsx"],
)
def test_file_validation_rejects_spoofed_or_broken_files(
    b5_client, token_for, filename, content
):
    response = b5_client.post(
        "/api/v1/attachments",
        data={
            "file": (io.BytesIO(content), filename),
            "metadata": json.dumps(attachment_metadata()),
        },
        headers=auth(token_for("admin")),
        content_type="multipart/form-data",
    )
    assert response.status_code == 415
    assert response.get_json()["code"] == "UNSUPPORTED_FILE_TYPE"


@pytest.mark.parametrize(
    ("filename", "content"),
    [
        ("legal.pdf", b"%PDF-1.4\\nB5"),
        ("legal.docx", docx_bytes()),
        ("legal.xlsx", xlsx_bytes()),
        ("legal.doc", cfb_bytes("WordDocument")),
        ("legal.xls", cfb_bytes("Workbook")),
    ],
    ids=["pdf", "docx", "xlsx", "doc", "xls"],
)
def test_all_allowed_file_structures_upload(
    b5_client, token_for, filename, content
):
    metadata = attachment_metadata("B5_TEST_" + filename)
    response = b5_client.post(
        "/api/v1/attachments",
        data={
            "file": (io.BytesIO(content), filename),
            "metadata": json.dumps(metadata),
        },
        headers=auth(token_for("admin")),
        content_type="multipart/form-data",
    )
    assert response.status_code == 201


def test_same_hash_creates_distinct_documents(b5_client, token_for):
    content = docx_bytes()
    token = token_for("admin")
    first = upload(
        b5_client,
        token,
        "/api/v1/attachments",
        content,
        attachment_metadata("B5_TEST_同哈希1"),
    )
    second = upload(
        b5_client,
        token,
        "/api/v1/attachments",
        content,
        attachment_metadata("B5_TEST_同哈希2"),
    )
    assert first.status_code == second.status_code == 201
    assert first.get_json()["data"]["id"] != second.get_json()["data"]["id"]
    assert first.get_json()["data"]["fileHash"] == second.get_json()["data"]["fileHash"]


def test_database_failure_removes_new_file_and_record(
    b5_app, b5_client, token_for, monkeypatch
):
    from dms.extensions import db
    from dms.models import Document

    original_dir = Path(b5_app.config["DMS_STORAGE_ROOT"]) / "original"
    before = set(original_dir.iterdir())
    monkeypatch.setattr(
        "dms.services.document_mutation_service.business_audit",
        lambda **_: (_ for _ in ()).throw(RuntimeError("audit failed")),
    )
    response = upload(
        b5_client,
        token_for("admin"),
        "/api/v1/attachments",
        docx_bytes(),
        attachment_metadata("B5_TEST_补偿"),
    )
    assert response.status_code == 500
    assert set(original_dir.iterdir()) == before
    with b5_app.app_context():
        assert (
            db.session.scalar(
                db.select(Document).where(
                    Document.document_name == "B5_TEST_补偿"
                )
            )
            is None
        )


def test_batch_partial_success_and_manifest_mismatch(
    b5_client, token_for
):
    token = token_for("admin")
    metadata = [
        attachment_metadata("B5_TEST_批量成功"),
        attachment_metadata("B5_TEST_批量失败"),
    ]
    response = b5_client.post(
        "/api/v1/attachments/batch-import",
        data={
            "files": [
                (io.BytesIO(docx_bytes()), "good.docx"),
                (io.BytesIO(b"bad"), "bad.docx"),
            ],
            "items": json.dumps(metadata),
        },
        headers=auth(token),
        content_type="multipart/form-data",
    )
    assert response.status_code == 200
    result = response.get_json()["data"]
    assert (result["successCount"], result["failureCount"]) == (1, 1)
    assert [item["index"] for item in result["items"]] == [0, 1]
    mismatch = b5_client.post(
        "/api/v1/attachments/batch-import",
        data={
            "files": [(io.BytesIO(docx_bytes()), "one.docx")],
            "items": "[]",
        },
        headers=auth(token),
        content_type="multipart/form-data",
    )
    assert mismatch.status_code == 400
    assert mismatch.get_json()["code"] == "BATCH_MANIFEST_MISMATCH"


def test_edit_version_conflict_and_forbidden_identity_fields(
    b5_app, b5_client, token_for
):
    token = token_for("admin")
    created = upload(
        b5_client,
        token,
        "/api/v1/documents",
        docx_bytes(),
        main_metadata(b5_app, "B5_TEST_编辑前"),
    ).get_json()["data"]
    payload = {
        "documentName": "B5_TEST_编辑后",
        "summary": "更新摘要",
        "categoryId": category_id(b5_app),
        "securityLevel": "SECRET",
        "visibilityType": "CUSTOM",
        "status": "DRAFT",
        "tags": ["更新"],
        "rowVersion": created["rowVersion"],
    }
    response = b5_client.put(
        f"/api/v1/documents/{created['id']}", json=payload, headers=auth(token)
    )
    assert response.status_code == 200
    assert response.get_json()["data"]["documentName"] == "B5_TEST_编辑后"
    conflict = b5_client.put(
        f"/api/v1/documents/{created['id']}", json=payload, headers=auth(token)
    )
    assert conflict.status_code == 409
    payload["fileHash"] = "b" * 64
    forbidden = b5_client.put(
        f"/api/v1/documents/{created['id']}", json=payload, headers=auth(token)
    )
    assert forbidden.status_code == 400


def test_delete_sub_updates_parent_and_keeps_file(
    b5_app, b5_client, token_for
):
    from dms.extensions import db
    from dms.models import Document

    token = token_for("admin")
    main = upload(
        b5_client, token, "/api/v1/documents", docx_bytes(),
        main_metadata(b5_app, "B5_TEST_删除父案")
    ).get_json()["data"]
    metadata = main_metadata(b5_app, "B5_TEST_删除子案")
    metadata.update(documentType="SUB_PLAN", parentDocumentId=main["id"])
    child = upload(
        b5_client, token, "/api/v1/documents", docx_bytes(), metadata
    ).get_json()["data"]
    with b5_app.app_context():
        document = db.session.get(Document, int(child["id"]))
        path = Path(b5_app.config["DMS_STORAGE_ROOT"]) / document.file_relative_path
    response = b5_client.delete(
        f"/api/v1/documents/{child['id']}?rowVersion={child['rowVersion']}",
        headers=auth(token),
    )
    assert response.status_code == 200
    assert path.is_file()
    with b5_app.app_context():
        assert db.session.get(Document, int(main["id"])).child_count == 0


def test_main_with_child_and_mounted_attachment_deletion_rules(
    b5_app, b5_client, token_for
):
    token = token_for("admin")
    main = upload(
        b5_client, token, "/api/v1/documents", docx_bytes(),
        main_metadata(b5_app, "B5_TEST_有子主案")
    ).get_json()["data"]
    metadata = main_metadata(b5_app, "B5_TEST_保留子案")
    metadata.update(documentType="SUB_PLAN", parentDocumentId=main["id"])
    upload(b5_client, token, "/api/v1/documents", docx_bytes(), metadata)
    response = b5_client.delete(
        f"/api/v1/documents/{main['id']}?rowVersion={main['rowVersion'] + 1}",
        headers=auth(token),
    )
    assert response.status_code == 409
    assert response.get_json()["code"] == "MAIN_PLAN_HAS_CHILDREN"


def test_attachment_in_use_conflict(b4_app, b4_client, token_for):
    from dms.extensions import db
    from dms.models import Document

    attachment_id = b4_app.config["B4_IDS"]["attachment_one"]
    with b4_app.app_context():
        version = db.session.get(Document, attachment_id).row_version
    response = b4_client.delete(
        f"/api/v1/attachments/{attachment_id}?rowVersion={version}",
        headers=auth(token_for("admin")),
    )
    assert response.status_code == 409
    details = response.get_json()["details"]
    assert details["mountedPlanCount"] == 2
    assert len(details["mainPlans"]) == 2


def test_pdf_preview_and_download_count_audit(
    b5_app, b5_client, token_for
):
    from dms.extensions import db
    from dms.models import AuditLog, Document

    token = token_for("admin")
    metadata = attachment_metadata("B5_TEST_PDF")
    created_response = b5_client.post(
        "/api/v1/attachments",
        data={
            "file": (io.BytesIO(b"%PDF-1.4\\nB5-CONTENT"), "safe.pdf"),
            "metadata": json.dumps(metadata),
        },
        headers=auth(token),
        content_type="multipart/form-data",
    )
    created = created_response.get_json()["data"]
    preview = b5_client.get(
        f"/api/v1/documents/{created['id']}/preview", headers=auth(token)
    )
    assert preview.status_code == 200
    assert preview.mimetype == "application/pdf"
    assert preview.headers["Content-Disposition"].startswith("inline")
    download = b5_client.get(
        f"/api/v1/documents/{created['id']}/download", headers=auth(token)
    )
    assert download.status_code == 200
    assert download.data == b"%PDF-1.4\\nB5-CONTENT"
    assert download.headers["X-Request-Id"]
    assert "\r" not in download.headers["Content-Disposition"]
    with b5_app.app_context():
        document = db.session.get(Document, int(created["id"]))
        assert document.download_count == 1
        assert (
            db.session.scalar(
                db.select(db.func.count(AuditLog.id)).where(
                    AuditLog.target_id == document.id,
                    AuditLog.action_type == "DOWNLOAD_DOCUMENT",
                )
            )
            == 1
        )


def test_office_preview_file_missing_and_no_restore_route(
    b5_app, b5_client, token_for
):
    from dms.extensions import db
    from dms.models import Document

    token = token_for("admin")
    created = upload(
        b5_client,
        token,
        "/api/v1/attachments",
        docx_bytes(),
        attachment_metadata("B5_TEST_Office预览"),
    ).get_json()["data"]
    preview = b5_client.get(
        f"/api/v1/documents/{created['id']}/preview", headers=auth(token)
    )
    assert preview.status_code == 415
    assert preview.get_json()["code"] == "PREVIEW_UNAVAILABLE"
    with b5_app.app_context():
        document = db.session.get(Document, int(created["id"]))
        path = Path(b5_app.config["DMS_STORAGE_ROOT"]) / document.file_relative_path
        path.unlink()
    missing = b5_client.get(
        f"/api/v1/documents/{created['id']}/download", headers=auth(token)
    )
    assert missing.status_code == 404
    assert missing.get_json()["code"] == "FILE_NOT_FOUND"
    restore = b5_client.post(
        f"/api/v1/documents/{created['id']}/restore", headers=auth(token)
    )
    assert restore.status_code == 404


def test_size_limit_and_temporary_cleanup(b5_app, b5_client, token_for):
    b5_app.config["DMS_MAX_FILE_SIZE"] = 8
    temporary = Path(b5_app.config["DMS_STORAGE_ROOT"]) / "temporary"
    response = b5_client.post(
        "/api/v1/attachments",
        data={
            "file": (io.BytesIO(b"%PDF-1.4\nlarge"), "large.pdf"),
            "metadata": json.dumps(attachment_metadata("B5_TEST_超限")),
        },
        headers=auth(token_for("admin")),
        content_type="multipart/form-data",
    )
    assert response.status_code == 413
    assert not list(temporary.glob("*.upload"))


def test_file_promotion_failure_rolls_back_database(
    b5_app, b5_client, token_for, monkeypatch
):
    from dms.extensions import db
    from dms.models import Document

    monkeypatch.setattr(
        "dms.storage.uploads.StagedUpload.promote",
        lambda _self: (_ for _ in ()).throw(OSError("disk failure")),
    )
    response = upload(
        b5_client,
        token_for("admin"),
        "/api/v1/attachments",
        docx_bytes(),
        attachment_metadata("B5_TEST_文件失败回滚"),
    )
    assert response.status_code == 500
    with b5_app.app_context():
        assert db.session.scalar(
            db.select(Document).where(
                Document.document_name == "B5_TEST_文件失败回滚"
            )
        ) is None
    assert not list(
        (Path(b5_app.config["DMS_STORAGE_ROOT"]) / "temporary").glob("*.upload")
    )


def test_batch_all_success_and_all_failure(b5_client, token_for):
    token = token_for("admin")
    success = b5_client.post(
        "/api/v1/attachments/batch-import",
        data={
            "files": [
                (io.BytesIO(docx_bytes()), "a.docx"),
                (io.BytesIO(xlsx_bytes()), "b.xlsx"),
            ],
            "items": json.dumps(
                [
                    attachment_metadata("B5_TEST_全成功1"),
                    attachment_metadata("B5_TEST_全成功2"),
                ]
            ),
        },
        headers=auth(token),
        content_type="multipart/form-data",
    ).get_json()["data"]
    assert (success["successCount"], success["failureCount"]) == (2, 0)
    failed = b5_client.post(
        "/api/v1/attachments/batch-import",
        data={
            "files": [
                (io.BytesIO(b"bad"), "a.docx"),
                (io.BytesIO(b"bad"), "b.xlsx"),
            ],
            "items": json.dumps(
                [
                    attachment_metadata("B5_TEST_全失败1"),
                    attachment_metadata("B5_TEST_全失败2"),
                ]
            ),
        },
        headers=auth(token),
        content_type="multipart/form-data",
    ).get_json()["data"]
    assert (failed["successCount"], failed["failureCount"]) == (0, 2)


def test_attachment_edit_and_unmounted_delete_keep_file(
    b5_app, b5_client, token_for
):
    from dms.extensions import db
    from dms.models import Document

    token = token_for("admin")
    created = upload(
        b5_client,
        token,
        "/api/v1/attachments",
        docx_bytes(),
        attachment_metadata("B5_TEST_待编辑附件"),
    ).get_json()["data"]
    edited = b5_client.put(
        f"/api/v1/attachments/{created['id']}",
        json={
            "documentName": "B5_TEST_已编辑附件",
            "attachmentType": "WORK_STANDARD",
            "summary": "已编辑",
            "tags": ["编辑"],
            "rowVersion": created["rowVersion"],
        },
        headers=auth(token),
    ).get_json()["data"]
    assert edited["securityLevel"] == "PUBLIC"
    with b5_app.app_context():
        document = db.session.get(Document, int(created["id"]))
        path = Path(b5_app.config["DMS_STORAGE_ROOT"]) / document.file_relative_path
    response = b5_client.delete(
        f"/api/v1/attachments/{created['id']}?rowVersion={edited['rowVersion']}",
        headers=auth(token),
    )
    assert response.status_code == 200
    assert path.is_file()
    assert b5_client.get(
        f"/api/v1/attachments/{created['id']}", headers=auth(token)
    ).status_code == 404


def test_plan_download_requires_acl_and_path_escape_is_rejected(
    b5_app, b5_client, token_for
):
    from dms.extensions import db
    from dms.models import Document

    admin_token = token_for("admin")
    metadata = main_metadata(b5_app, "B5_TEST_无下载权限")
    metadata["visibilityType"] = "CUSTOM"
    created = upload(
        b5_client, admin_token, "/api/v1/documents", docx_bytes(), metadata
    ).get_json()["data"]
    denied = b5_client.get(
        f"/api/v1/documents/{created['id']}/download",
        headers=auth(token_for("user")),
    )
    assert denied.status_code == 403
    assert denied.get_json()["code"] == "DOCUMENT_DOWNLOAD_FORBIDDEN"
    attachment = upload(
        b5_client,
        admin_token,
        "/api/v1/attachments",
        docx_bytes(),
        attachment_metadata("B5_TEST_越界路径"),
    ).get_json()["data"]
    with b5_app.app_context():
        document = db.session.get(Document, int(attachment["id"]))
        document.file_relative_path = "../outside.docx"
        db.session.commit()
    escaped = b5_client.get(
        f"/api/v1/documents/{attachment['id']}/download",
        headers=auth(admin_token),
    )
    assert escaped.status_code == 500
    assert "outside.docx" not in escaped.get_data(as_text=True)