#!/usr/bin/env python3 """web_login.py - Trae CN Relay 本地授权回调监听器 使用方式: python web_login.py [--relay http://192.168.5.246:9563] [--port 8765] 依赖: 仅用 Python 标准库。 """ import argparse import http.server import json import os import sys import urllib.parse import urllib.request import uuid import webbrowser from pathlib import Path def parse_args(): parser = argparse.ArgumentParser(description="Trae CN Relay 本地授权回调监听器") parser.add_argument("--relay", default="http://192.168.5.246:9563", help="Relay 地址") parser.add_argument( "--key", default=os.environ.get("RELAY_API_KEY", ""), help="relay 的 API key（RELAY_API_KEYS 中的任意一个）；也可用环境变量 RELAY_API_KEY", ) parser.add_argument("--port", type=int, default=8765, help="本地监听端口") parser.add_argument("--client-id", default="ono9krqynydwx5", help="Trae Client ID") parser.add_argument("--auth-url", default="https://www.trae.cn/authorization", help="Trae 授权 URL") parser.add_argument("--no-open", action="store_true", help="不自动打开浏览器") return parser.parse_args() def json_loads_safe(value: str) -> dict: if not value: return {} try: return json.loads(value) except (json.JSONDecodeError, TypeError): return {} def parse_oauth_params(query: dict) -> dict: """解析 Trae 授权回调参数。 Trae 网页授权页实际上会走两套流程： 1. 新流程 (code_challenge): callback 会带 authCodeInfo / code 等参数 2. 老流程 (refreshToken): callback 直接带 refreshToken=xxx 对于老流程，本地不会拥有 Cloud-IDE-JWT，需要拿到 refreshToken 后通过 oauth/ExchangeToken 向 api.trae.cn 兑换 Cloud-IDE-JWT。 """ user_jwt = json_loads_safe(query.get("userJwt", "")) token = user_jwt.get("Token") or user_jwt.get("token") or "" refresh = user_jwt.get("RefreshToken") or user_jwt.get("refreshToken") or query.get("refreshToken") or query.get("data") or "" token_exp = user_jwt.get("TokenExpireAt") or user_jwt.get("tokenExpireAt") or "" refresh_exp = user_jwt.get("RefreshExpireAt") or user_jwt.get("refreshExpireAt") or query.get("refreshExpireAt") or "" user_info = json_loads_safe(query.get("userInfo", "")) user_id = user_info.get("UserID") or user_info.get("userId") or user_info.get("userID") or query.get("userId") or "" region = user_info.get("Region") or user_info.get("region") or query.get("region") or "CN" ai_region = user_info.get("AIRegion") or user_info.get("aiRegion") or region client_id = user_jwt.get("ClientID") or user_jwt.get("clientId") or query.get("clientID") or query.get("clientId") or query.get("client_id") or "" uid = user_id or "" host = query.get("host") or user_info.get("Host") or user_info.get("host") or "" # 老流程：回调只带了 refreshToken，通过 oauth/ExchangeToken 兑换 # Cloud-IDE-JWT 和用户信息。 if not token and refresh: exchange = exchange_refresh_token( refresh_token=refresh, client_id=client_id or "ono9krqynydwx5", host=host, ) if exchange.get("token"): token = exchange["token"] refresh = exchange.get("refreshToken") or refresh token_exp = exchange.get("expiredAt") or token_exp refresh_exp = exchange.get("refreshExpiredAt") or refresh_exp client_id = exchange.get("clientId") or client_id host = exchange.get("host") or host user_id = exchange.get("userId") or user_id user_info = exchange.get("userInfo") or user_info region = exchange.get("region") or region ai_region = exchange.get("aiRegion") or region if not token: return {} return { "token": token, "refreshToken": refresh, "userId": user_id or "", "tenantId": user_info.get("TenantID") or user_info.get("tenantId") or "", "region": region, "aiRegion": ai_region, "host": host, "expiredAt": str(token_exp) if token_exp else "", "refreshExpiredAt": str(refresh_exp) if refresh_exp else "", "clientId": client_id, "webId": user_info.get("WebId") or user_info.get("webId") or uid, "bizUserId": user_info.get("BizUserId") or user_info.get("bizUserId") or uid, "userUniqueId": user_info.get("UserUniqueId") or user_info.get("userUniqueId") or uid, "scope": query.get("scope") or user_info.get("Scope") or user_info.get("scope") or "", "tenant": user_info.get("Tenant") or user_info.get("tenant") or "", "appLanguage": user_info.get("AppLanguage") or user_info.get("appLanguage") or "", "userRegion": query.get("userRegion") or user_info.get("UserRegion") or user_info.get("userRegion") or "", "userIdentity": user_info.get("UserIdentity") or user_info.get("userIdentity") or "", "screenName": user_info.get("ScreenName") or user_info.get("screenName") or "", } def exchange_refresh_token(refresh_token: str, client_id: str, host: str = "") -> dict: """使用 refreshToken 向 Trae CN 兑换 Cloud-IDE-JWT。 这与 Trae 官网的实现一致: POST https://api.trae.cn/cloudide/api/v3/trae/oauth/ExchangeToken {"ClientID":..., "RefreshToken":..., "ClientSecret":"-", "UserID":""} """ if not refresh_token: return {} base = host or "https://api.trae.cn" base = base.rstrip("/") url = base + "/cloudide/api/v3/trae/oauth/ExchangeToken" payload = { "ClientID": client_id, "RefreshToken": refresh_token, "ClientSecret": "-", "UserID": "", } try: req = urllib.request.Request( url, data=json.dumps(payload).encode("utf-8"), method="POST", ) req.add_header("Content-Type", "application/json") req.add_header("Accept", "application/json") with urllib.request.urlopen(req, timeout=20) as resp: body = resp.read().decode("utf-8") except urllib.error.HTTPError as e: try: body = e.read().decode("utf-8", errors="replace") except Exception: body = "" print(f"[web_login] ExchangeToken HTTP {e.code}: {body[:500]}") return {} except Exception as e: print(f"[web_login] ExchangeToken failed: {e}") return {} result = json_loads_safe(body) data = result.get("Result") or result.get("result") or result if not isinstance(data, dict): print(f"[web_login] ExchangeToken unexpected result: {body[:500]}") return {} token = data.get("Token") or data.get("token") or data.get("AccessToken") or data.get("accessToken") or "" if not token: print(f"[web_login] ExchangeToken missing Token: {body[:500]}") return {} refresh2 = data.get("RefreshToken") or data.get("refreshToken") or refresh_token user_info = json_loads_safe(str(data.get("UserInfo") or data.get("userInfo") or "")) return { "token": token, "refreshToken": refresh2, "expiredAt": str(data.get("TokenExpireAt") or data.get("tokenExpireAt") or ""), "refreshExpiredAt": str(data.get("RefreshExpireAt") or data.get("refreshExpireAt") or ""), "clientId": data.get("ClientID") or data.get("clientId") or client_id, "host": host or base, "userId": user_info.get("UserID") or user_info.get("userId") or data.get("UserID") or data.get("userId") or "", "userInfo": user_info, "region": user_info.get("Region") or user_info.get("region") or data.get("Region") or data.get("region") or "CN", "aiRegion": user_info.get("AIRegion") or user_info.get("aiRegion") or data.get("AIRegion") or data.get("aiRegion") or "", } def forward_to_relay(relay_url: str, creds: dict, api_key: str = "") -> tuple[bool, str]: """POST 凭证到 relay /api/web-auth，返回 (success, message)""" url = relay_url.rstrip("/") + "/api/web-auth" data = json.dumps(creds).encode("utf-8") req = urllib.request.Request(url, data=data, method="POST") req.add_header("Content-Type", "application/json") # /api/web-auth is gated like the rest of the console; without the key the # relay answers 401 and the credential never lands. if api_key: req.add_header("Authorization", "Bearer " + api_key) try: with urllib.request.urlopen(req, timeout=15) as resp: body = resp.read().decode("utf-8") result = json.loads(body) if result.get("success"): return True, "凭证已写入服务器" else: return False, result.get("error", "写入失败") except urllib.error.HTTPError as e: try: body = e.read().decode("utf-8", errors="replace") msg = json.loads(body).get("error", str(e)) except Exception: msg = str(e) return False, msg except Exception as e: return False, str(e) class RelayAuthHandler(http.server.BaseHTTPRequestHandler): relay_url = "http://192.168.5.246:9563" client_id = "ono9krqynydwx5" auth_url = "https://www.trae.cn/authorization" def log_message(self, fmt, *args): pass # 不输出请求日志 def _html(self, content: str) -> str: return f""" Trae CN Relay 授权
{content}
""" def _index_page(self) -> str: port = self.server.server_address[1] relay = self.relay_url return self._html(f"""

