gengx

wechat_rpa/archive_api.py

@@ -0,0 +1,498 @@
+# -*- coding: utf-8 -*-
+"""企业微信聊天归档 API。
+
+路由独立注册在 ``/api/v2/archive`` 下，不改动现有模型、用户和桌面端同步接口。
+"""
+
+from __future__ import annotations
+
+import hmac
+import ipaddress
+from pathlib import Path
+from typing import Any, Callable
+
+from fastapi import BackgroundTasks, Depends, FastAPI, HTTPException, Request
+from pydantic import BaseModel, Field
+from starlette.responses import FileResponse, RedirectResponse
+
+from archive_store import ArchiveStore
+
+
+class StorageBody(BaseModel):
+    bucket: str = ""
+    region: str = ""
+    custom_domain: str = ""
+    media_prefix: str = "archive/media"
+    export_prefix: str = "archive/exports"
+    encryption_mode: str = "AES256"
+    secret_id: str = ""
+    secret_key: str = ""
+    enabled: bool = False
+
+
+class MediaPrepareBody(BaseModel):
+    sha256: str
+    size_bytes: int
+    mime_type: str = "application/octet-stream"
+    original_filename: str = "file"
+
+
+class MultipartPartBody(BaseModel):
+    part_number: int = Field(ge=1, le=10000)
+    etag: str = Field(min_length=1, max_length=512)
+
+
+class MediaMultipartCompleteBody(BaseModel):
+    upload_id: str = Field(min_length=1, max_length=2048)
+    parts: list[MultipartPartBody] = Field(min_length=1, max_length=10000)
+
+
+class ImportBody(BaseModel):
+    source_account: dict[str, Any]
+    messages: list[dict[str, Any]] = Field(min_length=1, max_length=5000)
+    batch_id: str = ""
+    source_table: str = "message_table"
+    checkpoint: Any = None
+
+
+class MetadataBody(BaseModel):
+    source_account: dict[str, Any]
+    people: list[dict[str, Any]] = Field(default_factory=list, max_length=5000)
+    conversations: list[dict[str, Any]] = Field(default_factory=list, max_length=5000)
+
+
+class DesktopCheckpointBody(BaseModel):
+    source_account: dict[str, Any]
+    source_table: str = "message_table"
+    checkpoint: dict[str, Any]
+
+
+class MediaAccessBody(BaseModel):
+    media_ids: list[str] = Field(min_length=1, max_length=200)
+    expires: int = Field(default=300, ge=60, le=900)
+
+
+class ExportBody(BaseModel):
+    formats: list[str] = Field(default_factory=lambda: ["sql", "xlsx"])
+    filters: dict[str, Any] = Field(default_factory=dict)
+
+
+class IdentityBody(BaseModel):
+    identity_type: str
+    scope_id: str = ""
+    external_id: str
+    verified: bool = True
+
+
+def _http_error(exc: Exception) -> HTTPException:
+    if isinstance(exc, KeyError):
+        return HTTPException(status_code=404, detail=str(exc).strip("'"))
+    if isinstance(exc, ValueError):
+        return HTTPException(status_code=400, detail=str(exc))
+    if isinstance(exc, RuntimeError):
+        return HTTPException(status_code=409, detail=str(exc))
+    return HTTPException(status_code=502, detail=f"外部存储操作失败：{exc}")
+
+
+def register_archive_routes(
+    app: FastAPI,
+    database: Any,
+    current: Callable[..., Any],
+    require: Callable[..., Any],
+    client_ip: Callable[[Request], str],
+    desktop_sync_key: str = "",
+) -> ArchiveStore:
+    """将聊天归档路由附加到现有 FastAPI 应用。"""
+
+    store = ArchiveStore(database)
+    store.initialize()
+    app.state.archive_store = store
+
+    async def desktop_archive_ingest(request: Request) -> None:
+        """只允许本机桌面程序用机器凭证写归档。
+
+        管理员 Bearer Token 继续走原有接口。这个入口只用于打开桌面软件后的
+        自动备份，既不把管理员密码写进客户端，也不向局域网暴露批量写库能力。
+        """
+
+        # 这里必须看真实 TCP 对端，不能使用可能由调用方伪造的 X-Forwarded-For。
+        remote = request.client.host if request.client else ""
+        try:
+            is_loopback = ipaddress.ip_address(remote).is_loopback
+        except ValueError:
+            is_loopback = False
+        supplied = request.headers.get("x-desktop-sync-key", "")
+        if (
+            not is_loopback
+            or not supplied
+            or not desktop_sync_key
+            or not hmac.compare_digest(supplied, desktop_sync_key)
+        ):
+            raise HTTPException(status_code=401, detail="本机归档凭证无效")
+
+    @app.get("/api/v2/archive/stats")
+    async def archive_stats(_: Any = Depends(require("im:read"))) -> dict[str, Any]:
+        return store.stats()
+
+    @app.get("/api/v2/archive/conversations")
+    async def archive_conversations(
+        limit: int = 50,
+        cursor: str = "",
+        _: Any = Depends(require("im:content:read")),
+    ) -> dict[str, Any]:
+        try:
+            return store.conversations(limit, cursor)
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.get("/api/v2/archive/conversations/{conversation_id}/messages")
+    async def archive_messages(
+        conversation_id: str,
+        limit: int = 100,
+        cursor: str = "",
+        _: Any = Depends(require("im:content:read")),
+    ) -> dict[str, Any]:
+        try:
+            return store.messages(conversation_id, limit, cursor)
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.post("/api/v2/archive/imports/messages")
+    async def archive_import_messages(
+        body: ImportBody,
+        request: Request,
+        principal: Any = Depends(require("im:import")),
+    ) -> dict[str, Any]:
+        try:
+            return store.import_messages(
+                body.model_dump(), principal.id, client_ip(request)
+            )
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.post("/api/v2/archive/imports/metadata")
+    async def archive_import_metadata(
+        body: MetadataBody,
+        request: Request,
+        principal: Any = Depends(require("im:import")),
+    ) -> dict[str, Any]:
+        try:
+            return store.sync_metadata(
+                body.model_dump(), principal.id, client_ip(request)
+            )
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.get("/api/v2/archive/desktop/checkpoint")
+    async def desktop_archive_checkpoint(
+        external_account_id: str,
+        source_table: str = "message_table",
+        _: None = Depends(desktop_archive_ingest),
+    ) -> dict[str, Any]:
+        if not str(external_account_id or "").strip():
+            raise HTTPException(status_code=400, detail="账号标识不能为空")
+        return {
+            "checkpoint": store.source_checkpoint(external_account_id, source_table)
+        }
+
+    @app.post("/api/v2/archive/desktop/checkpoint")
+    async def desktop_archive_advance_checkpoint(
+        body: DesktopCheckpointBody,
+        _: None = Depends(desktop_archive_ingest),
+    ) -> dict[str, Any]:
+        return {
+            "checkpoint": store.advance_source_checkpoint(
+                body.source_account, body.source_table, body.checkpoint
+            )
+        }
+
+    @app.get("/api/v2/archive/desktop/pending-attachments")
+    async def desktop_pending_attachments(
+        external_account_id: str,
+        limit: int = 500,
+        _: None = Depends(desktop_archive_ingest),
+    ) -> dict[str, Any]:
+        if not str(external_account_id or "").strip():
+            raise HTTPException(status_code=400, detail="账号标识不能为空")
+        return {
+            "source_message_ids": store.claim_pending_attachment_source_ids(
+                external_account_id, limit
+            )
+        }
+
+    @app.post("/api/v2/archive/desktop/imports/messages")
+    async def desktop_archive_import_messages(
+        body: ImportBody,
+        request: Request,
+        _: None = Depends(desktop_archive_ingest),
+    ) -> dict[str, Any]:
+        try:
+            return store.import_messages(body.model_dump(), None, client_ip(request))
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.post("/api/v2/archive/desktop/imports/metadata")
+    async def desktop_archive_import_metadata(
+        body: MetadataBody,
+        request: Request,
+        _: None = Depends(desktop_archive_ingest),
+    ) -> dict[str, Any]:
+        try:
+            return store.sync_metadata(body.model_dump(), None, client_ip(request))
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.get("/api/v2/archive/people")
+    async def archive_people(
+        limit: int = 100,
+        keyword: str = "",
+        _: Any = Depends(require("im:identity:write")),
+    ) -> dict[str, Any]:
+        return {"items": store.people(limit, keyword)}
+
+    @app.get("/api/v2/archive/people/{person_id}")
+    async def archive_person(
+        person_id: str,
+        _: Any = Depends(require("im:identity:write")),
+    ) -> dict[str, Any]:
+        try:
+            return {"person": store.person_detail(person_id)}
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.post("/api/v2/archive/people/{person_id}/identities")
+    async def bind_archive_identity(
+        person_id: str,
+        body: IdentityBody,
+        request: Request,
+        principal: Any = Depends(require("im:identity:write")),
+    ) -> dict[str, Any]:
+        try:
+            return {
+                "person": store.bind_identity(
+                    person_id, body.model_dump(), principal.id, client_ip(request)
+                )
+            }
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.get("/api/v2/archive/storage")
+    async def archive_storage(
+        _: Any = Depends(require("im:storage:write")),
+    ) -> dict[str, Any]:
+        return {"storage": store.storage_config()}
+
+    @app.put("/api/v2/archive/storage")
+    async def save_archive_storage(
+        body: StorageBody,
+        request: Request,
+        principal: Any = Depends(require("im:storage:write")),
+    ) -> dict[str, Any]:
+        try:
+            return {
+                "storage": store.save_storage_config(
+                    body.model_dump(), principal.id, client_ip(request)
+                )
+            }
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.post("/api/v2/archive/storage/test")
+    async def test_archive_storage(
+        request: Request,
+        principal: Any = Depends(require("im:storage:write")),
+    ) -> dict[str, Any]:
+        try:
+            result = store.test_storage()
+            database.audit(
+                principal.id,
+                "archive.storage.test",
+                f"ok={int(bool(result.get('ok')))} bucket={result.get('bucket', '')}",
+                client_ip(request),
+            )
+            return {"result": result}
+        except Exception as exc:
+            database.audit(
+                principal.id,
+                "archive.storage.test",
+                f"ok=0 error={str(exc)[:500]}",
+                client_ip(request),
+            )
+            raise _http_error(exc) from exc
+
+    @app.get("/api/v2/archive/media")
+    async def archive_media(
+        limit: int = 100,
+        _: Any = Depends(require("im:content:read")),
+    ) -> dict[str, Any]:
+        return {"items": store.media_items(limit)}
+
+    @app.post("/api/v2/archive/media/access-urls")
+    async def archive_media_access_urls(
+        body: MediaAccessBody,
+        request: Request,
+        principal: Any = Depends(require("im:content:read")),
+    ) -> dict[str, Any]:
+        try:
+            items = store.media_download_urls(body.media_ids, body.expires)
+            database.audit(
+                principal.id,
+                "archive.media.access",
+                f"count={len(items)}",
+                client_ip(request),
+            )
+            return {"items": items}
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.post("/api/v2/archive/media/prepare")
+    async def prepare_archive_media(
+        body: MediaPrepareBody,
+        _: Any = Depends(require("im:import")),
+    ) -> dict[str, Any]:
+        try:
+            return store.prepare_media(body.model_dump())
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.post("/api/v2/archive/desktop/media/prepare")
+    async def desktop_prepare_archive_media(
+        body: MediaPrepareBody,
+        _: None = Depends(desktop_archive_ingest),
+    ) -> dict[str, Any]:
+        try:
+            return store.prepare_media(body.model_dump())
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.post("/api/v2/archive/media/{media_id}/complete")
+    async def complete_archive_media(
+        media_id: str,
+        _: Any = Depends(require("im:import")),
+    ) -> dict[str, Any]:
+        try:
+            return {"media": store.complete_media(media_id)}
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.post("/api/v2/archive/media/{media_id}/multipart-complete")
+    async def complete_archive_multipart_media(
+        media_id: str,
+        body: MediaMultipartCompleteBody,
+        _: Any = Depends(require("im:import")),
+    ) -> dict[str, Any]:
+        try:
+            return {
+                "media": store.complete_multipart_media(
+                    media_id,
+                    body.upload_id,
+                    [part.model_dump() for part in body.parts],
+                )
+            }
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.post("/api/v2/archive/desktop/media/{media_id}/complete")
+    async def desktop_complete_archive_media(
+        media_id: str,
+        _: None = Depends(desktop_archive_ingest),
+    ) -> dict[str, Any]:
+        try:
+            return {"media": store.complete_media(media_id)}
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.post("/api/v2/archive/desktop/media/{media_id}/multipart-complete")
+    async def desktop_complete_archive_multipart_media(
+        media_id: str,
+        body: MediaMultipartCompleteBody,
+        _: None = Depends(desktop_archive_ingest),
+    ) -> dict[str, Any]:
+        try:
+            return {
+                "media": store.complete_multipart_media(
+                    media_id,
+                    body.upload_id,
+                    [part.model_dump() for part in body.parts],
+                )
+            }
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.get("/api/v2/archive/media/{media_id}/view")
+    async def view_archive_media(
+        media_id: str,
+        request: Request,
+        principal: Any = Depends(require("im:content:read")),
+    ) -> RedirectResponse:
+        try:
+            url = store.media_download_url(media_id)
+            database.audit(
+                principal.id, "archive.media.view", f"media={media_id}", client_ip(request)
+            )
+            return RedirectResponse(url=url, status_code=307)
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.post("/api/v2/archive/exports")
+    async def create_archive_export(
+        body: ExportBody,
+        request: Request,
+        background: BackgroundTasks,
+        principal: Any = Depends(require("im:export")),
+    ) -> dict[str, Any]:
+        try:
+            job = store.create_export_job(
+                body.formats, body.filters, principal.id, client_ip(request)
+            )
+            background.add_task(store.run_export_job, job["id"])
+            return {"job": job}
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.get("/api/v2/archive/exports")
+    async def archive_exports(
+        limit: int = 100,
+        _: Any = Depends(require("im:export")),
+    ) -> dict[str, Any]:
+        return {"jobs": store.export_jobs(limit)}
+
+    @app.get("/api/v2/archive/exports/{job_id}")
+    async def archive_export(
+        job_id: str,
+        _: Any = Depends(require("im:export")),
+    ) -> dict[str, Any]:
+        try:
+            return {"job": store.export_job(job_id)}
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    @app.get("/api/v2/archive/export-files/{file_id}/download")
+    async def download_archive_export(
+        file_id: str,
+        request: Request,
+        principal: Any = Depends(require("im:export")),
+    ):
+        try:
+            item = store.export_file(file_id)
+            database.audit(
+                principal.id,
+                "archive.export.download",
+                f"file={file_id} job={item['job_id']}",
+                client_ip(request),
+            )
+            if item["storage_status"] == "cos":
+                return RedirectResponse(
+                    url=store.export_download_url(file_id), status_code=307
+                )
+            path = Path(item["local_path"]).resolve()
+            root = store.export_root.resolve()
+            if not path.is_relative_to(root) or not path.is_file():
+                raise KeyError("导出文件已不在本机")
+            return FileResponse(
+                path, filename=item["file_name"], media_type="application/octet-stream"
+            )
+        except Exception as exc:
+            raise _http_error(exc) from exc
+
+    return store