X-Financial/server/tests/test_knowledge_onlyoffice_tenant_security.py

from __future__ import annotations

import json
from datetime import UTC, datetime, timedelta
from io import BytesIO
from urllib.parse import parse_qs, urlsplit
from zipfile import ZIP_DEFLATED, ZipFile

import jwt
import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import StaticPool

from app.api.deps import CurrentUserContext
from app.db.base import Base
from app.models.knowledge_security import KnowledgeOnlyOfficeSession
from app.models.tenant import Tenant
from app.services import knowledge_onlyoffice_security as security_module
from app.services.knowledge import KnowledgeService
from app.services.knowledge_onlyoffice_callback import (
    handle_onlyoffice_callback,
    resolve_onlyoffice_content,
)
from app.services.knowledge_onlyoffice_security import (
    ONLYOFFICE_TOKEN_AUDIENCE,
    KnowledgeOnlyOfficeSessionService,
    OnlyOfficeReplayError,
    OnlyOfficeSecurityError,
    download_onlyoffice_document,
)
from app.services.knowledge_rag import KnowledgeRagService
from app.services.knowledge_tenant_scope import PLATFORM_KNOWLEDGE_SCOPE
from app.services.settings import OnlyOfficeRuntimeConfig

JWT_SECRET = "test-onlyoffice-security-secret-32bytes"


def _docx_bytes(text: str) -> bytes:
    stream = BytesIO()
    with ZipFile(stream, mode="w", compression=ZIP_DEFLATED) as archive:
        archive.writestr(
            "[Content_Types].xml",
            """<?xml version="1.0" encoding="UTF-8"?>
<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
  <Default Extension="xml" ContentType="application/xml" />
</Types>""",
        )
        archive.writestr(
            "word/document.xml",
            """<?xml version="1.0" encoding="UTF-8"?>
<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
  <w:body><w:p><w:r><w:t>"""
            + text
            + """</w:t></w:r></w:p></w:body>
</w:document>""",
        )
    return stream.getvalue()


def _user(tenant_id: str, *, admin: bool = True) -> CurrentUserContext:
    return CurrentUserContext(
        username=f"user-{tenant_id}",
        name=f"用户 {tenant_id}",
        role_codes=["manager"] if admin else ["employee"],
        is_admin=admin,
        tenant_id=tenant_id,
    )