Trae CN Relay 授权

监听端口: {port}
中转站: {relay}

1. 确保浏览器已登录 trae.cn
2. 点击下方按钮

打开 trae.cn
等待授权中…
""") def do_GET(self): parsed = urllib.parse.urlparse(self.path) params = dict(urllib.parse.parse_qsl(parsed.query)) if parsed.path == "/healthz": self._send_json({"status": "ok", "relay": self.relay_url, "port": self.server.server_address[1]}) elif parsed.path == "/relay-url": self._send_json({"relay": self.relay_url}) elif parsed.path == "/authorize": self._handle_authorize(params) elif parsed.path in ("/", "/index.html", ""): self._send_html(self._index_page()) else: self.send_response(404) self.end_headers() self.wfile.write(b"Not found") def _send_json(self, obj: dict): body = json.dumps(obj).encode("utf-8") self.send_response(200) self.send_header("Content-Type", "application/json; charset=utf-8") self.send_header("Access-Control-Allow-Origin", "*") self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS") self.send_header("Access-Control-Allow-Headers", "Content-Type") self.send_header("Cache-Control", "no-store") self.send_header("Content-Length", str(len(body))) self.end_headers() self.wfile.write(body) def do_OPTIONS(self): self.send_response(204) self.send_header("Access-Control-Allow-Origin", "*") self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS") self.send_header("Access-Control-Allow-Headers", "Content-Type") self.send_header("Access-Control-Max-Age", "86400") self.end_headers() def _handle_authorize(self, query: dict): creds = parse_oauth_params(query) trace_id = query.get("loginTraceID") or query.get("login_trace_id") or "" if not creds.get("token"): html = self._oauth_result_page(False, "未收到有效的 userJwt，请确认已登录 trae.cn", trace_id) self._send_html(html) return # 转发到 relay success, msg = forward_to_relay(self.relay_url, creds, getattr(self, "api_key", "")) html = self._oauth_result_page(success, msg, trace_id) self._send_html(html) def _oauth_result_page(self, success: bool, message: str, login_trace_id: str = "") -> str: safe_msg = message.replace("&", "&").replace("<", "<").replace(">", ">").replace('"', """) safe_msg = message.replace("&", "&").replace("<", "<").replace(">", ">").replace('"', """).replace("'", "'") safe_trace = login_trace_id.replace("&", "&").replace("<", "<") return self._html(f"""

{"成功" if success else "失败"}

{safe_msg}

""") def _send_html(self, html: str): body = html.encode("utf-8") self.send_response(200) self.send_header("Content-Type", "text/html; charset=utf-8") self.send_header("Access-Control-Allow-Origin", "*") self.send_header("Cache-Control", "no-store") self.send_header("Content-Length", str(len(body))) self.end_headers() self.wfile.write(body) def main(): args = parse_args() relay = args.relay.rstrip("/") port = args.port client_id = args.client_id auth_url = args.auth_url handler = RelayAuthHandler handler.relay_url = relay handler.client_id = client_id handler.auth_url = auth_url handler.api_key = args.key.strip() server = http.server.HTTPServer(("127.0.0.1", port), handler) print(f"Trae CN Relay 授权回调监听器已启动") print(f" 本地监听: http://127.0.0.1:{port}") print(f" 中转站: {relay}") print(f" 授权 URL: {auth_url}") print(f" ClientID: {client_id}") if handler.api_key: print(" API key: 已设置（最后一位 %s）" % handler.api_key[-1]) else: print(" API key: 未设置 —— relay 已开启鉴权，授权结果会被 401 拒绝") print(" 请追加 --key <你的 RELAY_API_KEYS> 或设置环境变量 RELAY_API_KEY") print() print(f"请在浏览器中打开 http://127.0.0.1:{port}") print(f"或直接在中转站页面点击授权按钮") print() if not args.no_open: try: # Trae 授权完成后回调到本机 /authorize；打开 relay 控制台方便用户点授权。 webbrowser.open(f"{relay}/web/login") except Exception: pass try: server.serve_forever() except KeyboardInterrupt: print("\n已停止") sys.exit(0) if __name__ == "__main__": main()