mirror of
https://github.com/zc-zhangchen/any-auto-register.git
synced 2026-05-10 01:04:15 +08:00
115 lines
4.4 KiB
Python
115 lines
4.4 KiB
Python
"""Trae.ai 平台插件"""
|
||
from core.base_platform import BasePlatform, Account, AccountStatus, RegisterConfig
|
||
from core.base_mailbox import BaseMailbox
|
||
from core.registry import register
|
||
|
||
|
||
@register
|
||
class TraePlatform(BasePlatform):
|
||
name = "trae"
|
||
display_name = "Trae.ai"
|
||
version = "1.0.0"
|
||
|
||
def __init__(self, config: RegisterConfig = None, mailbox: BaseMailbox = None):
|
||
super().__init__(config)
|
||
self.mailbox = mailbox
|
||
|
||
def register(self, email: str, password: str = None) -> Account:
|
||
from platforms.trae.core import TraeRegister
|
||
log = getattr(self, '_log_fn', print)
|
||
|
||
mail_acct = self.mailbox.get_email() if self.mailbox else None
|
||
email = email or (mail_acct.email if mail_acct else None)
|
||
log(f"邮箱: {email}")
|
||
before_ids = self.mailbox.get_current_ids(mail_acct) if mail_acct else set()
|
||
|
||
def otp_cb():
|
||
log("等待验证码...")
|
||
code = self.mailbox.wait_for_code(mail_acct, keyword="", before_ids=before_ids)
|
||
if code: log(f"验证码: {code}")
|
||
return code
|
||
|
||
with self._make_executor() as ex:
|
||
reg = TraeRegister(executor=ex, log_fn=log)
|
||
result = reg.register(
|
||
email=email,
|
||
password=password,
|
||
otp_callback=otp_cb if self.mailbox else None,
|
||
)
|
||
|
||
return Account(
|
||
platform="trae",
|
||
email=result["email"],
|
||
password=result["password"],
|
||
user_id=result["user_id"],
|
||
token=result["token"],
|
||
region=result["region"],
|
||
status=AccountStatus.REGISTERED,
|
||
extra={"cashier_url": result["cashier_url"],
|
||
"ai_pay_host": result["ai_pay_host"]},
|
||
)
|
||
|
||
def check_valid(self, account: Account) -> bool:
|
||
return bool(account.token)
|
||
|
||
def get_platform_actions(self) -> list:
|
||
"""返回平台支持的操作列表"""
|
||
return [
|
||
{"id": "switch_account", "label": "切换到桌面应用", "params": []},
|
||
{"id": "get_user_info", "label": "获取用户信息", "params": []},
|
||
{"id": "get_cashier_url", "label": "获取升级链接", "params": []},
|
||
]
|
||
|
||
def execute_action(self, action_id: str, account: Account, params: dict) -> dict:
|
||
"""执行平台操作"""
|
||
if action_id == "switch_account":
|
||
from platforms.trae.switch import switch_trae_account, restart_trae_ide
|
||
|
||
token = account.token
|
||
user_id = account.user_id or ""
|
||
email = account.email or ""
|
||
region = account.region or ""
|
||
|
||
if not token:
|
||
return {"ok": False, "error": "账号缺少 token"}
|
||
|
||
ok, msg = switch_trae_account(token, user_id, email, region)
|
||
if not ok:
|
||
return {"ok": False, "error": msg}
|
||
|
||
restart_ok, restart_msg = restart_trae_ide()
|
||
return {
|
||
"ok": True,
|
||
"data": {
|
||
"message": f"{msg}。{restart_msg}" if restart_ok else msg,
|
||
}
|
||
}
|
||
|
||
elif action_id == "get_user_info":
|
||
from platforms.trae.switch import get_trae_user_info
|
||
|
||
token = account.token
|
||
if not token:
|
||
return {"ok": False, "error": "账号缺少 token"}
|
||
|
||
user_info = get_trae_user_info(token)
|
||
if user_info:
|
||
return {"ok": True, "data": user_info}
|
||
return {"ok": False, "error": "获取用户信息失败"}
|
||
|
||
elif action_id == "get_cashier_url":
|
||
from platforms.trae.core import TraeRegister
|
||
with self._make_executor() as ex:
|
||
reg = TraeRegister(executor=ex)
|
||
# 重新登录刷新 session，再获取新 token 和 cashier_url
|
||
reg.step4_trae_login()
|
||
token = reg.step5_get_token()
|
||
if not token:
|
||
token = account.token
|
||
cashier_url = reg.step7_create_order(token)
|
||
if not cashier_url:
|
||
return {"ok": False, "error": "获取升级链接失败，token 可能已过期，请重新注册"}
|
||
return {"ok": True, "data": {"cashier_url": cashier_url, "message": "请在浏览器中打开升级链接完成 Pro 订阅"}}
|
||
|
||
raise NotImplementedError(f"未知操作: {action_id}")
|