ForcePilot/backend/package/yuxi/channel/extensions/feishu/tools/drive.py

from __future__ import annotations

import logging

from yuxi.channel.extensions.feishu.client import get_client

logger = logging.getLogger(__name__)

FEISHU_DRIVE_TOOLS = [
    {
        "name": "feishu_drive_list",
        "description": "列出飞书云盘文件和文件夹",
        "parameters": {"folder_token": "文件夹 token（可选，默认根目录）"},
    },
    {
        "name": "feishu_drive_info",
        "description": "获取飞书云盘文件信息",
        "parameters": {"file_token": "文件 token"},
    },
    {
        "name": "feishu_drive_create_folder",
        "description": "创建飞书云盘文件夹",
        "parameters": {"name": "文件夹名称", "folder_token": "父文件夹 token"},
    },
    {
        "name": "feishu_drive_upload",
        "description": "上传文件到飞书云盘",
        "parameters": {"file_path": "本地文件路径", "parent_token": "目标文件夹 token（可选）"},
    },
    {
        "name": "feishu_drive_download",
        "description": "下载飞书云盘文件",
        "parameters": {"file_token": "文件 token", "save_path": "本地保存路径（可选）"},
    },
    {
        "name": "feishu_drive_delete",
        "description": "删除飞书云盘文件或文件夹",
        "parameters": {"file_token": "文件 token", "file_type": "文件类型（file/folder，默认 file）"},
    },
    {
        "name": "feishu_drive_move",
        "description": "移动飞书云盘文件到指定文件夹",
        "parameters": {"file_token": "文件 token", "target_folder_token": "目标文件夹 token"},
    },
]


async def list_drive_files(
    app_id: str,
    app_secret: str,
    domain: str,
    folder_token: str = "",
    http_timeout_ms: int = 30000,
) -> dict:
    try:
        client = get_client(app_id, app_secret, domain, http_timeout_ms)

        import lark_oapi
        req = lark_oapi.drive.v1.ListFileRequest(
            folder_token=folder_token if folder_token else None,
            page_size=50,
        )
        resp = await client.drive.v1.file.list_async(req)

        if resp.success() and resp.data:
            files = getattr(resp.data, "files", []) or []
            return {
                "success": True,
                "files": [
                    {
                        "token": getattr(f, "token", ""),
                        "name": getattr(f, "name", ""),
                        "type": getattr(f, "type", ""),
                        "parent_token": getattr(f, "parent_token", ""),
                    }
                    for f in files
                ],
                "has_more": getattr(resp.data, "has_more", False),
            }
        else:
            return {"success": False, "error": getattr(resp, "msg", "")}
    except ImportError:
        return {"success": False, "error": "lark-oapi SDK not installed"}
    except Exception:
        logger.exception("Feishu list_drive_files error")
        return {"success": False, "error": "Internal error"}


async def get_drive_file_info(
    app_id: str,
    app_secret: str,
    domain: str,
    file_token: str,
    http_timeout_ms: int = 30000,
) -> dict:
    try:
        client = get_client(app_id, app_secret, domain, http_timeout_ms)

        import lark_oapi
        req = lark_oapi.drive.v1.GetFileRequest(file_token=file_token)
        resp = await client.drive.v1.file.get_async(req)

        if resp.success() and resp.data:
            f = resp.data
            return {
                "success": True,
                "file": {
                    "token": getattr(f, "token", ""),
                    "name": getattr(f, "name", ""),
                    "type": getattr(f, "type", ""),
                    "size": getattr(f, "size", 0),
                    "url": getattr(f, "url", ""),
                },
            }
        else:
            return {"success": False, "error": getattr(resp, "msg", "")}
    except ImportError:
        return {"success": False, "error": "lark-oapi SDK not installed"}
    except Exception:
        logger.exception("Feishu get_drive_file_info error")
        return {"success": False, "error": "Internal error"}


async def create_drive_folder(
    app_id: str,
    app_secret: str,
    domain: str,
    name: str,
    folder_token: str,
    http_timeout_ms: int = 30000,
) -> dict:
    try:
        client = get_client(app_id, app_secret, domain, http_timeout_ms)

        import lark_oapi
        req = lark_oapi.drive.v1.CreateFolderFileRequest(
            request_body=lark_oapi.drive.v1.CreateFolderFileRequestBody(
                name=name,
                folder_token=folder_token,
            ),
        )
        resp = await client.drive.v1.file.create_folder_async(req)

        if resp.success() and resp.data:
            return {
                "success": True,
                "token": getattr(resp.data, "token", ""),
                "url": getattr(resp.data, "url", ""),
            }
        else:
            return {"success": False, "error": getattr(resp, "msg", "")}
    except ImportError:
        return {"success": False, "error": "lark-oapi SDK not installed"}
    except Exception:
        logger.exception("Feishu create_drive_folder error")
        return {"success": False, "error": "Internal error"}


