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

"""方案密码锁：判定语义单元测试 + HTTP 集成测试。"""

from __future__ import annotations

import io
import json
from types import SimpleNamespace

import pytest

from dms.common.enums import DocumentPasswordAction, RoleCode
from dms.common.errors import (
    DocumentPasswordInvalidError,
    DocumentPasswordRequiredError,
    InvalidArgumentError,
)
from dms.security.passwords import hash_password
from dms.services.document_password_service import (
    _pop_password_spec,
    apply_password_lock,
    clear_password_lock,
    parse_protected_actions,
    password_required_for,
    protected_actions_of,
    require_document_password,
)

ALL_ACTIONS = [action.value for action in DocumentPasswordAction]


# ---------------------------------------------------------------------------
# 纯单元测试（无数据库依赖）
# ---------------------------------------------------------------------------


def stub_document(**overrides):
    base = dict(
        id=1,
        document_type="MAIN",
        password_hash=None,
        protect_view=False,
        protect_download=False,
        protect_edit=False,
        protect_delete=False,
        protect_config_permission=False,
        protect_bind_attachment=False,
        protect_unbind_attachment=False,
    )
    base.update(overrides)
    return SimpleNamespace(**base)


def admin_ctx():
    return SimpleNamespace(role_code=RoleCode.ADMIN.value)


def user_ctx():
    return SimpleNamespace(role_code=RoleCode.USER.value)


@pytest.mark.parametrize("value", [None, [], "VIEW", [1], ["VIEW", "HACK"]])
def test_parse_protected_actions_rejects_invalid(value):
    with pytest.raises(InvalidArgumentError):
        parse_protected_actions(value)


def test_parse_protected_actions_dedup_and_stable_order():
    assert parse_protected_actions(
        ["EDIT", "VIEW", "VIEW", "DOWNLOAD"]
    ) == ["VIEW", "DOWNLOAD", "EDIT"]


def test_protected_actions_of_empty_when_no_password():
    doc = stub_document(protect_view=True)
    assert protected_actions_of(doc) == []


def test_protected_actions_of_lists_checked_columns():
    doc = stub_document(
        password_hash="x",
        protect_view=True,
        protect_unbind_attachment=True,
    )
    assert protected_actions_of(doc) == ["VIEW", "UNBIND_ATTACHMENT"]


def test_password_required_for_user_only_checked_actions():
    doc = stub_document(password_hash="x", protect_view=True)
    ctx = user_ctx()
    assert password_required_for(doc, "VIEW", ctx) is True
    assert password_required_for(doc, "DOWNLOAD", ctx) is False


def test_password_required_for_admin_all_but_delete():
    doc = stub_document(password_hash="x")  # 未勾选任何操作
    ctx = admin_ctx()
    for action in ALL_ACTIONS:
        expected = action != "DELETE"
        assert password_required_for(doc, action, ctx) is expected, action


def test_password_required_for_unlocked_document_is_false():
    assert password_required_for(stub_document(), "VIEW", admin_ctx()) is False


def test_require_document_password_missing_raises_required():
    doc = stub_document(
        password_hash=hash_password("s3cret"), protect_view=True
    )
    with pytest.raises(DocumentPasswordRequiredError):
        require_document_password(doc, "VIEW", None, user_ctx())


def test_require_document_password_wrong_raises_invalid():
    doc = stub_document(
        password_hash=hash_password("s3cret"), protect_view=True
    )
    with pytest.raises(DocumentPasswordInvalidError):
        require_document_password(doc, "VIEW", "wrong", user_ctx())


def test_require_document_password_correct_passes():
    doc = stub_document(
        password_hash=hash_password("s3cret"), protect_view=True
    )
    require_document_password(doc, "VIEW", "s3cret", user_ctx())


def test_apply_and_clear_password_lock():
    doc = stub_document()
    apply_password_lock(doc, "s3cret", ["VIEW", "EDIT"])
    assert doc.password_hash is not None
    assert doc.protect_view is True
    assert doc.protect_edit is True
    assert doc.protect_download is False
    clear_password_lock(doc)
    assert doc.password_hash is None
    assert doc.protect_view is False
    assert doc.protect_edit is False


