From de26ed706218fb44516b224ce76ceac264427abd Mon Sep 17 00:00:00 2001 From: omg-xtao <100690902+omg-xtao@users.noreply.github.com> Date: Sat, 4 Nov 2023 16:03:41 +0800 Subject: [PATCH] :sparkles: Support get all token by stoken --- simnet/client/components/auth.py | 42 ++++++++++++++- simnet/client/cookies.py | 93 +++++++++++++++++++++++++++++++- 2 files changed, 132 insertions(+), 3 deletions(-) diff --git a/simnet/client/components/auth.py b/simnet/client/components/auth.py index 91aa3d4..42ef48d 100644 --- a/simnet/client/components/auth.py +++ b/simnet/client/components/auth.py @@ -2,6 +2,7 @@ import json as jsonlib from typing import Optional, Tuple, Union from simnet.client.base import BaseClient +from simnet.client.cookies import CookiesModel from simnet.client.routes import ( AUTH_URL, AUTH_KEY_URL, @@ -372,7 +373,7 @@ class AuthClient(BaseClient): "account_id": self.account_id, "game_token": game_token, } - headers = {"x-rpc-app_id": "bll8iq97cem8" if self.region == Region.CHINESE else "c9oqaq3s3gu8"} + headers = {"x-rpc-app_id": "bll8iq97cem8"} data = await self.request_lab(url, data=data, headers=headers) mid = data.get("user_info", {}).get("mid", "") stoken_v2 = data.get("token", {}).get("token", "") @@ -409,6 +410,45 @@ class AuthClient(BaseClient): self.cookies.set("stoken", stoken_v2) return stoken_v2, mid + async def get_all_token_by_stoken( + self, + stoken: Optional[str] = None, + account_id: Optional[int] = None, + mid: Optional[str] = None, + ) -> CookiesModel: + """ + Get stoken_v2, mid, ltoken and cookie_token by stoken v1 or v2 + + Args: + stoken (Optional[str]): The stoken_v1 to use to retrieve the stoken_v2 and mid. + If not provided, the `stoken` attribute value will be used. + account_id (Optional[int]): The account ID to use to retrieve the stoken_v2 and mid. + If not provided, the `account_id` attribute value will be used. + mid (Optional[str]): The mid to use to retrieve the stoken_v2 and mid. + + Returns: + CookiesModel: The stoken_v2, mid, ltoken and cookie_token. + """ + if self.region != Region.OVERSEAS: + raise RegionNotSupported() + self.check_stoken(stoken, account_id, mid) + url = AUTH_KEY_URL.get_url(self.region) / "../../../account/ma-passport/token/getBySToken" + headers = {"x-rpc-app_id": "c9oqaq3s3gu8"} + data_ = {"dst_token_types": [1, 2, 4]} + data = await self.request_lab(url, method="POST", headers=headers, data=data_) + tokens = data.get("tokens", []) + token_map = {1: "", 2: "", 4: ""} + for token in tokens: + token_type = token.get("token_type", 0) + token_map[token_type] = token.get("token", "") + mid = data.get("user_info", {}).get("mid", "") + model = CookiesModel(mid=mid, stoken=token_map[1], ltoken=token_map[2], cookie_token=token_map[4]) + self.cookies.set("mid", model.mid) + self.cookies.set("stoken", model.stoken) + self.cookies.set("ltoken", model.ltoken) + self.cookies.set("cookie_token", model.cookie_token) + return model + async def get_game_token_by_stoken( self, stoken: Optional[str] = None, diff --git a/simnet/client/cookies.py b/simnet/client/cookies.py index b4dffef..0cbff48 100644 --- a/simnet/client/cookies.py +++ b/simnet/client/cookies.py @@ -1,12 +1,19 @@ from http.cookiejar import CookieJar from http.cookies import SimpleCookie -from typing import Optional +from typing import Optional, TypeVar from httpx import Cookies as _Cookies from simnet.utils.types import CookieTypes -__all__ = ("Cookies",) +from pydantic import BaseModel + +IntStr = TypeVar("IntStr", int, str) + +__all__ = ( + "Cookies", + "CookiesModel", +) class Cookies(_Cookies): @@ -79,3 +86,85 @@ class Cookies(_Cookies): if value is None: return default return value + + +class CookiesModel(BaseModel, frozen=False): + """A model that represents the cookies used by the client.""" + + login_uid: Optional[IntStr] = None + login_ticket: Optional[str] = None + + stoken: Optional[str] = None + stuid: Optional[IntStr] = None + mid: Optional[str] = None + + account_id: Optional[IntStr] = None + cookie_token: Optional[str] = None + + ltoken: Optional[str] = None + ltuid: Optional[IntStr] = None + + account_mid_v2: Optional[str] = None + cookie_token_v2: Optional[str] = None + account_id_v2: Optional[IntStr] = None + + ltoken_v2: Optional[str] = None + ltmid_v2: Optional[str] = None + ltuid_v2: Optional[IntStr] = None + + @property + def is_v1(self) -> bool: + if self.account_id or self.cookie_token or self.ltoken or self.ltuid: + return True + return False + + @property + def is_v2(self) -> bool: + if self.account_mid_v2 or self.cookie_token_v2 or self.ltoken_v2 or self.ltmid_v2: + return True + return False + + def remove_v2(self): + """Remove the v2 cookies.""" + self.account_mid_v2 = None + self.cookie_token_v2 = None + self.ltoken_v2 = None + self.ltmid_v2 = None + + def to_dict(self): + """Return the cookies as a dictionary.""" + return self.dict(exclude_defaults=True) + + def to_json(self): + """Return the cookies as a JSON string.""" + return self.json(exclude_defaults=True) + + @property + def user_id(self) -> Optional[int]: + if self.ltuid: + return self.ltuid + if self.account_id: + return self.account_id + if self.login_uid: + return self.login_uid + if self.stuid: + return self.stuid + if self.account_id_v2: + return self.account_id_v2 + if self.ltuid_v2: + return self.ltuid_v2 + return None + + def set_v2_uid(self, user_id: int): + """Set the user ID for the v2 cookies.""" + if self.ltuid_v2 is None and self.ltoken_v2: + self.ltuid_v2 = user_id + if self.account_id_v2 is None and self.account_mid_v2: + self.account_id_v2 = user_id + + def set_uid(self, user_id: int): + """Set the user ID for the v1 cookies.""" + if self.account_id is None and self.cookie_token: + self.account_id = user_id + if self.ltuid is None and self.ltoken: + self.ltuid = user_id