def _factory():
    assert "tenants" in Base.metadata.tables
    engine = create_engine(
        "sqlite+pysqlite:///:memory:",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
    Tenant.__table__.create(bind=engine)
    KnowledgeOnlyOfficeSession.__table__.create(bind=engine)
    factory = sessionmaker(bind=engine, autoflush=False, autocommit=False)
    with factory() as db:
        db.add_all(
            [
                Tenant(tenant_id="tenant-a", tenant_code="A", name="A", status="active"),
                Tenant(tenant_id="tenant-b", tenant_code="B", name="B", status="active"),
            ]
        )
        db.commit()
    return factory


def _configure_onlyoffice(monkeypatch) -> OnlyOfficeRuntimeConfig:
    runtime = OnlyOfficeRuntimeConfig(
        enabled=True,
        public_url="https://docs.example.com",
        backend_url="https://app.example.com",
        jwt_secret=JWT_SECRET,
    )
    monkeypatch.setattr(
        "app.services.knowledge_onlyoffice.resolve_onlyoffice_settings",
        lambda *_args, **_kwargs: runtime,
    )
    monkeypatch.setattr(
        security_module,
        "resolve_onlyoffice_settings",
        lambda *_args, **_kwargs: runtime,
    )
    monkeypatch.setattr(
        KnowledgeRagService,
        "get_document_status_map",
        lambda _self, _document_ids: {},
    )
    monkeypatch.setattr(KnowledgeRagService, "delete_document", lambda *_args: None)
    return runtime


def _token_from_url(url: str, parameter: str) -> str:
    return parse_qs(urlsplit(url).query)[parameter][0]


def test_onlyoffice_tokens_bind_tenant_resource_key_version_and_audience(
    tmp_path,
    monkeypatch,
) -> None:
    _configure_onlyoffice(monkeypatch)
    factory = _factory()
    with factory() as db:
        service = KnowledgeService(storage_root=tmp_path, db=db, tenant_id="tenant-a")
        uploaded = service.upload_document(
            "制度政策",
            "制度.docx",
            _docx_bytes("version one"),
            _user("tenant-a"),
        )
        config = service.build_onlyoffice_config(uploaded.id, _user("tenant-a"))
        content_token = _token_from_url(config.config["document"]["url"], "access_token")
        callback_token = _token_from_url(
            config.config["editorConfig"]["callbackUrl"],
            "callback_token",
        )
        claims = jwt.decode(
            content_token,
            JWT_SECRET,
            algorithms=["HS256"],
            audience=ONLYOFFICE_TOKEN_AUDIENCE,
        )
        callback_claims = jwt.decode(
            callback_token,
            JWT_SECRET,
            algorithms=["HS256"],
            audience=ONLYOFFICE_TOKEN_AUDIENCE,
        )

        assert claims["tenant_id"] == "tenant-a"
        assert claims["document_id"] == uploaded.id
        assert claims["document_key"] == config.config["document"]["key"]
        assert claims["document_version"] == 1
        assert claims["editable"] is False
        assert callback_claims["jti"] == claims["jti"]
        assert callback_claims["exp"] - claims["exp"] >= 3 * 60 * 60
        row = db.get(KnowledgeOnlyOfficeSession, claims["jti"])
        assert row is not None and row.tenant_id == "tenant-a" and row.status == "active"
        content_path, _, _ = resolve_onlyoffice_content(
            db=db,
            storage_root=tmp_path,
            document_id=uploaded.id,
            access_token=content_token,
        )
        assert content_path.is_relative_to(
            tmp_path / "knowledge" / "tenants" / "tenant-a"
        )
        with ZipFile(content_path) as archive:
            assert b"version one" in archive.read("word/document.xml")

        tampered_claims = dict(claims)
        tampered_claims["tenant_id"] = "tenant-b"
        tampered_token = jwt.encode(tampered_claims, JWT_SECRET, algorithm="HS256")
        with pytest.raises(OnlyOfficeSecurityError, match="不匹配"):
            KnowledgeOnlyOfficeSessionService(db).validate_content(
                document_id=uploaded.id,
                token=tampered_token,
            )
        assert callback_token


def test_platform_document_session_is_tenant_bound_and_strictly_read_only(
    tmp_path,
    monkeypatch,
) -> None:
    _configure_onlyoffice(monkeypatch)
    document_id = "platform-doc"
    content = _docx_bytes("platform policy")
    filename = "平台制度.docx"
    stored_name = f"{document_id}__{filename}"
    platform_root = tmp_path / "knowledge" / "platform"
    folder_root = platform_root / "制度政策"
    folder_root.mkdir(parents=True)
    (folder_root / stored_name).write_bytes(content)
    (platform_root / ".index.json").write_text(
        json.dumps(
            {
                "version": 1,
                "documents": [
                    {
                        "id": document_id,
                        "folder": "制度政策",
                        "original_name": filename,
                        "stored_name": stored_name,
                        "mime_type": (
                            "application/vnd.openxmlformats-officedocument."
                            "wordprocessingml.document"
                        ),
                        "extension": "docx",
                        "size_bytes": len(content),
                        "sha256": "platform-checksum",
                        "created_at": "2026-07-17T00:00:00+00:00",
                        "updated_at": "2026-07-17T00:00:00+00:00",
                        "uploaded_by": "平台",
                        "version_number": 1,
                        "ingest_status": 1,
                    }
                ],
            },
            ensure_ascii=False,
        ),
        encoding="utf-8",
    )
    factory = _factory()
    with factory() as db:
        service = KnowledgeService(storage_root=tmp_path, db=db, tenant_id="tenant-a")
        config = service.build_onlyoffice_config(document_id, _user("tenant-a"))
        content_token = _token_from_url(config.config["document"]["url"], "access_token")
        claims = jwt.decode(
            content_token,
            JWT_SECRET,
            algorithms=["HS256"],
            audience=ONLYOFFICE_TOKEN_AUDIENCE,
        )
        assert claims["tenant_id"] == "tenant-a"
        assert claims["resource_scope"] == PLATFORM_KNOWLEDGE_SCOPE
        assert claims["editable"] is False
        resolved, _, _ = resolve_onlyoffice_content(
            db=db,
            storage_root=tmp_path,
            document_id=document_id,
            access_token=content_token,
        )
        assert resolved.read_bytes() == content
        assert resolved.is_relative_to(platform_root)
        with pytest.raises(ValueError, match="只读"):
            service.build_onlyoffice_config(
                document_id,
                _user("tenant-a"),
                editable=True,
            )


def test_view_session_never_writes_and_wrong_key_does_not_claim_session(
    tmp_path,
    monkeypatch,
) -> None:
    _configure_onlyoffice(monkeypatch)
    factory = _factory()
    with factory() as db:
        service = KnowledgeService(storage_root=tmp_path, db=db, tenant_id="tenant-a")
        uploaded = service.upload_document(
            "制度政策",
            "制度.docx",
            _docx_bytes("original"),
            _user("tenant-a"),
        )
        view_config = service.build_onlyoffice_config(uploaded.id, _user("tenant-a"))
        view_token = _token_from_url(
            view_config.config["editorConfig"]["callbackUrl"],
            "callback_token",
        )
        with pytest.raises(OnlyOfficeSecurityError, match="只读"):
            handle_onlyoffice_callback(
                db=db,
                storage_root=tmp_path,
                document_id=uploaded.id,
                callback_token=view_token,
                payload={
                    "status": 2,
                    "key": view_config.config["document"]["key"],
                    "url": "https://docs.example.com/download/view",
                },
            )

        edit_config = service.build_onlyoffice_config(
            uploaded.id,
            _user("tenant-a"),
            editable=True,
        )
        edit_token = _token_from_url(
            edit_config.config["editorConfig"]["callbackUrl"],
            "callback_token",
        )
        with pytest.raises(OnlyOfficeSecurityError, match="key"):
            handle_onlyoffice_callback(
                db=db,
                storage_root=tmp_path,
                document_id=uploaded.id,
                callback_token=edit_token,
                payload={
                    "status": 2,
                    "key": "wrong-key",
                    "url": "https://docs.example.com/download/edit",
                },
            )
        edit_claims = jwt.decode(
            edit_token,
            JWT_SECRET,
            algorithms=["HS256"],
            audience=ONLYOFFICE_TOKEN_AUDIENCE,
        )
        assert db.get(KnowledgeOnlyOfficeSession, edit_claims["jti"]).status == "active"
        assert service.get_document_entry(uploaded.id)["version_number"] == 1


def test_edit_callback_is_one_time_and_replay_is_rejected(tmp_path, monkeypatch) -> None:
    _configure_onlyoffice(monkeypatch)
    replacement = _docx_bytes("replacement")
    monkeypatch.setattr(
        "app.services.knowledge_onlyoffice_callback.download_onlyoffice_document",
        lambda _url, *, expected_filename: replacement,
    )
    factory = _factory()
    with factory() as db:
        service = KnowledgeService(storage_root=tmp_path, db=db, tenant_id="tenant-a")
        uploaded = service.upload_document(
            "制度政策",
            "制度.docx",
            _docx_bytes("original"),
            _user("tenant-a"),
        )
        config = service.build_onlyoffice_config(
            uploaded.id,
            _user("tenant-a"),
            editable=True,
        )
        callback_token = _token_from_url(
            config.config["editorConfig"]["callbackUrl"],
            "callback_token",
        )
        payload = {
            "status": 2,
            "key": config.config["document"]["key"],
            "url": "https://docs.example.com/download/final",
            "users": ["editor"],
        }
        handle_onlyoffice_callback(
            db=db,
            storage_root=tmp_path,
            document_id=uploaded.id,
            callback_token=callback_token,
            payload=payload,
        )

        claims = jwt.decode(
            callback_token,
            JWT_SECRET,
            algorithms=["HS256"],
            audience=ONLYOFFICE_TOKEN_AUDIENCE,
        )
        assert db.get(KnowledgeOnlyOfficeSession, claims["jti"]).status == "consumed"
        assert service.get_document_entry(uploaded.id)["version_number"] == 2
        assert service.get_document_content(uploaded.id)[0].read_bytes() == replacement
        with pytest.raises(OnlyOfficeReplayError):
            handle_onlyoffice_callback(
                db=db,
                storage_root=tmp_path,
                document_id=uploaded.id,
                callback_token=callback_token,
                payload=payload,
            )


def test_expired_or_stale_version_session_is_rejected(tmp_path, monkeypatch) -> None:
    _configure_onlyoffice(monkeypatch)
    factory = _factory()
    with factory() as db:
        service = KnowledgeService(storage_root=tmp_path, db=db, tenant_id="tenant-a")
        uploaded = service.upload_document(
            "制度政策",
            "制度.docx",
            _docx_bytes("original"),
            _user("tenant-a"),
        )
        config = service.build_onlyoffice_config(uploaded.id, _user("tenant-a"))
        content_token = _token_from_url(config.config["document"]["url"], "access_token")
        claims = jwt.decode(
            content_token,
            JWT_SECRET,
            algorithms=["HS256"],
            audience=ONLYOFFICE_TOKEN_AUDIENCE,
        )
        row = db.get(KnowledgeOnlyOfficeSession, claims["jti"])
        row.expires_at = datetime.now(UTC) - timedelta(seconds=1)
        db.commit()
        with pytest.raises(OnlyOfficeSecurityError):
            resolve_onlyoffice_content(
                db=db,
                storage_root=tmp_path,
                document_id=uploaded.id,
                access_token=content_token,
            )


def test_callback_ssrf_attempt_fails_session_without_overwriting_document(
    tmp_path,
    monkeypatch,
) -> None:
    _configure_onlyoffice(monkeypatch)
    factory = _factory()
    with factory() as db:
        service = KnowledgeService(storage_root=tmp_path, db=db, tenant_id="tenant-a")
        original = _docx_bytes("original")
        uploaded = service.upload_document(
            "制度政策",
            "制度.docx",
            original,
            _user("tenant-a"),
        )
        config = service.build_onlyoffice_config(
            uploaded.id,
            _user("tenant-a"),
            editable=True,
        )
        callback_token = _token_from_url(
            config.config["editorConfig"]["callbackUrl"],
            "callback_token",
        )
        with pytest.raises(OnlyOfficeSecurityError, match="不属于"):
            handle_onlyoffice_callback(
                db=db,
                storage_root=tmp_path,
                document_id=uploaded.id,
                callback_token=callback_token,
                payload={
                    "status": 2,
                    "key": config.config["document"]["key"],
                    "url": "https://evil.example/internal-metadata",
                },
            )
        claims = jwt.decode(
            callback_token,
            JWT_SECRET,
            algorithms=["HS256"],
            audience=ONLYOFFICE_TOKEN_AUDIENCE,
        )
        assert db.get(KnowledgeOnlyOfficeSession, claims["jti"]).status == "failed"
        assert service.get_document_content(uploaded.id)[0].read_bytes() == original
        assert service.get_document_entry(uploaded.id)["version_number"] == 1


def test_download_target_rejects_wrong_origin_private_dns_and_redirects(monkeypatch) -> None:
    _configure_onlyoffice(monkeypatch)
    with pytest.raises(OnlyOfficeSecurityError, match="不属于"):
        security_module._validate_download_target("https://evil.example/download")

    monkeypatch.setattr(
        security_module.socket,
        "getaddrinfo",
        lambda *_args, **_kwargs: [
            (2, 1, 6, "", ("127.0.0.1", 443)),
        ],
    )
    with pytest.raises(OnlyOfficeSecurityError, match="非公网"):
        security_module._validate_download_target("https://docs.example.com/download")

    monkeypatch.setattr(
        security_module.socket,
        "getaddrinfo",
        lambda *_args, **_kwargs: [
            (2, 1, 6, "", ("8.8.8.8", 443)),
        ],
    )

    class RedirectResponse:
        status = 302

        @staticmethod
        def getheader(_name):
            return None

    class FakeConnection:
        def request(self, *_args, **_kwargs) -> None:
            pass

        @staticmethod
        def getresponse():
            return RedirectResponse()

        def close(self) -> None:
            pass

    pinned: list[str] = []

    def fake_open(_parsed, resolved_ip):
        pinned.append(resolved_ip)
        return FakeConnection()

    monkeypatch.setattr(security_module, "_open_pinned_connection", fake_open)
    with pytest.raises(OnlyOfficeSecurityError, match="状态码 302"):
        download_onlyoffice_document(
            "https://docs.example.com/download",
            expected_filename="制度.docx",
        )
    assert pinned == ["8.8.8.8"]


def test_download_enforces_mime_size_and_ooxml_structure(monkeypatch) -> None:
    _configure_onlyoffice(monkeypatch)
    monkeypatch.setattr(
        security_module.socket,
        "getaddrinfo",
        lambda *_args, **_kwargs: [(2, 1, 6, "", ("8.8.8.8", 443))],
    )

    class Response:
        status = 200

        def __init__(self, body: bytes, content_type: str, declared_size: int | None = None):
            self.body = body
            self.content_type = content_type
            self.declared_size = declared_size

        def getheader(self, name):
            if name == "Content-Type":
                return self.content_type
            if name == "Content-Length":
                return None if self.declared_size is None else str(self.declared_size)
            return None

        def read(self, limit):
            return self.body[:limit]

    class Connection:
        def __init__(self, response):
            self.response = response

        def request(self, *_args, **_kwargs) -> None:
            pass

        def getresponse(self):
            return self.response

        def close(self) -> None:
            pass

    response = Response(
        _docx_bytes("safe"),
        "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    )
    monkeypatch.setattr(
        security_module,
        "_open_pinned_connection",
        lambda *_args: Connection(response),
    )
    downloaded = download_onlyoffice_document(
        "https://docs.example.com/download",
        expected_filename="制度.docx",
    )
    with ZipFile(BytesIO(downloaded)) as archive:
        assert b"safe" in archive.read("word/document.xml")

    response.content_type = "text/html"
    with pytest.raises(OnlyOfficeSecurityError, match="MIME"):
        download_onlyoffice_document(
            "https://docs.example.com/download",
            expected_filename="制度.docx",
        )
    response.content_type = "application/octet-stream"
    response.declared_size = 200 * 1024 * 1024
    with pytest.raises(OnlyOfficeSecurityError, match="大小"):
        download_onlyoffice_document(
            "https://docs.example.com/download",
            expected_filename="制度.docx",
        )
    response.declared_size = None
    response.body = b"not-a-zip"
    with pytest.raises(OnlyOfficeSecurityError, match="OOXML"):
        download_onlyoffice_document(
            "https://docs.example.com/download",
            expected_filename="制度.docx",
        )