def test_pop_password_spec_absent_returns_none():
    assert _pop_password_spec(
        {"documentName": "x"}, attachment=False
    ) == (None, None)


@pytest.mark.parametrize(
    "payload",
    [
        {"password": "p", "documentType": "ATTACHMENT"},  # 附件不允许
        {"password": "p", "documentType": "SUB_PLAN"},  # 子案不允许
        {"password": "p"},  # 缺 protectedActions
        {"protectedActions": ["VIEW"]},  # 缺 password
        {"password": "", "protectedActions": ["VIEW"]},  # 空密码
        {"password": "x" * 129, "protectedActions": ["VIEW"]},  # 超长
        {"password": "p", "protectedActions": ["HACK"]},  # 未知操作码
    ],
)
def test_pop_password_spec_rejects_invalid(payload):
    with pytest.raises(InvalidArgumentError):
        _pop_password_spec(payload, attachment=False)


def test_pop_password_spec_rejects_overwrite_upload():
    payload = {
        "documentType": "MAIN",
        "password": "p",
        "protectedActions": ["VIEW"],
        "overwriteDocumentId": "123",
    }
    with pytest.raises(InvalidArgumentError):
        _pop_password_spec(payload, attachment=False)


def test_pop_password_spec_valid_pops_fields():
    payload = {
        "documentType": "MAIN",
        "password": "p",
        "protectedActions": ["EDIT", "VIEW"],
    }
    password, actions = _pop_password_spec(payload, attachment=False)
    assert password == "p"
    assert actions == ["VIEW", "EDIT"]
    assert "password" not in payload
    assert "protectedActions" not in payload


def test_pop_password_spec_attachment_rejected_even_for_main_type():
    payload = {
        "documentType": "MAIN",
        "password": "p",
        "protectedActions": ["VIEW"],
    }
    with pytest.raises(InvalidArgumentError):
        _pop_password_spec(payload, attachment=True)


# ---------------------------------------------------------------------------
# HTTP 集成测试（依赖独立测试库，本机无MySQL时自动跳过）
# ---------------------------------------------------------------------------

PASSWORD = "pw-test-1234"
MINIMAL_PDF = b"%PDF-1.4\n1 0 obj<</Type/Catalog>>endobj\ntrailer<</Root 1 0 R>>\n%%EOF"


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


def pwd_auth(token: str, password: str | None = None) -> dict[str, str]:
    headers = auth(token)
    if password is not None:
        headers["X-DMS-Document-Password"] = password
    return headers


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 == "CHILD_A1")
        )
        return str(category.id)


def upload_locked_main(b4_client, app, token_for, name="PW_TEST_主案", actions=("VIEW", "DOWNLOAD")):
    metadata = {
        "documentName": name,
        "documentType": "MAIN",
        "summary": "密码锁测试",
        "categoryId": category_id(app),
        "securityLevel": "INTERNAL",
        "tags": ["PW"],
        "password": PASSWORD,
        "protectedActions": list(actions),
    }
    response = b4_client.post(
        "/api/v1/documents",
        data={
            "file": (io.BytesIO(MINIMAL_PDF), name + ".pdf"),
            "metadata": json.dumps(metadata, ensure_ascii=False),
        },
        headers=auth(token_for("admin")),
        content_type="multipart/form-data",
    )
    assert response.status_code == 201, response.get_json()
    return response.get_json()["data"]


def test_upload_with_password_sets_lock(b4_client, b4_app, token_for):
    detail = upload_locked_main(b4_client, b4_app, token_for)
    assert detail["hasPassword"] is True
    assert detail["passwordProtectedActions"] == ["VIEW", "DOWNLOAD"]
    # ADMIN对设密主案查看恒需密码
    assert detail["viewLocked"] is True


