Add support for agent pinning & authenticated WS connections

TikTokLive/client/client.py

@@ -103,7 +103,8 @@ async def start(
             fetch_room_info: bool = False,
             fetch_gift_info: bool = False,
             fetch_live_check: bool = True,
-            room_id: Optional[int] = None
+            room_id: Optional[int] = None,
+            preferred_agent_id: Optional[str] = None
     ) -> Task:
         """
         Create a non-blocking connection to TikTok LIVE and return the task
@@ -115,6 +116,7 @@ async def start(
         :param room_id: An override to the room ID to connect directly to the livestream and skip scraping the live.
                         Useful when trying to scale, as scraping the HTML can result in TikTok blocks.
         :param compress_ws_events: Whether to compress the WebSocket events using gzip compression (you should probably have this on)
+        :param preferred_agent_id: The preferred agent ID to use when connecting to the WebSocket
         :return: Task containing the heartbeat of the client
 
         """
@@ -152,7 +154,7 @@ async def start(
             self._gift_info = await self._web.fetch_gift_list()
 
         # <Required> Fetch the first response
-        initial_webcast_response: WebcastResponse = await self._web.fetch_signed_websocket()
+        initial_webcast_response: WebcastResponse = await self._web.fetch_signed_websocket(preferred_agent_id=preferred_agent_id)
 
         # Start the websocket connection & return it
         self._event_loop_task = self._asyncio_loop.create_task(
@@ -314,7 +316,8 @@ async def _ws_client_loop(
                 compress_ws_events=compress_ws_events,
                 cookies=self._web.cookies,
                 room_id=self._room_id,
-                user_agent=self._web.headers['User-Agent']
+                user_agent=self._web.headers['User-Agent'],
+                authenticate_websocket=self._web.authenticate_websocket
         ):
 
             # Iterate over the events extracted

TikTokLive/client/errors.py

@@ -73,6 +73,7 @@ class ErrorReason(enum.Enum):
         SIGN_NOT_200 = 4
         EMPTY_COOKIES = 5
         PREMIUM_ENDPOINT = 6
+        AUTHENTICATED_WS = 7
 
     def __init__(
             self,
@@ -183,3 +184,13 @@ def __init__(self, *args, api_message: str):
         _args.append(self.format_sign_server_message(api_message))
 
         super().__init__(SignAPIError.ErrorReason.PREMIUM_ENDPOINT, *_args)
+
+
+class AuthenticatedWebSocketConnectionError(SignAPIError):
+    """
+    Thrown when sending the session ID to the sign server as this is deemed a risky operation that could lead to an account being banned.
+
+    """
+
+    def __init__(self, *args):
+        super().__init__(SignAPIError.ErrorReason.AUTHENTICATED_WS, *args)

TikTokLive/client/web/routes/fetch_signed_websocket.py

@@ -1,11 +1,11 @@
 import os
 from http.cookies import SimpleCookie
-from typing import Optional, Union
+from typing import Optional
 
 import httpx
 from httpx import Response
 
-from TikTokLive.client.errors import SignAPIError, SignatureRateLimitError
+from TikTokLive.client.errors import SignAPIError, SignatureRateLimitError, AuthenticatedWebSocketConnectionError
 from TikTokLive.client.web.web_base import ClientRoute
 from TikTokLive.client.web.web_settings import WebDefaults, CLIENT_NAME
 from TikTokLive.client.ws.ws_utils import extract_webcast_response_message
@@ -20,22 +20,60 @@ class FetchSignedWebSocketRoute(ClientRoute):
 
     async def __call__(
             self,
-            room_id: Optional[int] = None
+            room_id: Optional[int] = None,
+            preferred_agent_id: Optional[int] = None,
+            session_id: Optional[str] = None
     ) -> WebcastResponse:
         """
         Call the method to get the first WebcastResponse (as bytes) to use to upgrade to WebSocket & perform the first ack
 
         :param room_id: Override the room ID to fetch the webcast for
+        :param preferred_agent_id: The preferred agent ID to use for the request
         :return: The WebcastResponse forwarded from the sign server proxy, as raw bytes
 
         """
 
-        extra_headers: dict = {}
-        extra_params: dict = {'client': CLIENT_NAME}
+        extra_headers: dict = {
+        }
+
+        extra_params: dict = {
+            'client': CLIENT_NAME,
+        }
 
         if room_id is not None:
             extra_params['room_id'] = room_id
 
+        if preferred_agent_id is not None:
+            extra_params['preferred_agent_id'] = preferred_agent_id
+
+        # The session ID we want to add to the request
+        session_id: str = session_id or self._web.cookies.get('sessionid')
+
+        if self._web.authenticate_websocket and session_id:
+            if not os.getenv('WHITELIST_AUTHENTICATED_SESSION_ID_HOST'):
+                raise AuthenticatedWebSocketConnectionError(
+                    "For your safety, this request has been BLOCKED. To understand why, see the reason below:\n\t"
+                    "You set 'authenticate_websocket' to True, which allows your Session ID to be sent to the Sign Server when connecting to TikTok LIVE.\n\t"
+                    "This is risky, because a session ID grants a user complete access to your account.\n\t"
+                    "You should ONLY enable this setting if you trust the Sign Server. The Euler Stream sign server does NOT store your session ID. Third party servers MAY."
+                    "\n\n\t>> THIRD PARTY SIGN SERVERS MAY STEAL YOUR SESSION ID. <<\n\n\t"
+                    "It should also be noted that since there are a limited number of sign servers, your session ID will\n\t"
+                    "connect to TikTok with the same IP address as other users. This could potentially lead to a ban of the account.\n\t"
+                    "With that said, there has never been a case of a ban due to this feature.\n\t"
+                    "You are only recommended to use this setting if you are aware of the risks and are willing to take them.\n\t"
+                    "If you are sure you want to enable this setting, set the environment variable 'WHITELIST_AUTHENTICATED_SESSION_ID_HOST' to the HOST you want to authorize (e.g. 'tiktok.eulerstream.com').\n\t"
+                    "By doing so, you acknowledge the risks and agree to take responsibility for any consequences."
+                )
+
+            if os.getenv('WHITELIST_AUTHENTICATED_SESSION_ID_HOST', '') != WebDefaults.tiktok_sign_url.split("://")[1]:
+                raise AuthenticatedWebSocketConnectionError(
+                    f"The host '{os.getenv('WHITELIST_AUTHENTICATED_SESSION_ID_HOST')}' you set in 'WHITELIST_AUTHENTICATED_SESSION_ID_HOST' does not match the host '{WebDefaults.tiktok_sign_url.split('://')[1]}' of the Sign Server. "
+                    f"Please set the correct host in 'WHITELIST_AUTHENTICATED_SESSION_ID_HOST' to authorize the Sign Server."
+                )
+
+            extra_params['session_id'] = session_id
+            self._logger.warning("Sending session ID to sign server for WebSocket connection. This is a risky operation.")
+
         # Add the API key if it exists
         if self._web.signer.sign_api_key is not None:
             extra_headers['X-Api-Key'] = self._web.signer.sign_api_key
@@ -52,6 +90,7 @@ async def __call__(
                 "Failed to connect to the sign server due to an httpx.ConnectError!"
             ) from ex
 
+        self._logger.debug(f"Sign API Fetched from agent {response.headers.get('X-Agent-Id')}: {response.status_code}")
         data: bytes = await response.aread()
 
         if response.status_code == 429:

TikTokLive/client/web/web_base.py

@@ -53,6 +53,20 @@ def __init__(
         # Special client for requests that check the TLS certificate
         self._curl_cffi: Optional[curl_cffi.requests.AsyncSession] = curl_cffi.requests.AsyncSession(**(curl_cffi_kwargs or {})) if SUPPORTS_CURL_CFFI else None
 
+        # Whether to allow using the session ID with the WebSocket connection
+        self._authenticate_websocket: bool = False
+
+    @property
+    def authenticate_websocket(self) -> bool:
+        """
+        Get whether to use the session ID with the WebSocket connection
+
+        :return: Whether to use the session ID with the WebSocket connection
+
+        """
+
+        return self._authenticate_websocket
+
     @property
     def httpx_client(self) -> AsyncClient:
         """
@@ -105,7 +119,6 @@ def _create_httpx_client(
             **httpx_kwargs.pop("cookies", {})
         })
 
-
         # Create the headers
         self.headers = {
             **WebDefaults.web_client_headers,
@@ -135,15 +148,27 @@ async def close(self) -> None:
         await self._httpx.aclose()
         await self._curl_cffi.close()
 
-    def set_session_id(self, session_id: str) -> None:
+    def set_session_id(
+            self,
+            session_id: str,
+            authenticate_websocket: bool = False
+    ) -> None:
         """
         Set the session id cookies for the HTTP client and Websocket connection
 
         :param session_id: The (must be valid) session ID
+
+        :param authenticate_websocket: Whether to use the session ID with the WebSocket connection.
+            Enable this at your own risk, as it passes the session ID to Euler Stream's servers.
+            This COULD lead to a hypothetical ban, but the odds are low & this has never occurred that we know of.
+
         :return: None
 
         """
 
+        # Set the authenticate WS setting
+        self._authenticate_websocket = authenticate_websocket
+
         self.cookies.set("sessionid", session_id)
         self.cookies.set("sessionid_ss", session_id)
         self.cookies.set("sid_tt", session_id)

TikTokLive/client/web/web_presets.py

@@ -30,12 +30,6 @@ class ScreenPreset(TypedDict):
 
 
 Locations: List[LocationPreset] = [
-    {
-        "lang_country": "en-US",
-        "lang": "en",
-        "country": "US",
-        "tz_name": "America/New_York",
-    },
     {
         "lang_country": "en-GB",
         "lang": "en",

TikTokLive/client/ws/ws_client.py

@@ -128,10 +128,47 @@ async def disconnect(self) -> None:
 
         await self.ws.close()
 
+    def get_ws_cookie_string(self, cookies: httpx.Cookies, authenticate_websocket: bool) -> str:
+        """
+        Get the cookie string for the WebSocket connection.
+
+        :param cookies: Cookies to pass to the WebSocket connection
+        :param authenticate_websocket: Whether the WebSocket is authenticated
+        :return: Cookie string
+
+        """
+
+        cookie_values = []
+        session_id: str | None = cookies.get("sessionid")
+
+        # Exclude all session ID's
+        for key, value in cookies.items():
+            if not authenticate_websocket:
+                if session_id in value:
+                    continue
+            cookie_values.append(f"{key}={value};")
+
+        # Create the cookie string
+        cookie_string: str = " ".join(cookie_values)
+
+        # Create a redacted SID & cookie string
+        redacted_sid = session_id[:8] + "*" * (len(session_id) - 8)
+        redacted_cookie_string: str = cookie_string.replace(session_id, redacted_sid)
+
+        # Log that we're creating a cookie string for a logged in session
+        if session_id is not None and authenticate_websocket:
+            self._logger.warning(f"Created WS Cookie string for a LOGGED IN TikTok LIVE WebSocket session (Session ID: {redacted_sid}). Cookies: {redacted_cookie_string}")
+        else:
+            self._logger.debug(f"Created WS Cookie string for an ANONYMOUS TikTok Live WebSocket session. Cookies: {redacted_cookie_string}")
+
+        # Return the cookie string
+        return cookie_string
+
     async def connect(
             self,
             room_id: int,
             cookies: httpx.Cookies,
+            authenticate_websocket: bool,
             user_agent: str,
             initial_webcast_response: WebcastResponse,
             process_connect_events: bool = True,
@@ -165,6 +202,7 @@ async def connect(
         :param room_id: The room ID to connect to
         :param user_agent: The user agent to pass to the WebSocket connection
         :param cookies: The cookies to pass to the WebSocket connection
+        :param authenticate_websocket: Whether to authenticate the WebSocket connection
         :param process_connect_events: Whether to process the initial events sent in the first fetch
         :param compress_ws_events: Whether to ask TikTok to gzip the WebSocket events
         :return: Yields WebcastResponseMessage, the messages within WebcastResponse.messages
@@ -201,7 +239,7 @@ async def connect(
             # Extra headers
             extra_headers={
                 # Must pass cookies to connect to the WebSocket
-                "Cookie": " ".join(f"{k}={v};" for k, v in cookies.items()),
+                "Cookie": self.get_ws_cookie_string(cookies=cookies, authenticate_websocket=authenticate_websocket),
                 "User-Agent": user_agent,
 
                 # Optional override for the headers
@@ -269,7 +307,6 @@ async def _ping_loop_fn(self) -> None:
         # Ping Loop
         try:
             while self.connected:
-
                 # Send the ping
                 await self.send(message=self.PING_MESSAGE)
 

setup.py

@@ -7,7 +7,7 @@
     "name": "TikTokLive",
     "license": "MIT",
     "author": "Isaac Kogan",
-    "version": "6.3.1",
+    "version": "6.4.0",
     "email": "info@isaackogan.com"
 }
 
@@ -47,7 +47,6 @@
             "mashumaro>=3.5",  # JSON Deserialization
             "protobuf3-to-dict>=0.1.5",
             "protobuf>=3.19.4",
-
         ],
         classifiers=[
             "Development Status :: 4 - Beta",