task_worker_v2_portable/platforms/kimi.py

import json
import struct


class KimiClient:
    def __init__(self, token):
        from curl_cffi import requests as cffi_req
        self.requests = cffi_req
        self.url = "https://www.kimi.com/apiv2/kimi.gateway.chat.v1.ChatService/Chat"
        self._refresh_token = ""

        # 兼容 JSON 格式: {"access_token":"...", "refresh_token":"..."}
        raw = token.strip()
        if raw.startswith("{"):
            try:
                import json
                obj = json.loads(raw)
                self._refresh_token = obj.get("refresh_token") or ""
                raw = obj.get("access_token") or ""
            except Exception:
                pass

        self.token = raw if raw.startswith("Bearer") else f"Bearer {raw}"

        import jwt as jwt_lib
        payload = jwt_lib.decode(raw.replace("Bearer ", ""), options={"verify_signature": False})
        self.device_id = payload.get("device_id", "")
        self.traffic_id = payload.get("sub", "")
        self.session_id = self._get_tobid()

    def refresh_access_token(self) -> bool:
        if not self._refresh_token:
            return False
        try:
            resp = self.requests.get(
                "https://www.kimi.com/api/auth/token/refresh",
                headers={
                    "authorization": f"Bearer {self._refresh_token}",
                    "origin": "https://www.kimi.com",
                    "referer": "https://www.kimi.com/",
                },
                timeout=15, impersonate="chrome146",
            )
            if resp.status_code == 200:
                data = resp.json()
                new_token = data.get("access_token") or ""
                new_refresh = data.get("refresh_token") or ""
                if new_token:
                    self.token = f"Bearer {new_token}" if not new_token.startswith("Bearer") else new_token
                    if new_refresh:
                        self._refresh_token = new_refresh
                    print(f"    [kimi] access_token 已刷新")
                    return True
        except Exception as e:
            print(f"    [kimi] refresh 异常: {e}")
        return False

    def _get_tobid(self):
        resp = self.requests.post(
            "https://gator.volces.com/tobid",
            json={"app_id": 20001731, "user_unique_id": self.traffic_id, "web_id": self.device_id},
            timeout=10,
        )
        return resp.json().get("tobid", "")

    def _build_headers(self):
        return {
            "authorization": self.token,
            "connect-protocol-version": "1",
            "content-type": "application/connect+json",
            "origin": "https://www.kimi.com",
            "referer": "https://www.kimi.com/",
            "x-msh-device-id": self.device_id,
            "x-msh-session-id": self.session_id,
            "x-msh-platform": "web",
            "x-msh-version": "2.0.0",
            "x-traffic-id": self.traffic_id,
            "x-language": "zh-CN",
            "r-timezone": "Asia/Shanghai",
            "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/149.0.0.0 Safari/537.36",
        }

    def chat(self, question, thinking=True):
        scenario = "SCENARIO_K3" if thinking else "SCENARIO_K2D5"
        body = {
            "scenario": scenario,
            "tools": [
                {"type": "TOOL_TYPE_SEARCH", "search": {}},
                {"type": "TOOL_TYPE_CRON_JOB"},
            ],
            "message": {
                "role": "user",
                "blocks": [{"message_id": "", "text": {"content": question}}],
                "scenario": scenario,
                "is_goal": False,
            },
            "options": {
                "thinking": bool(thinking),
                "enable_plugin": True,
                "reasoning_effort": "REASONING_EFFORT_NONE",
            },
            "project_id": "",
        }
        json_bytes = json.dumps(body, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
        packet = b"\x00" + struct.pack(">I", len(json_bytes)) + json_bytes

        resp = self.requests.post(
            self.url, headers=self._build_headers(), data=packet,
            timeout=360, impersonate="chrome146",
        )
        raw_req = json.dumps(body, ensure_ascii=False)
        return resp.content, f"\x00{raw_req}"


async def extract_kimi_from_browser():
    from playwright.async_api import async_playwright
    from .base import CDP_URL
    async with async_playwright() as pw:
        try:
            browser = await pw.chromium.connect_over_cdp(CDP_URL)
            for page in browser.contexts[0].pages:
                if "kimi.com" in page.url:
                    token = await page.evaluate("""() => {
                        const at = localStorage.getItem('access_token');
                        if (at && at.startsWith('eyJ') && at.length > 400) return at;
                        const ka = document.cookie.match(/kimi-auth=([^;]+)/);
                        if (ka && ka[1].startsWith('eyJ') && ka[1].length > 400) return ka[1];
                        const m = document.cookie.match(/access_token=([^;]+)/);
                        if (m && m[1].startsWith('eyJ') && m[1].length > 400) return m[1];
                        return '';
                    }""")
                    if token:
                        print(f"  [Kimi] Token 从浏览器提取: {token[:30]}...")
                        return {"token": token, "id": "-1"}
        except Exception:
            pass
    return None