def test_preview_blocked_without_or_wrong_password(
    b4_client, b4_app, token_for
):
    detail = upload_locked_main(b4_client, b4_app, token_for)
    url = f"/api/v1/documents/{detail['id']}/preview"
    blocked = b4_client.get(url, headers=auth(token_for("admin")))
    assert blocked.status_code == 403
    assert blocked.get_json()["code"] == "DOCUMENT_PASSWORD_REQUIRED"
    wrong = b4_client.get(
        url, headers=pwd_auth(token_for("admin"), "bad-password")
    )
    assert wrong.status_code == 403
    assert wrong.get_json()["code"] == "DOCUMENT_PASSWORD_INVALID"
    ok = b4_client.get(
        url, headers=pwd_auth(token_for("admin"), PASSWORD)
    )
    assert ok.status_code == 200


def test_download_blocked_until_password_given(b4_client, b4_app, token_for):
    detail = upload_locked_main(b4_client, b4_app, token_for)
    url = f"/api/v1/documents/{detail['id']}/download"
    blocked = b4_client.get(url, headers=auth(token_for("admin")))
    assert blocked.status_code == 403
    assert blocked.get_json()["code"] == "DOCUMENT_PASSWORD_REQUIRED"
    ok = b4_client.get(url, headers=pwd_auth(token_for("admin"), PASSWORD))
    assert ok.status_code == 200


def test_detail_shields_content_until_verified(b4_client, b4_app, token_for):
    detail = upload_locked_main(b4_client, b4_app, token_for)
    url = f"/api/v1/documents/{detail['id']}"
    shielded = b4_client.get(url, headers=auth(token_for("admin")))
    assert shielded.status_code == 200
    body = shielded.get_json()["data"]
    assert body["viewLocked"] is True
    assert body["viewPasswordVerified"] is False
    assert body["contentText"] is None
    verified = b4_client.get(
        url, headers=pwd_auth(token_for("admin"), PASSWORD)
    )
    assert verified.get_json()["data"]["viewPasswordVerified"] is True


def test_admin_delete_exempt_from_password(b4_client, b4_app, token_for):
    detail = upload_locked_main(
        b4_client, b4_app, token_for, actions=("VIEW", "DELETE")
    )
    response = b4_client.delete(
        f"/api/v1/documents/{detail['id']}?rowVersion={detail['rowVersion']}",
        headers=auth(token_for("admin")),
    )
    assert response.status_code == 200


