|
3 | 3 | import datetime |
4 | 4 | import heapq |
5 | 5 | import itertools |
| 6 | +import json |
6 | 7 | import os # noqa |
7 | 8 | import pathlib |
8 | 9 | import pickle |
|
37 | 38 | _SIMPLE_COOKIE = SimpleCookie() |
38 | 39 |
|
39 | 40 |
|
| 41 | +class _RestrictedCookieUnpickler(pickle.Unpickler): |
| 42 | + """A restricted unpickler that only allows cookie-related types. |
| 43 | +
|
| 44 | + This prevents arbitrary code execution when loading pickled cookie data |
| 45 | + from untrusted sources. Only types that are expected in a serialized |
| 46 | + CookieJar are permitted. |
| 47 | +
|
| 48 | + See: https://docs.python.org/3/library/pickle.html#restricting-globals |
| 49 | + """ |
| 50 | + |
| 51 | + _ALLOWED_CLASSES: frozenset[tuple[str, str]] = frozenset( |
| 52 | + { |
| 53 | + # Core cookie types |
| 54 | + ("http.cookies", "SimpleCookie"), |
| 55 | + ("http.cookies", "Morsel"), |
| 56 | + # Container types used by CookieJar._cookies |
| 57 | + ("collections", "defaultdict"), |
| 58 | + # builtins that pickle uses for reconstruction |
| 59 | + ("builtins", "tuple"), |
| 60 | + ("builtins", "set"), |
| 61 | + ("builtins", "frozenset"), |
| 62 | + ("builtins", "dict"), |
| 63 | + } |
| 64 | + ) |
| 65 | + |
| 66 | + def find_class(self, module: str, name: str) -> type: |
| 67 | + if (module, name) not in self._ALLOWED_CLASSES: |
| 68 | + raise pickle.UnpicklingError( |
| 69 | + f"Forbidden class: {module}.{name}. " |
| 70 | + "CookieJar.load() only allows cookie-related types for security. " |
| 71 | + "See https://docs.python.org/3/library/pickle.html#restricting-globals" |
| 72 | + ) |
| 73 | + return super().find_class(module, name) # type: ignore[no-any-return] |
| 74 | + |
| 75 | + |
40 | 76 | class CookieJar(AbstractCookieJar): |
41 | 77 | """Implements cookie storage adhering to RFC 6265.""" |
42 | 78 |
|
@@ -112,14 +148,84 @@ def quote_cookie(self) -> bool: |
112 | 148 | return self._quote_cookie |
113 | 149 |
|
114 | 150 | def save(self, file_path: PathLike) -> None: |
| 151 | + """Save cookies to a file using JSON format. |
| 152 | +
|
| 153 | + :param file_path: Path to file where cookies will be serialized, |
| 154 | + :class:`str` or :class:`pathlib.Path` instance. |
| 155 | + """ |
115 | 156 | file_path = pathlib.Path(file_path) |
116 | | - with file_path.open(mode="wb") as f: |
117 | | - pickle.dump(self._cookies, f, pickle.HIGHEST_PROTOCOL) |
| 157 | + data: dict[str, dict[str, dict[str, str | bool]]] = {} |
| 158 | + for (domain, path), cookie in self._cookies.items(): |
| 159 | + key = f"{domain}|{path}" |
| 160 | + data[key] = {} |
| 161 | + for name, morsel in cookie.items(): |
| 162 | + morsel_data: dict[str, str | bool] = { |
| 163 | + "key": morsel.key, |
| 164 | + "value": morsel.value, |
| 165 | + "coded_value": morsel.coded_value, |
| 166 | + } |
| 167 | + # Save all morsel attributes that have values |
| 168 | + for attr in morsel._reserved: # type: ignore[attr-defined] |
| 169 | + attr_val = morsel[attr] |
| 170 | + if attr_val: |
| 171 | + morsel_data[attr] = attr_val |
| 172 | + data[key][name] = morsel_data |
| 173 | + with file_path.open(mode="w", encoding="utf-8") as f: |
| 174 | + json.dump(data, f, indent=2) |
118 | 175 |
|
119 | 176 | def load(self, file_path: PathLike) -> None: |
| 177 | + """Load cookies from a file. |
| 178 | +
|
| 179 | + Tries to load JSON format first. Falls back to loading legacy |
| 180 | + pickle format (using a restricted unpickler) for backward |
| 181 | + compatibility with existing cookie files. |
| 182 | +
|
| 183 | + :param file_path: Path to file from where cookies will be |
| 184 | + imported, :class:`str` or :class:`pathlib.Path` instance. |
| 185 | + """ |
120 | 186 | file_path = pathlib.Path(file_path) |
121 | | - with file_path.open(mode="rb") as f: |
122 | | - self._cookies = pickle.load(f) |
| 187 | + # Try JSON format first |
| 188 | + try: |
| 189 | + with file_path.open(mode="r", encoding="utf-8") as f: |
| 190 | + data = json.load(f) |
| 191 | + self._cookies = self._load_json_data(data) |
| 192 | + except (json.JSONDecodeError, UnicodeDecodeError, ValueError): |
| 193 | + # Fall back to legacy pickle format with restricted unpickler |
| 194 | + with file_path.open(mode="rb") as f: |
| 195 | + self._cookies = _RestrictedCookieUnpickler(f).load() |
| 196 | + |
| 197 | + def _load_json_data( |
| 198 | + self, data: dict[str, dict[str, dict[str, str | bool]]] |
| 199 | + ) -> defaultdict[tuple[str, str], SimpleCookie]: |
| 200 | + """Load cookies from parsed JSON data.""" |
| 201 | + cookies: defaultdict[tuple[str, str], SimpleCookie] = defaultdict(SimpleCookie) |
| 202 | + for compound_key, cookie_data in data.items(): |
| 203 | + domain, path = compound_key.split("|", 1) |
| 204 | + key = (domain, path) |
| 205 | + for name, morsel_data in cookie_data.items(): |
| 206 | + morsel: Morsel[str] = Morsel() |
| 207 | + morsel_key = morsel_data["key"] |
| 208 | + morsel_value = morsel_data["value"] |
| 209 | + morsel_coded_value = morsel_data["coded_value"] |
| 210 | + # Use __setstate__ to bypass validation, same pattern |
| 211 | + # used in _build_morsel and _cookie_helpers. |
| 212 | + morsel.__setstate__( # type: ignore[attr-defined] |
| 213 | + { |
| 214 | + "key": morsel_key, |
| 215 | + "value": morsel_value, |
| 216 | + "coded_value": morsel_coded_value, |
| 217 | + } |
| 218 | + ) |
| 219 | + # Restore morsel attributes |
| 220 | + for attr in morsel._reserved: # type: ignore[attr-defined] |
| 221 | + if attr in morsel_data and attr not in ( |
| 222 | + "key", |
| 223 | + "value", |
| 224 | + "coded_value", |
| 225 | + ): |
| 226 | + morsel[attr] = morsel_data[attr] |
| 227 | + cookies[key][name] = morsel |
| 228 | + return cookies |
123 | 229 |
|
124 | 230 | def clear(self, predicate: ClearCookiePredicate | None = None) -> None: |
125 | 231 | if predicate is None: |
|
0 commit comments