Anush
e7626ba0c1
Refactors image mode logic to unify text-to-image and image-to-image header and payload construction, ensuring both use the same extended headers and 67-element payload structure. Enhances robustness of image container extraction to support both list and dict-based formats (e.g., gemini-3-pro), improving compatibility with different model responses. Adds an 'auth_method' parameter to get_access_token to allow strict manual cookie usage or fallback to browser cookies, and updates GeminiClient to set this based on initialization. Bumps version to 1.0.6.
213 lines
7.9 KiB
Python
213 lines
7.9 KiB
Python
import os
|
|
import re
|
|
import asyncio
|
|
from asyncio import Task
|
|
from pathlib import Path
|
|
|
|
from httpx import AsyncClient, Response
|
|
|
|
from ..constants import Endpoint, Headers
|
|
from ..exceptions import AuthError
|
|
from .load_browser_cookies import load_browser_cookies
|
|
from .logger import logger
|
|
|
|
|
|
async def send_request(
|
|
cookies: dict, proxy: str | None = None
|
|
) -> tuple[Response | None, dict]:
|
|
"""
|
|
Send http request with provided cookies.
|
|
"""
|
|
|
|
async with AsyncClient(
|
|
proxy=proxy,
|
|
headers=Headers.GEMINI.value,
|
|
cookies=cookies,
|
|
follow_redirects=True,
|
|
verify=False,
|
|
) as client:
|
|
response = await client.get(Endpoint.INIT.value)
|
|
response.raise_for_status()
|
|
return response, cookies
|
|
|
|
|
|
async def get_access_token(
|
|
base_cookies: dict, proxy: str | None = None, verbose: bool = False,
|
|
auth_method: str | None = None
|
|
) -> tuple[str, dict]:
|
|
"""
|
|
Send a get request to gemini.google.com for each group of available cookies and return
|
|
the value of "SNlM0e" as access token on the first successful request.
|
|
|
|
Cookie sources depend on auth_method:
|
|
- "manual": Only use base_cookies (no browser fallback)
|
|
- "auto_cookies" or None: Try base cookies, cache, and browser in order
|
|
|
|
Parameters
|
|
----------
|
|
base_cookies : `dict`
|
|
Base cookies to be used in the request.
|
|
proxy: `str`, optional
|
|
Proxy URL.
|
|
verbose: `bool`, optional
|
|
If `True`, will print more infomation in logs.
|
|
auth_method: `str`, optional
|
|
Authentication method: "manual" (strict) or "auto_cookies" (try all sources).
|
|
If None, defaults to "auto_cookies" behavior.
|
|
|
|
Returns
|
|
-------
|
|
`str`
|
|
Access token.
|
|
`dict`
|
|
Cookies of the successful request.
|
|
|
|
Raises
|
|
------
|
|
`gemini_webapi.AuthError`
|
|
If all requests failed.
|
|
"""
|
|
|
|
async with AsyncClient(proxy=proxy, follow_redirects=True, verify=False) as client:
|
|
response = await client.get(Endpoint.GOOGLE.value)
|
|
|
|
extra_cookies = {}
|
|
if response.status_code == 200:
|
|
extra_cookies = response.cookies
|
|
|
|
tasks = []
|
|
|
|
# Base cookies passed directly on initializing client
|
|
if "__Secure-1PSID" in base_cookies and "__Secure-1PSIDTS" in base_cookies:
|
|
tasks.append(Task(send_request({**extra_cookies, **base_cookies}, proxy=proxy)))
|
|
elif verbose:
|
|
logger.debug(
|
|
"Skipping loading base cookies. Either __Secure-1PSID or __Secure-1PSIDTS is not provided."
|
|
)
|
|
|
|
# Cached cookies in local file
|
|
cache_dir = (
|
|
(GEMINI_COOKIE_PATH := os.getenv("GEMINI_COOKIE_PATH"))
|
|
and Path(GEMINI_COOKIE_PATH)
|
|
or (Path(__file__).parent / "temp")
|
|
)
|
|
if "__Secure-1PSID" in base_cookies:
|
|
filename = f".cached_1psidts_{base_cookies['__Secure-1PSID']}.txt"
|
|
cache_file = cache_dir / filename
|
|
if cache_file.is_file():
|
|
cached_1psidts = cache_file.read_text()
|
|
if cached_1psidts:
|
|
cached_cookies = {
|
|
**extra_cookies,
|
|
**base_cookies,
|
|
"__Secure-1PSIDTS": cached_1psidts,
|
|
}
|
|
tasks.append(Task(send_request(cached_cookies, proxy=proxy)))
|
|
elif verbose:
|
|
logger.debug("Skipping loading cached cookies. Cache file is empty.")
|
|
elif verbose:
|
|
logger.debug("Skipping loading cached cookies. Cache file not found.")
|
|
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 = {
|
|
**extra_cookies,
|
|
"__Secure-1PSID": cache_file.stem[16:],
|
|
"__Secure-1PSIDTS": cached_1psidts,
|
|
}
|
|
tasks.append(Task(send_request(cached_cookies, proxy=proxy)))
|
|
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)
|
|
# Only check browser cookies if auth_method is NOT "manual"
|
|
browser_failure_reason = ""
|
|
if auth_method != "manual":
|
|
try:
|
|
valid_browser_cookies = 0
|
|
# Use '.google.com' to match all Google subdomains (gemini.google.com, etc.)
|
|
browser_cookies, browser_failure_reason = load_browser_cookies(
|
|
domain_name=".google.com", verbose=verbose
|
|
)
|
|
if browser_cookies:
|
|
for browser, cookies in browser_cookies.items():
|
|
if secure_1psid := cookies.get("__Secure-1PSID"):
|
|
if (
|
|
"__Secure-1PSID" in base_cookies
|
|
and base_cookies["__Secure-1PSID"] != secure_1psid
|
|
):
|
|
if verbose:
|
|
logger.debug(
|
|
f"Skipping loading local browser cookies from {browser}. "
|
|
f"__Secure-1PSID does not match the one provided."
|
|
)
|
|
continue
|
|
|
|
local_cookies = {"__Secure-1PSID": secure_1psid}
|
|
if secure_1psidts := cookies.get("__Secure-1PSIDTS"):
|
|
local_cookies["__Secure-1PSIDTS"] = secure_1psidts
|
|
if nid := cookies.get("NID"):
|
|
local_cookies["NID"] = nid
|
|
tasks.append(Task(send_request(local_cookies, proxy=proxy)))
|
|
valid_browser_cookies += 1
|
|
if verbose:
|
|
logger.debug(f"Loaded local browser cookies from {browser}")
|
|
|
|
if valid_browser_cookies == 0 and verbose:
|
|
logger.debug(
|
|
"Skipping loading local browser cookies. Login to gemini.google.com in your browser first."
|
|
)
|
|
except ImportError:
|
|
if verbose:
|
|
logger.debug(
|
|
"Skipping loading local browser cookies. Optional dependency 'browser-cookie3' is not installed."
|
|
)
|
|
except Exception as e:
|
|
if verbose:
|
|
logger.warning(f"Skipping loading local browser cookies. {e}")
|
|
elif verbose:
|
|
logger.debug("Skipping browser cookies check (auth_method='manual')")
|
|
|
|
if not tasks:
|
|
# Include the detailed browser failure reason if available
|
|
if browser_failure_reason:
|
|
raise AuthError(
|
|
f"Cookie Authentication Failed!\n\n{browser_failure_reason}"
|
|
)
|
|
else:
|
|
raise AuthError(
|
|
"No valid cookies available for initialization. Please pass __Secure-1PSID and __Secure-1PSIDTS manually."
|
|
)
|
|
|
|
for i, future in enumerate(asyncio.as_completed(tasks)):
|
|
try:
|
|
response, request_cookies = await future
|
|
match = re.search(r'"SNlM0e":"(.*?)"', response.text)
|
|
if match:
|
|
if verbose:
|
|
logger.debug(
|
|
f"Init attempt ({i + 1}/{len(tasks)}) succeeded. Initializing client..."
|
|
)
|
|
return match.group(1), request_cookies
|
|
elif verbose:
|
|
logger.debug(
|
|
f"Init attempt ({i + 1}/{len(tasks)}) failed. Cookies invalid."
|
|
)
|
|
except Exception as e:
|
|
if verbose:
|
|
logger.debug(
|
|
f"Init attempt ({i + 1}/{len(tasks)}) failed with error: {e}"
|
|
)
|
|
|
|
raise AuthError(
|
|
"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)})"
|
|
)
|