summaryrefslogtreecommitdiffstats
path: root/g4f/cookies.py
diff options
context:
space:
mode:
Diffstat (limited to 'g4f/cookies.py')
-rw-r--r--g4f/cookies.py97
1 files changed, 97 insertions, 0 deletions
diff --git a/g4f/cookies.py b/g4f/cookies.py
new file mode 100644
index 00000000..b5c869e7
--- /dev/null
+++ b/g4f/cookies.py
@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import os
+
+try:
+ from platformdirs import user_config_dir
+ has_platformdirs = True
+except ImportError:
+ has_platformdirs = False
+try:
+ from browser_cookie3 import (
+ chrome, chromium, opera, opera_gx,
+ brave, edge, vivaldi, firefox,
+ _LinuxPasswordManager, BrowserCookieError
+ )
+ has_browser_cookie3 = True
+except ImportError:
+ has_browser_cookie3 = False
+
+from .typing import Dict, Cookies
+from .errors import MissingRequirementsError
+from . import debug
+
+# Global variable to store cookies
+_cookies: Dict[str, Cookies] = {}
+
+if has_browser_cookie3 and os.environ.get('DBUS_SESSION_BUS_ADDRESS') == "/dev/null":
+ _LinuxPasswordManager.get_password = lambda a, b: b"secret"
+
+def get_cookies(domain_name: str = '', raise_requirements_error: bool = True) -> Dict[str, str]:
+ """
+ Load cookies for a given domain from all supported browsers and cache the results.
+
+ Args:
+ domain_name (str): The domain for which to load cookies.
+
+ Returns:
+ Dict[str, str]: A dictionary of cookie names and values.
+ """
+ if domain_name in _cookies:
+ return _cookies[domain_name]
+
+ cookies = load_cookies_from_browsers(domain_name, raise_requirements_error)
+ _cookies[domain_name] = cookies
+ return cookies
+
+def set_cookies(domain_name: str, cookies: Cookies = None) -> None:
+ if cookies:
+ _cookies[domain_name] = cookies
+ elif domain_name in _cookies:
+ _cookies.pop(domain_name)
+
+def load_cookies_from_browsers(domain_name: str, raise_requirements_error: bool = True) -> Cookies:
+ """
+ Helper function to load cookies from various browsers.
+
+ Args:
+ domain_name (str): The domain for which to load cookies.
+
+ Returns:
+ Dict[str, str]: A dictionary of cookie names and values.
+ """
+ if not has_browser_cookie3:
+ if raise_requirements_error:
+ raise MissingRequirementsError('Install "browser_cookie3" package')
+ return {}
+ cookies = {}
+ for cookie_fn in [_g4f, chrome, chromium, opera, opera_gx, brave, edge, vivaldi, firefox]:
+ try:
+ cookie_jar = cookie_fn(domain_name=domain_name)
+ if len(cookie_jar) and debug.logging:
+ print(f"Read cookies from {cookie_fn.__name__} for {domain_name}")
+ for cookie in cookie_jar:
+ if cookie.name not in cookies:
+ cookies[cookie.name] = cookie.value
+ except BrowserCookieError:
+ pass
+ except Exception as e:
+ if debug.logging:
+ print(f"Error reading cookies from {cookie_fn.__name__} for {domain_name}: {e}")
+ return cookies
+
+def _g4f(domain_name: str) -> list:
+ """
+ Load cookies from the 'g4f' browser (if exists).
+
+ Args:
+ domain_name (str): The domain for which to load cookies.
+
+ Returns:
+ list: List of cookies.
+ """
+ if not has_platformdirs:
+ return []
+ user_data_dir = user_config_dir("g4f")
+ cookie_file = os.path.join(user_data_dir, "Default", "Cookies")
+ return [] if not os.path.exists(cookie_file) else chrome(cookie_file, domain_name) \ No newline at end of file