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