async def upload_drive_file(
    app_id: str,
    app_secret: str,
    domain: str,
    file_path: str,
    parent_token: str = "",
    file_name: str = "",
    http_timeout_ms: int = 60000,
) -> dict:
    try:
        client = get_client(app_id, app_secret, domain, http_timeout_ms)

        import lark_oapi

        with open(file_path, "rb") as f:
            req = lark_oapi.drive.v1.UploadAllMediaRequest(
                request_body=lark_oapi.drive.v1.UploadAllMediaRequestBody(
                    file_name=file_name or file_path.rsplit("\\", 1)[-1].rsplit("/", 1)[-1],
                    parent_type="explorer",
                    parent_node=parent_token if parent_token else None,
                    size=str(len(f.read())),
                    file=f,
                ),
            )
            f.seek(0)
            resp = await client.drive.v1.media.upload_all_async(req)

        if resp.success() and resp.data:
            return {
                "success": True,
                "file_token": getattr(resp.data, "file_token", ""),
                "url": getattr(resp.data, "url", ""),
            }
        else:
            return {"success": False, "error": getattr(resp, "msg", "")}
    except ImportError:
        return {"success": False, "error": "lark-oapi SDK not installed"}
    except Exception:
        logger.exception("Feishu upload_drive_file error")
        return {"success": False, "error": "Internal error"}


async def download_drive_file(
    app_id: str,
    app_secret: str,
    domain: str,
    file_token: str,
    save_path: str = "",
    http_timeout_ms: int = 30000,
) -> dict:
    try:
        client = get_client(app_id, app_secret, domain, http_timeout_ms)

        import lark_oapi
        req = lark_oapi.drive.v1.DownloadMediaRequest(file_token=file_token)
        resp = await client.drive.v1.media.download_async(req)

        if resp.success():
            result = {"success": True, "file_token": file_token}
            if resp.file and save_path:
                with open(save_path, "wb") as f:
                    f.write(resp.file.read())
                result["saved_to"] = save_path
            return result
        else:
            return {"success": False, "error": getattr(resp, "msg", "")}
    except ImportError:
        return {"success": False, "error": "lark-oapi SDK not installed"}
    except Exception:
        logger.exception("Feishu download_drive_file error")
        return {"success": False, "error": "Internal error"}


async def delete_drive_file(
    app_id: str,
    app_secret: str,
    domain: str,
    file_token: str,
    file_type: str = "file",
    http_timeout_ms: int = 30000,
) -> dict:
    try:
        client = get_client(app_id, app_secret, domain, http_timeout_ms)

        import lark_oapi
        req = lark_oapi.drive.v1.DeleteFileRequest(
            file_token=file_token,
            type=file_type,
        )
        resp = await client.drive.v1.file.delete_async(req)

        if resp.success():
            return {"success": True}
        else:
            return {"success": False, "error": getattr(resp, "msg", "")}
    except ImportError:
        return {"success": False, "error": "lark-oapi SDK not installed"}
    except Exception:
        logger.exception("Feishu delete_drive_file error")
        return {"success": False, "error": "Internal error"}


async def move_drive_file(
    app_id: str,
    app_secret: str,
    domain: str,
    file_token: str,
    target_folder_token: str,
    http_timeout_ms: int = 30000,
) -> dict:
    try:
        client = get_client(app_id, app_secret, domain, http_timeout_ms)

        import lark_oapi
        req = lark_oapi.drive.v1.MoveFileRequest(
            file_token=file_token,
            request_body=lark_oapi.drive.v1.MoveFileRequestBody(
                target_parent_token=target_folder_token,
            ),
        )
        resp = await client.drive.v1.file.move_async(req)

        if resp.success():
            return {"success": True}
        else:
            return {"success": False, "error": getattr(resp, "msg", "")}
    except ImportError:
        return {"success": False, "error": "lark-oapi SDK not installed"}
    except Exception:
        logger.exception("Feishu move_drive_file error")
        return {"success": False, "error": "Internal error"}