summaryrefslogtreecommitdiffstats
path: root/g4f/cookies.py
blob: 0c62d697ecfcb7eeb7649b11f8144e7ae8792f93 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
from __future__ import annotations

import os
import time
import json

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
    )

    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)

    browsers = [
        _g4f,
        chrome, chromium, opera, opera_gx,
        brave, edge, vivaldi, firefox,
    ]
    has_browser_cookie3 = True
except ImportError:
    has_browser_cookie3 = False
    browsers = []

from .typing import Dict, Cookies
from .errors import MissingRequirementsError
from . import debug

class CookiesConfig():
    cookies: Dict[str, Cookies] = {}
    cookies_dir: str = "./har_and_cookies"

DOMAINS = [
    ".bing.com",
    ".meta.ai",
    ".google.com",
    "www.whiterabbitneo.com",
    "huggingface.co",
    "chat.reka.ai",
    "chatgpt.com"
]

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, single_browser: bool = False) -> 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 CookiesConfig.cookies:
        return CookiesConfig.cookies[domain_name]

    cookies = load_cookies_from_browsers(domain_name, raise_requirements_error, single_browser)
    CookiesConfig.cookies[domain_name] = cookies
    return cookies

def set_cookies(domain_name: str, cookies: Cookies = None) -> None:
    if cookies:
        CookiesConfig.cookies[domain_name] = cookies
    elif domain_name in CookiesConfig.cookies:
        CookiesConfig.cookies.pop(domain_name)

def load_cookies_from_browsers(domain_name: str, raise_requirements_error: bool = True, single_browser: bool = False) -> 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:
                    if not cookie.expires or cookie.expires > time.time():
                        cookies[cookie.name] = cookie.value
            if single_browser and len(cookie_jar):
                break
        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 set_cookies_dir(dir: str) -> None:
    CookiesConfig.cookies_dir = dir

def get_cookies_dir() -> str:
    return CookiesConfig.cookies_dir

def read_cookie_files(dirPath: str = None):
    def get_domain(v: dict) -> str:
        host = [h["value"] for h in v['request']['headers'] if h["name"].lower() in ("host", ":authority")]
        if not host:
            return
        host = host.pop()
        for d in DOMAINS:
            if d in host:
                return d

    harFiles = []
    cookieFiles = []
    for root, _, files in os.walk(CookiesConfig.cookies_dir if dirPath is None else dirPath):
        for file in files:
            if file.endswith(".har"):
                harFiles.append(os.path.join(root, file))
            elif file.endswith(".json"):
                cookieFiles.append(os.path.join(root, file))

    CookiesConfig.cookies = {}
    for path in harFiles:
        with open(path, 'rb') as file:
            try:
                harFile = json.load(file)
            except json.JSONDecodeError:
                # Error: not a HAR file!
                continue
            if debug.logging:
                print("Read .har file:", path)
            new_cookies = {}
            for v in harFile['log']['entries']:
                domain = get_domain(v)
                if domain is None:
                    continue
                v_cookies = {}
                for c in v['request']['cookies']:
                    v_cookies[c['name']] = c['value']
                if len(v_cookies) > 0:
                    CookiesConfig.cookies[domain] = v_cookies
                    new_cookies[domain] = len(v_cookies)
            if debug.logging:
                for domain, new_values in new_cookies.items():
                    print(f"Cookies added: {new_values} from {domain}")
    for path in cookieFiles:
        with open(path, 'rb') as file:
            try:
                cookieFile = json.load(file)
            except json.JSONDecodeError:
                # Error: not a json file!
                continue
            if not isinstance(cookieFile, list):
                continue
            if debug.logging:
                print("Read cookie file:", path)
            new_cookies = {}
            for c in cookieFile:
                if isinstance(c, dict) and "domain" in c:
                    if c["domain"] not in new_cookies:
                        new_cookies[c["domain"]] = {}
                    new_cookies[c["domain"]][c["name"]] = c["value"]
            for domain, new_values in new_cookies.items():
                if debug.logging:
                    print(f"Cookies added: {len(new_values)} from {domain}")
                CookiesConfig.cookies[domain] = new_values

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)