From 1de92200f0e54baad9480791016167643901ec96 Mon Sep 17 00:00:00 2001 From: Seto <61304189@qq.com> Date: Thu, 28 Dec 2023 12:57:28 +0800 Subject: [PATCH] Use the whole cookies to avoid the "SNlM0e value not found" error. --- bardapi/core.py | 13 ++++++++++--- bardapi/core_async.py | 32 ++++++++++++++++++++++---------- bardapi/utils.py | 25 +++++++++++++------------ 3 files changed, 45 insertions(+), 25 deletions(-) diff --git a/bardapi/core.py b/bardapi/core.py index b00f87ca5..486e30d8a 100644 --- a/bardapi/core.py +++ b/bardapi/core.py @@ -76,7 +76,7 @@ def __init__( self.conversation_id = conversation_id or "" self.response_id = "" self.choice_id = "" - self.session = self._get_session(session) + self.session = self._get_session(session, token_from_browser) self.SNlM0e = self._get_snim0e() self.language = language or os.getenv("_BARD_API_LANG") self.run_code = run_code @@ -116,7 +116,7 @@ def _get_token(self, token: str, token_from_browser: bool) -> str: "Bard API Key must be provided as token argument or extracted from browser." ) - def _get_session(self, session: Optional[requests.Session]) -> requests.Session: + def _get_session(self, session: Optional[requests.Session], token_from_browser) -> requests.Session: """ Get the requests Session object. @@ -129,7 +129,14 @@ def _get_session(self, session: Optional[requests.Session]) -> requests.Session: if session is None: new_session = requests.Session() new_session.headers = SESSION_HEADERS - new_session.cookies.set("__Secure-1PSID", self.token) + if(token_from_browser): + extracted_cookie_dict = extract_bard_cookie(cookies=False) + if not extracted_cookie_dict: + raise Exception("Failed to extract cookie from browsers.") + for e in extracted_cookie_dict: + new_session.cookies.set(e, extracted_cookie_dict[e]) + else: + new_session.cookies.set("__Secure-1PSID", self.token) new_session.proxies = self.proxies return new_session else: diff --git a/bardapi/core_async.py b/bardapi/core_async.py index 7e17e721b..9d4b72cfc 100644 --- a/bardapi/core_async.py +++ b/bardapi/core_async.py @@ -72,7 +72,7 @@ def __init__( self.conversation_id = conversation_id or "" self.response_id = "" self.choice_id = "" - self.client = self._get_client(client) # Creating an httpx async client for asynchronous core code + self.client = self._get_client(client, token_from_browser) # Creating an httpx async client for asynchronous core code self.language = language self.cookie_dict = {"__Secure-1PSID": self.token} self.run_code = run_code or False @@ -95,7 +95,7 @@ def _get_snim0e(self): :return: The SNlM0e value as a string. """ - if isinstance(self.SNlM0e, str): + if hasattr(self, 'SNlM0e') and isinstance(self.SNlM0e, str): return self.SNlM0e if not self.token or self.token[-1] != ".": @@ -145,7 +145,7 @@ def _get_token(self, token: str, token_from_browser: bool) -> str: "Bard API Key must be provided as token argument or extracted from browser." ) - async def _get_client(self, session: Optional[AsyncClient]) -> AsyncClient: + async def _get_client(self, session: Optional[AsyncClient], token_from_browser) -> AsyncClient: """ The _get_snim0e function is used to get the SNlM0e value from the Bard website. @@ -156,13 +156,25 @@ async def _get_client(self, session: Optional[AsyncClient]) -> AsyncClient: :return: (`str`) The SNlM0e value """ if session is None: - async_client = AsyncClient( - http2=True, - headers=SESSION_HEADERS, - cookies={"__Secure-1PSID": self.token}, - timeout=self.timeout, - proxies=self.proxies, - ) + if(token_from_browser): + extracted_cookie_dict = extract_bard_cookie(cookies=False) + if not extracted_cookie_dict: + raise Exception("Failed to extract cookie from browsers.") + async_client = AsyncClient( + http2=True, + headers=SESSION_HEADERS, + cookies=extracted_cookie_dict, + timeout=self.timeout, + proxies=self.proxies, + ) + else: + async_client = AsyncClient( + http2=True, + headers=SESSION_HEADERS, + cookies={"__Secure-1PSID": self.token}, + timeout=self.timeout, + proxies=self.proxies, + ) return async_client else: assert type(session)==AsyncClient diff --git a/bardapi/utils.py b/bardapi/utils.py index de32d92b1..becc2159e 100644 --- a/bardapi/utils.py +++ b/bardapi/utils.py @@ -103,18 +103,19 @@ def extract_bard_cookie(cookies: bool = False) -> dict: cj = browser_fn(domain_name=".google.com") for cookie in cj: - print(cookie.name) - if cookie.name == "__Secure-1PSID" and cookie.value.endswith("."): - cookie_dict["__Secure-1PSID"] = cookie.value - if cookies: - if cookie.name == "__Secure-1PSIDTS": - print(cookie.value) - cookie_dict["__Secure-1PSIDTS"] = cookie.value - elif cookie.name == "__Secure-1PSIDCC": - print(cookie.value) - cookie_dict["__Secure-1PSIDCC"] = cookie.value - if len(cookie_dict) == 3: - return cookie_dict + cookie_dict[cookie.name] = cookie.value + # print(cookie.name) + # if cookie.name == "__Secure-1PSID" and cookie.value.endswith("."): + # cookie_dict["__Secure-1PSID"] = cookie.value + # if cookies: + # if cookie.name == "__Secure-1PSIDTS": + # print(cookie.value) + # cookie_dict["__Secure-1PSIDTS"] = cookie.value + # elif cookie.name == "__Secure-1PSIDCC": + # print(cookie.value) + # cookie_dict["__Secure-1PSIDCC"] = cookie.value + # if len(cookie_dict) == 3: + # return cookie_dict except Exception as e: # Ignore exceptions and try the next browser function continue