ForcePilot/backend/package/yuxi/channel/extensions/wecom/directory.py

import logging

import httpx

logger = logging.getLogger(__name__)

DEPARTMENT_LIST_URL = "https://qyapi.weixin.qq.com/cgi-bin/department/list"
USER_LIST_URL = "https://qyapi.weixin.qq.com/cgi-bin/user/list"
USER_SIMPLE_LIST_URL = "https://qyapi.weixin.qq.com/cgi-bin/user/simplelist"
USER_GET_URL = "https://qyapi.weixin.qq.com/cgi-bin/user/get"
USER_LIST_ID_URL = "https://qyapi.weixin.qq.com/cgi-bin/user/list_id"


class WeComDirectory:
    def __init__(self, gateway):
        self._gateway = gateway
        self._http: httpx.AsyncClient | None = None

    @property
    def _token(self) -> str | None:
        return self._gateway.access_token if self._gateway else None

    async def list_departments(self, parent_id: int = 0) -> list[dict]:
        token = self._token
        if not token:
            return []
        url = f"{DEPARTMENT_LIST_URL}?access_token={token}&id={parent_id}"
        try:
            resp = await self._get_http().get(url)
            data = resp.json()
            if data.get("errcode") == 0:
                return data.get("department", [])
            logger.warning("WeCom list_departments failed: %s", data)
            return []
        except Exception:
            logger.exception("WeCom list_departments error")
            return []

    async def list_users_simple(self, department_id: int, fetch_child: int = 0) -> list[dict]:
        token = self._token
        if not token:
            return []
        url = f"{USER_SIMPLE_LIST_URL}?access_token={token}&department_id={department_id}&fetch_child={fetch_child}"
        try:
            resp = await self._get_http().get(url)
            data = resp.json()
            if data.get("errcode") == 0:
                return data.get("userlist", [])
            logger.warning("WeCom list_users_simple failed: %s", data)
            return []
        except Exception:
            logger.exception("WeCom list_users_simple error")
            return []

    async def list_users(self, department_id: int, fetch_child: int = 0) -> list[dict]:
        token = self._token
        if not token:
            return []
        url = f"{USER_LIST_URL}?access_token={token}&department_id={department_id}&fetch_child={fetch_child}"
        try:
            resp = await self._get_http().get(url)
            data = resp.json()
            if data.get("errcode") == 0:
                return data.get("userlist", [])
            logger.warning("WeCom list_users failed: %s", data)
            return []
        except Exception:
            logger.exception("WeCom list_users error")
            return []

    async def get_user(self, user_id: str) -> dict | None:
        token = self._token
        if not token:
            return None
        url = f"{USER_GET_URL}?access_token={token}&userid={user_id}"
        try:
            resp = await self._get_http().get(url)
            data = resp.json()
            if data.get("errcode") == 0:
                return data
            logger.warning("WeCom get_user failed: %s", data)
            return None
        except Exception:
            logger.exception("WeCom get_user error")
            return None

    async def list_user_ids(self, department_id: int = 0, cursor: str = "", limit: int = 10000) -> dict:
        token = self._token
        if not token:
            return {}
        body: dict = {"cursor": cursor, "limit": limit}
        if department_id > 0:
            body["department_id"] = department_id

        url = f"{USER_LIST_ID_URL}?access_token={token}"
        try:
            resp = await self._get_http().post(url, json=body)
            data = resp.json()
            if data.get("errcode") == 0:
                return data
            logger.warning("WeCom list_user_ids failed: %s", data)
            return {}
        except Exception:
            logger.exception("WeCom list_user_ids error")
            return {}

    def _get_http(self) -> httpx.AsyncClient:
        if self._http is None:
            self._http = httpx.AsyncClient(timeout=15.0)
        return self._http

    async def close(self):
        if self._http:
            await self._http.aclose()
            self._http = None