def test_admin_edit_requires_password(b4_client, b4_app, token_for):
    detail = upload_locked_main(
        b4_client, b4_app, token_for, actions=("EDIT",)
    )
    payload = {
        "documentName": "PW_TEST_主案",
        "summary": "改摘要",
        "categoryId": category_id(b4_app),
        "securityLevel": "INTERNAL",
        "tags": ["PW"],
        "rowVersion": detail["rowVersion"],
    }
    blocked = b4_client.put(
        f"/api/v1/documents/{detail['id']}",
        json=payload,
        headers=auth(token_for("admin")),
    )
    assert blocked.status_code == 403
    assert blocked.get_json()["code"] == "DOCUMENT_PASSWORD_REQUIRED"
    ok = b4_client.put(
        f"/api/v1/documents/{detail['id']}",
        json=payload,
        headers=pwd_auth(token_for("admin"), PASSWORD),
    )
    assert ok.status_code == 200


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

    metadata = {
        "documentName": "PW_TEST_无密主案",
        "documentType": "MAIN",
        "summary": "密码管理",
        "categoryId": category_id(b4_app),
        "securityLevel": "INTERNAL",
        "tags": ["PW"],
    }
    created = b4_client.post(
        "/api/v1/documents",
        data={
            "file": (io.BytesIO(MINIMAL_PDF), "PW_TEST_无密主案.pdf"),
            "metadata": json.dumps(metadata, ensure_ascii=False),
        },
        headers=auth(token_for("admin")),
        content_type="multipart/form-data",
    )
    assert created.status_code == 201
    data = created.get_json()["data"]
    document_id, url = data["id"], f"/api/v1/documents/{data['id']}/password"

    # 未设密时禁止提供oldPassword
    with_old = b4_client.put(
        url,
        json={"rowVersion": data["rowVersion"], "oldPassword": "x"},
        headers=auth(token_for("admin")),
    )
    assert with_old.status_code == 400

    # 设置：无需oldPassword
    set_response = b4_client.put(
        url,
        json={
            "rowVersion": data["rowVersion"],
            "password": "first-pw",
            "protectedActions": ["VIEW"],
        },
        headers=auth(token_for("admin")),
    )
    assert set_response.status_code == 200
    summary = set_response.get_json()["data"]
    assert summary["hasPassword"] is True
    version = summary["rowVersion"]

    # 修改：oldPassword错误被拒
    wrong = b4_client.put(
        url,
        json={
            "rowVersion": version,
            "oldPassword": "nope",
            "password": "second-pw",
            "protectedActions": ["VIEW", "EDIT"],
        },
        headers=auth(token_for("admin")),
    )
    assert wrong.status_code == 403
    assert wrong.get_json()["code"] == "DOCUMENT_PASSWORD_INVALID"

    # rowVersion冲突被拒
    stale = b4_client.put(
        url,
        json={
            "rowVersion": data["rowVersion"],
            "oldPassword": "first-pw",
            "password": "second-pw",
            "protectedActions": ["VIEW"],
        },
        headers=auth(token_for("admin")),
    )
    assert stale.status_code == 409

    # 修改成功
    modified = b4_client.put(
        url,
        json={
            "rowVersion": version,
            "oldPassword": "first-pw",
            "password": "second-pw",
            "protectedActions": ["VIEW", "EDIT"],
        },
        headers=auth(token_for("admin")),
    )
    assert modified.status_code == 200
    assert modified.get_json()["data"]["protectedActions"] == ["VIEW", "EDIT"]
    version = modified.get_json()["data"]["rowVersion"]

    # 清除：仅oldPassword
    removed = b4_client.put(
        url,
        json={"rowVersion": version, "oldPassword": "second-pw"},
        headers=auth(token_for("admin")),
    )
    assert removed.status_code == 200
    assert removed.get_json()["data"]["hasPassword"] is False

    # 密码锁清除后校验列复位
    with b4_app.app_context():
        document = db.session.get(Document, int(document_id))
        assert document.password_hash is None
        assert document.protect_view is False


def test_attachment_upload_rejects_password_fields(
    b4_client, token_for
):
    metadata = {
        "documentName": "PW_TEST_附件",
        "attachmentType": "SUPPORT_ACHIEVEMENT",
        "summary": "附件禁用密码",
        "tags": ["PW"],
        "password": PASSWORD,
        "protectedActions": ["VIEW"],
    }
    response = b4_client.post(
        "/api/v1/attachments",
        data={
            "file": (io.BytesIO(MINIMAL_PDF), "PW_TEST_附件.pdf"),
            "metadata": json.dumps(metadata, ensure_ascii=False),
        },
        headers=auth(token_for("admin")),
        content_type="multipart/form-data",
    )
    assert response.status_code == 400


def test_overwrite_upload_rejects_password_fields(
    b4_client, b4_app, token_for
):
    detail = upload_locked_main(b4_client, b4_app, token_for)
    metadata = {
        "documentName": "PW_TEST_主案",
        "documentType": "MAIN",
        "summary": "覆盖",
        "categoryId": category_id(b4_app),
        "securityLevel": "INTERNAL",
        "tags": ["PW"],
        "overwriteDocumentId": detail["id"],
        "rowVersion": detail["rowVersion"],
        "password": "new-pw",
        "protectedActions": ["VIEW"],
    }
    response = b4_client.post(
        "/api/v1/documents",
        data={
            "file": (io.BytesIO(MINIMAL_PDF), "PW_TEST_主案.pdf"),
            "metadata": json.dumps(metadata, ensure_ascii=False),
        },
        headers=auth(token_for("admin")),
        content_type="multipart/form-data",
    )
    assert response.status_code == 400