data_query_skill/skills/feishu-data-publisher/scripts/feishu_client.py

"""Small environment-driven Feishu client for publishing tabular files."""

from __future__ import annotations

import json
import os
import time
from pathlib import Path
from typing import Any

import httpx


SUPPORTED_EXTENSIONS = {"csv", "xls", "xlsx"}


class FeishuApiError(RuntimeError):
    """Raised when Feishu returns a failed response without leaking credentials."""


class FeishuPublisher:
    def __init__(
        self,
        app_id: str,
        app_secret: str,
        *,
        base_url: str = "https://open.feishu.cn/open-apis",
        timeout: float = 30.0,
        client: httpx.Client | None = None,
    ) -> None:
        if not app_id.strip():
            raise ValueError("FEISHU_APP_ID 未配置")
        if not app_secret.strip():
            raise ValueError("FEISHU_APP_SECRET 未配置")
        self._app_id = app_id.strip()
        self._app_secret = app_secret.strip()
        self._base_url = base_url.rstrip("/")
        self._owns_client = client is None
        self._client = client or httpx.Client(timeout=timeout)

    @classmethod
    def from_env(cls, **kwargs: Any) -> "FeishuPublisher":
        return cls(
            os.getenv("FEISHU_APP_ID", ""),
            os.getenv("FEISHU_APP_SECRET", ""),
            base_url=os.getenv(
                "FEISHU_BASE_URL", "https://open.feishu.cn/open-apis"
            ),
            **kwargs,
        )

    def close(self) -> None:
        if self._owns_client:
            self._client.close()

    def __enter__(self) -> "FeishuPublisher":
        return self

    def __exit__(self, *_: object) -> None:
        self.close()

    @staticmethod
    def _payload(response: httpx.Response, action: str) -> dict[str, Any]:
        try:
            response.raise_for_status()
            payload = response.json()
        except (httpx.HTTPError, ValueError) as exc:
            raise FeishuApiError(f"{action}失败：HTTP 响应无效") from exc
        if payload.get("code") != 0:
            code = payload.get("code", "unknown")
            message = payload.get("msg", "unknown error")
            raise FeishuApiError(f"{action}失败：code={code}, msg={message}")
        return payload

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

    def tenant_access_token(self) -> str:
        response = self._client.post(
            f"{self._base_url}/auth/v3/tenant_access_token/internal",
            json={"app_id": self._app_id, "app_secret": self._app_secret},
        )
        return str(
            self._payload(response, "获取 tenant access token")[
                "tenant_access_token"
            ]
        )

    def list_chats(self, token: str, name: str = "") -> list[dict[str, str]]:
        matches: list[dict[str, str]] = []
        page_token = ""
        while True:
            params: dict[str, str | int] = {"page_size": 100}
            if page_token:
                params["page_token"] = page_token
            response = self._client.get(
                f"{self._base_url}/im/v1/chats",
                headers=self._headers(token),
                params=params,
            )
            data = self._payload(response, "查询机器人可见群聊").get("data", {})
            for item in data.get("items", []) or []:
                chat_name = str(item.get("name") or "")
                chat_id = str(item.get("chat_id") or "")
                if chat_id and (not name or name in chat_name):
                    matches.append({"name": chat_name, "chat_id": chat_id})
            if not data.get("has_more"):
                break
            page_token = str(data.get("page_token") or "")
            if not page_token:
                break
        return matches

    def resolve_chat_id(self, token: str, chat_name: str) -> str:
        exact = [
            chat
            for chat in self.list_chats(token, chat_name)
            if chat["name"] == chat_name
        ]
        if not exact:
            raise FeishuApiError(f"机器人可见群聊中未找到：{chat_name}")
        if len(exact) > 1:
            ids = ", ".join(chat["chat_id"] for chat in exact)
            raise FeishuApiError(
                f"群名匹配多个群聊：{chat_name}；请改用 chat_id。候选：{ids}"
            )
        return exact[0]["chat_id"]

    @staticmethod
    def validate_file(file_path: Path) -> str:
        if not file_path.is_file():
            raise FileNotFoundError(f"结果文件不存在：{file_path}")
        extension = file_path.suffix.lower().lstrip(".")
        if extension not in SUPPORTED_EXTENSIONS:
            allowed = ", ".join(sorted(SUPPORTED_EXTENSIONS))
            raise ValueError(f"不支持的文件类型：.{extension}；仅支持 {allowed}")
        if file_path.stat().st_size == 0:
            raise ValueError(f"结果文件为空：{file_path}")
        return extension

    def upload_import_source(
        self, token: str, file_path: Path, extension: str
    ) -> str:
        extra = json.dumps(
            {"obj_type": "sheet", "file_extension": extension},
            ensure_ascii=False,
        )
        with file_path.open("rb") as handle:
            response = self._client.post(
                f"{self._base_url}/drive/v1/medias/upload_all",
                headers=self._headers(token),
                data={
                    "file_name": file_path.name,
                    "parent_type": "ccm_import_open",
                    "size": str(file_path.stat().st_size),
                    "extra": extra,
                },
                files={
                    "file": (
                        file_path.name,
                        handle,
                        "application/octet-stream",
                    )
                },
            )
        return str(
            self._payload(response, "上传待导入文件")["data"]["file_token"]
        )

    def create_import_task(
        self,
        token: str,
        file_token: str,
        *,
        extension: str,
        title: str,
    ) -> str:
        response = self._client.post(
            f"{self._base_url}/drive/v1/import_tasks",
            headers={**self._headers(token), "Content-Type": "application/json"},
            json={
                "file_extension": extension,
                "file_token": file_token,
                "type": "sheet",
                "file_name": title,
                "point": {"mount_type": 1, "mount_key": ""},
            },
        )
        return str(self._payload(response, "创建表格导入任务")["data"]["ticket"])

    def wait_import_result(
        self,
        token: str,
        ticket: str,
        *,
        max_wait: float = 90.0,
        poll_interval: float = 2.0,
    ) -> dict[str, Any]:
        deadline = time.monotonic() + max_wait
        while time.monotonic() < deadline:
            response = self._client.get(
                f"{self._base_url}/drive/v1/import_tasks/{ticket}",
                headers=self._headers(token),
            )
            result = (
                self._payload(response, "查询表格导入结果")
                .get("data", {})
                .get("result", {})
            )
            status = result.get("job_status")
            if status == 0:
                return result
            if status == 3:
                message = result.get("job_error_msg", "unknown error")
                raise FeishuApiError(f"导入在线表格失败：{message}")
            time.sleep(poll_interval)
        raise TimeoutError(f"导入在线表格超过 {max_wait:g} 秒仍未完成")

    def set_anyone_editable(
        self, token: str, sheet_token: str, file_type: str = "sheet"
    ) -> None:
        response = self._client.patch(
            f"{self._base_url}/drive/v1/permissions/{sheet_token}/public",
            headers={**self._headers(token), "Content-Type": "application/json"},
            params={"type": file_type},
            json={
                "external_access_entity": "open",
                "link_share_entity": "anyone_editable",
            },
        )
        self._payload(response, "设置任何人可编辑权限")

    def send_sheet_card(
        self,
        token: str,
        *,
        chat_id: str,
        title: str,
        message: str,
        url: str,
    ) -> str:
        card = {
            "config": {"wide_screen_mode": True},
            "header": {
                "template": "blue",
                "title": {"tag": "plain_text", "content": title},
            },
            "elements": [
                {
                    "tag": "div",
                    "text": {"tag": "lark_md", "content": message},
                },
                {"tag": "hr"},
                {
                    "tag": "action",
                    "actions": [
                        {
                            "tag": "button",
                            "type": "primary",
                            "text": {
                                "tag": "plain_text",
                                "content": "打开在线表格",
                            },
                            "url": url,
                        }
                    ],
                },
            ],
        }
        response = self._client.post(
            f"{self._base_url}/im/v1/messages",
            headers={**self._headers(token), "Content-Type": "application/json"},
            params={"receive_id_type": "chat_id"},
            json={
                "receive_id": chat_id,
                "msg_type": "interactive",
                "content": json.dumps(card, ensure_ascii=False),
            },
        )
        return str(self._payload(response, "发送飞书群卡片")["data"]["message_id"])

    def publish(
        self,
        file_path: Path,
        *,
        title: str,
        message: str,
        chat_id: str = "",
        chat_name: str = "",
        notify: bool = True,
    ) -> dict[str, Any]:
        extension = self.validate_file(file_path)
        token = self.tenant_access_token()
        resolved_chat_id = chat_id.strip()
        if notify and not resolved_chat_id:
            if not chat_name.strip():
                raise ValueError(
                    "未配置目标群；请设置 FEISHU_TARGET_CHAT_ID、"
                    "FEISHU_TARGET_CHAT_NAME 或命令行参数"
                )
            resolved_chat_id = self.resolve_chat_id(token, chat_name.strip())

        file_token = self.upload_import_source(token, file_path, extension)
        ticket = self.create_import_task(
            token,
            file_token,
            extension=extension,
            title=title,
        )
        imported = self.wait_import_result(token, ticket)
        url = str(imported.get("url") or "")
        sheet_token = str(imported.get("token") or "")
        file_type = str(imported.get("type") or "sheet")
        if not url or not sheet_token:
            raise FeishuApiError("导入成功响应缺少在线表格 URL 或 token")

        self.set_anyone_editable(token, sheet_token, file_type)
        message_id = ""
        if notify:
            message_id = self.send_sheet_card(
                token,
                chat_id=resolved_chat_id,
                title=title,
                message=message,
                url=url,
            )
        return {
            "url": url,
            "permission": "anyone_editable",
            "message_sent": bool(message_id),
            "chat_id": resolved_chat_id if notify else "",
            "message_id": message_id,
        }