From 84021abcc8a450f12c6cebecedfb1ec66ad58d92 Mon Sep 17 00:00:00 2001
From: UZQueen <157540577+HanaokaYuzu@users.noreply.github.com>
Date: Wed, 11 Sep 2024 00:25:48 -0500
Subject: [PATCH] feat: when secure_1psid is not provided, loop through all
 local cached cookies instead of raising value error

ref #33
---
 src/gemini_webapi/client.py                 |  7 +---
 src/gemini_webapi/utils/get_access_token.py | 41 ++++++++++++++++-----
 tests/test_client_features.py               |  4 +-
 3 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/src/gemini_webapi/client.py b/src/gemini_webapi/client.py
index 35ec81c..f4d300d 100644
--- a/src/gemini_webapi/client.py
+++ b/src/gemini_webapi/client.py
@@ -83,8 +83,7 @@ class GeminiClient:
     Raises
     ------
     `ValueError`
-        If `secure_1psid` is not provided and optional dependency `browser-cookie3` is not installed, or
-        `browser-cookie3` is installed but cookies for google.com are not found in your local browser storage.
+        If `browser-cookie3` is installed but cookies for google.com are not found in your local browser storage.
     """
 
     __slots__ = [
@@ -132,9 +131,7 @@ def __init__(
                         "Failed to load cookies from local browser. Please pass cookie values manually."
                     )
             except ImportError:
-                raise ValueError(
-                    "'secure_1psid' must be provided if optional dependency 'browser-cookie3' is not installed."
-                )
+                pass
 
     async def init(
         self,
diff --git a/src/gemini_webapi/utils/get_access_token.py b/src/gemini_webapi/utils/get_access_token.py
index 5543517..38a3640 100644
--- a/src/gemini_webapi/utils/get_access_token.py
+++ b/src/gemini_webapi/utils/get_access_token.py
@@ -58,13 +58,21 @@ async def send_request(cookies: dict) -> tuple[Response | None, dict]:
 
     tasks = []
 
-    if "__Secure-1PSID" in base_cookies:
+    # Base cookies passed directly on initializing client
+    if "__Secure-1PSID" in base_cookies and "__Secure-1PSIDTS" in base_cookies:
         tasks.append(Task(send_request(base_cookies)))
+    elif verbose:
+        logger.debug(
+            "Skipping loading base cookies. Either __Secure-1PSID or __Secure-1PSIDTS is not provided."
+        )
 
+    # Cached cookies in local file
+    cache_dir = Path(__file__).parent / "temp"
+    if "__Secure-1PSID" in base_cookies:
         filename = f".cached_1psidts_{base_cookies['__Secure-1PSID']}.txt"
-        path = Path(__file__).parent / "temp" / filename
-        if path.is_file():
-            cached_1psidts = path.read_text()
+        cache_file = cache_dir / filename
+        if cache_file.is_file():
+            cached_1psidts = cache_file.read_text()
             if cached_1psidts:
                 cached_cookies = {**base_cookies, "__Secure-1PSIDTS": cached_1psidts}
                 tasks.append(Task(send_request(cached_cookies)))
@@ -72,11 +80,25 @@ async def send_request(cookies: dict) -> tuple[Response | None, dict]:
                 logger.debug("Skipping loading cached cookies. Cache file is empty.")
         elif verbose:
             logger.debug("Skipping loading cached cookies. Cache file not found.")
-    elif verbose:
-        logger.debug(
-            "Skipping loading base cookies and cached cookies. __Secure-1PSID is not provided."
-        )
+    else:
+        valid_caches = 0
+        cache_files = cache_dir.glob(".cached_1psidts_*.txt")
+        for cache_file in cache_files:
+            cached_1psidts = cache_file.read_text()
+            if cached_1psidts:
+                cached_cookies = {
+                    "__Secure-1PSID": cache_file.stem[16:],
+                    "__Secure-1PSIDTS": cached_1psidts,
+                }
+                tasks.append(Task(send_request(cached_cookies)))
+                valid_caches += 1
+
+        if valid_caches == 0 and verbose:
+            logger.debug(
+                "Skipping loading cached cookies. Cookies will be cached after successful initialization."
+            )
 
+    # Browser cookies (if browser-cookie3 is installed)
     try:
         browser_cookies = load_browser_cookies(
             domain_name="google.com", verbose=verbose
@@ -120,5 +142,6 @@ async def send_request(cookies: dict) -> tuple[Response | None, dict]:
                 )
 
     raise AuthError(
-        "Failed to initialize client. SECURE_1PSIDTS could get expired frequently, please make sure cookie values are up to date."
+        "Failed to initialize client. SECURE_1PSIDTS could get expired frequently, please make sure cookie values are up to date. "
+        f"(Failed initialization attempts: {len(tasks)})"
     )
diff --git a/tests/test_client_features.py b/tests/test_client_features.py
index 9bb1bf9..758e1e0 100644
--- a/tests/test_client_features.py
+++ b/tests/test_client_features.py
@@ -19,8 +19,8 @@ async def asyncSetUp(self):
 
         try:
             await self.geminiclient.init()
-        except AuthError:
-            self.skipTest("Test was skipped due to invalid cookies")
+        except AuthError as e:
+            self.skipTest(e)
 
     @logger.catch(reraise=True)
     async def test_successful_request(self):