task_worker_v2_portable/platforms/qianwen.py

import base64
import hashlib
import hmac as hmac_mod
import json
import time
import uuid
from urllib.parse import urlparse, quote


class QianwenClient:
    def __init__(self, cookies=None):
        self.url = "https://chat2.qianwen.com/api/v2/chat"
        self.share_url = "https://chat2-api.qianwen.com/api/v1/share/create"
        self.ua = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/131.0.0.0 Safari/537.36"
        self._cookie_str = ""
        self._xsrf_token = ""
        if cookies:
            self.set_cookies(cookies)

    def set_cookies(self, cookies):
        self._cookie_str = "; ".join(f'{c["name"]}={c["value"]}' for c in cookies)
        self._xsrf_token = next(
            (c["value"] for c in cookies if c["name"] == "XSRF-TOKEN" and "chat2" in c.get("domain", "")),
            next((c["value"] for c in cookies if c["name"] == "XSRF-TOKEN"), ""),
        )

    @staticmethod
    def _get_sign(method, url, params, data):
        path = urlparse(url).path
        content_md5 = hashlib.md5(data.encode("utf-8")).hexdigest()
        sign_params = {k: v for k, v in params.items() if k != "sign"}
        sorted_q = "&".join(f"{k}={quote(str(sign_params[k]), safe='')}" for k in sorted(sign_params))
        base_string = f"{method.upper()}{path}{sorted_q}{content_md5}quark"
        return "3a34" + hmac_mod.new(b"af54041a93cd4f6a757f", base_string.encode(), hashlib.sha1).hexdigest()

    @staticmethod
    def _generate_ut():
        from deps.utdid import generate_utdid
        from Crypto.Cipher import AES
        from Crypto.Util.Padding import pad
        key = iv = b"a89a69c4bbe039e8"
        cipher = AES.new(key, AES.MODE_CBC, iv)
        return base64.b64encode(b"\x3a\x34" + cipher.encrypt(pad(generate_utdid().encode("utf-8"), AES.block_size))).decode()

    def _build_params(self):
        ts = int(time.time() * 1000)
        return {
            "ut": self._generate_ut(), "ve": "2.12.0", "bi": "37261", "pf": "3300",
            "pr": "qwen", "fr": "pc", "sv": "release", "cg": "default",
            "chat_client": "h5", "biz_id": "ai_qwen",
            "vcode": str(ts), "nonce": str(ts + 2), "timestamp": str(ts + 2), "sign_type": "2",
        }

    def _build_headers(self):
        headers = {
            "accept": "application/json, text/event-stream, text/plain, */*",
            "content-type": "application/json",
            "x-platform": "pc_tongyi",
            "user-agent": self.ua,
        }
        if self._cookie_str:
            headers["cookie"] = self._cookie_str
        if self._xsrf_token:
            headers["x-xsrf-token"] = self._xsrf_token
        return headers

    async def chat(self, question, thinking=False):
        import aiohttp
        from yarl import URL

        params = self._build_params()
        params["protocol_version"] = "v2"
        payload = {
            "req_id": uuid.uuid4().hex, "parent_req_id": "0",
            "session_id": uuid.uuid4().hex, "topic_id": uuid.uuid4().hex,
            "scene": "chat", "sub_scene": "", "scene_param": "first_turn",
            "biz_id": "ai_qwen", "model": "Qwen", "from": "default",
            "protocol_version": "v2", "messages_merge": False, "chat_client": "h5",
            "deep_search": "1" if thinking else "0", "ai_tool_scene": "",
            "temporary": False,
            "messages": [{"mime_type": "text/plain", "content": question,
                          "meta_data": {"ori_query": question}, "status": "complete"}],
        }
        body = json.dumps(payload, separators=(",", ":"), ensure_ascii=False)
        params["sign"] = self._get_sign("POST", self.url, params, body)
        qs = "&".join(f"{k}={quote(str(params[k]), safe='')}" for k in sorted(params))
        full_url = URL(f"{self.url}?{qs}", encoded=True)

        raw_text = ""
        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=180)) as session:
            async with session.post(full_url, headers=self._build_headers(), data=body, ssl=False) as resp:
                buf = ""
                async for chunk in resp.content.iter_any():
                    buf += chunk.decode("utf-8", errors="replace")
                    raw_text = buf

        return raw_text, body

    async def create_share(self, session_id):
        import aiohttp
        from yarl import URL

        params = self._build_params()
        body = json.dumps({"session_id": session_id, "share_type": 1}, separators=(",", ":"))
        params["sign"] = self._get_sign("POST", self.share_url, params, body)
        qs = "&".join(f"{k}={quote(str(params[k]), safe='')}" for k in sorted(params))
        full_url = URL(f"{self.share_url}?{qs}", encoded=True)

        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=30)) as session:
            async with session.post(full_url, headers=self._build_headers(), data=body, ssl=False) as resp:
                data = await resp.json()
                if data.get("code") == 0 and data.get("data"):
                    return data["data"].get("share_id", "")
        return ""


async def extract_qianwen_cookies():
    from playwright.async_api import async_playwright
    from .base import CDP_URL
    try:
        async with async_playwright() as pw:
            browser = await pw.chromium.connect_over_cdp(CDP_URL)
            ctx = browser.contexts[0]
            cookies = await ctx.cookies([
                "https://www.qianwen.com",
                "https://chat2.qianwen.com",
                "https://chat2-api.qianwen.com",
            ])
            if any(c["name"] == "tongyi_sso_ticket" for c in cookies):
                return cookies
    except Exception:
        pass
    return None


def get_qw_session_id(result):
    raw_req = result.get("_raw_req", "")
    try:
        req = json.loads(raw_req) if raw_req else {}
        return req.get("session_id", "")
    except Exception:
        return ""