返回 last30days-skill
test_cookie_extract.py
根目录 / tests / test_cookie_extract.py
1 """Tests for browser cookie extraction module."""
2
3 import configparser
4 import os
5 import sqlite3
6 import textwrap
7 from pathlib import Path
8 from typing import Dict, List, Optional, Tuple
9 from unittest.mock import patch
10
11 import pytest
12
13 from lib import cookie_extract
14 from lib.cookie_extract import (
15 extract_cookies,
16 extract_firefox_cookies,
17 _query_cookies_db,
18 _find_default_profile,
19 _get_firefox_profiles_dir,
20 )
21
22 @pytest.fixture
23 def mock_firefox_env(tmp_path):
24 """Create a mock Firefox profiles directory with cookies.sqlite.
25
26 Returns (profiles_dir, profile_dir) for patching.
27 """
28
29 def _make(
30 *,
31 profiles_ini=None, # type: Optional[str]
32 profiles=None, # type: Optional[Dict[str, List[Tuple[str, str, str]]]]
33 default_profile="abc123.default-release", # type: str
34 ):
35 profiles_dir = tmp_path / "Firefox"
36 profiles_dir.mkdir(parents=True, exist_ok=True)
37
38 # Default: one profile with X cookies
39 if profiles is None:
40 profiles = {
41 default_profile: [
42 (".x.com", "auth_token", "tok_abc123"),
43 (".x.com", "ct0", "ct0_xyz789"),
44 (".example.com", "session", "sess_other"),
45 ],
46 }
47
48 # Create profile directories with cookies databases
49 for profile_name, cookies in profiles.items():
50 profile_dir = profiles_dir / profile_name
51 profile_dir.mkdir(parents=True, exist_ok=True)
52 db_path = profile_dir / "cookies.sqlite"
53 conn = sqlite3.connect(str(db_path))
54 conn.execute(
55 "CREATE TABLE moz_cookies ("
56 " id INTEGER PRIMARY KEY,"
57 " name TEXT NOT NULL,"
58 " value TEXT NOT NULL,"
59 " host TEXT NOT NULL,"
60 " path TEXT DEFAULT '/',"
61 " expiry INTEGER DEFAULT 0,"
62 " isSecure INTEGER DEFAULT 1,"
63 " isHttpOnly INTEGER DEFAULT 1,"
64 " sameSite INTEGER DEFAULT 0,"
65 " schemeMap INTEGER DEFAULT 0"
66 ")"
67 )
68 for host, name, value in cookies:
69 conn.execute(
70 "INSERT INTO moz_cookies (name, value, host) VALUES (?, ?, ?)",
71 (name, value, host),
72 )
73 conn.commit()
74 conn.close()
75
76 # Write profiles.ini
77 if profiles_ini is None:
78 profiles_ini = textwrap.dedent(f"""\
79 [General]
80 StartWithLastProfile=1
81
82 [Profile0]
83 Name=default-release
84 IsRelative=1
85 Path={default_profile}
86 Default=1
87 """)
88
89 (profiles_dir / "profiles.ini").write_text(profiles_ini)
90
91 return profiles_dir
92
93 return _make
94
95
96 class TestExtractFirefoxCookies:
97 """Tests for extract_firefox_cookies."""
98
99 @pytest.mark.skipif(os.name == "nt", reason="POSIX permission bits are not reliable on Windows")
100 def test_temp_cookie_db_copy_is_owner_only(self, tmp_path):
101 """Copied cookie DB temp files are chmodded owner-only before read."""
102 db_path = tmp_path / "cookies.sqlite"
103 conn = sqlite3.connect(str(db_path))
104 conn.execute(
105 "CREATE TABLE moz_cookies (name TEXT NOT NULL, value TEXT NOT NULL, host TEXT NOT NULL)"
106 )
107 conn.execute(
108 "INSERT INTO moz_cookies (name, value, host) VALUES (?, ?, ?)",
109 ("auth_token", "tok_abc123", ".x.com"),
110 )
111 conn.commit()
112 conn.close()
113 os.chmod(db_path, 0o644)
114
115 real_connect = sqlite3.connect
116
117 def assert_temp_copy_locked(path, *args, **kwargs):
118 if Path(str(path)) != db_path:
119 assert Path(str(path)).stat().st_mode & 0o777 == 0o600
120 return real_connect(path, *args, **kwargs)
121
122 with patch("lib.cookie_extract.sqlite3.connect", side_effect=assert_temp_copy_locked):
123 result = _query_cookies_db(db_path, ".x.com", ["auth_token"])
124
125 assert result == {"auth_token": "tok_abc123"}
126
127 @pytest.mark.skipif(os.name == "nt", reason="POSIX permission model does not apply on Windows; mkstemp is 0o666 there")
128 def test_temp_cookie_copy_never_world_readable(self, tmp_path):
129 """The temp copy must be private the instant it exists, not only after
130 the lock chmod. Regression for the TOCTOU window where copy2 widened the
131 0600 mkstemp file to the source's 0644 before _lock_temp_cookie_copy ran.
132 """
133 db_path = tmp_path / "cookies.sqlite"
134 conn = sqlite3.connect(str(db_path))
135 conn.execute(
136 "CREATE TABLE moz_cookies (name TEXT NOT NULL, value TEXT NOT NULL, host TEXT NOT NULL)"
137 )
138 conn.execute(
139 "INSERT INTO moz_cookies (name, value, host) VALUES (?, ?, ?)",
140 ("auth_token", "tok_abc123", ".x.com"),
141 )
142 conn.commit()
143 conn.close()
144 os.chmod(db_path, 0o644) # loose source perms, as Firefox ships them
145
146 observed = {}
147 real_lock = cookie_extract._lock_temp_cookie_copy
148
149 def spy(path):
150 # Mode of the copy as it exists right after copyfile, before chmod.
151 observed["mode_after_copy"] = os.stat(path).st_mode & 0o777
152 return real_lock(path)
153
154 with patch.object(cookie_extract, "_lock_temp_cookie_copy", side_effect=spy):
155 _query_cookies_db(db_path, ".x.com", ["auth_token"])
156
157 assert observed["mode_after_copy"] == 0o600
158
159 def test_valid_cookies_extracted(self, mock_firefox_env):
160 """Cookies for the target domain are returned correctly."""
161 profiles_dir = mock_firefox_env()
162
163 with patch(
164 "lib.cookie_extract._get_firefox_profiles_dir",
165 return_value=profiles_dir,
166 ):
167 result = extract_firefox_cookies(".x.com", ["auth_token", "ct0"])
168
169 assert result is not None
170 assert result["auth_token"] == "tok_abc123"
171 assert result["ct0"] == "ct0_xyz789"
172 assert "session" not in result # different domain cookie not included
173
174 def test_multiple_profiles_selects_default(self, mock_firefox_env):
175 """When multiple profiles exist, the one with Default=1 is used."""
176 profiles_dir = mock_firefox_env(
177 profiles={
178 "aaa111.other": [
179 (".x.com", "auth_token", "wrong_token"),
180 ],
181 "bbb222.default-release": [
182 (".x.com", "auth_token", "correct_token"),
183 (".x.com", "ct0", "correct_ct0"),
184 ],
185 },
186 profiles_ini=textwrap.dedent("""\
187 [General]
188 StartWithLastProfile=1
189
190 [Profile0]
191 Name=other
192 IsRelative=1
193 Path=aaa111.other
194
195 [Profile1]
196 Name=default-release
197 IsRelative=1
198 Path=bbb222.default-release
199 Default=1
200 """),
201 )
202
203 with patch(
204 "lib.cookie_extract._get_firefox_profiles_dir",
205 return_value=profiles_dir,
206 ):
207 result = extract_firefox_cookies(".x.com", ["auth_token", "ct0"])
208
209 assert result is not None
210 assert result["auth_token"] == "correct_token"
211 assert result["ct0"] == "correct_ct0"
212
213 def test_firefox_not_installed(self):
214 """Returns None when Firefox profiles directory doesn't exist."""
215 with patch(
216 "lib.cookie_extract._get_firefox_profiles_dir",
217 return_value=None,
218 ), patch(
219 "lib.cookie_extract._is_wsl",
220 return_value=False,
221 ):
222 result = extract_firefox_cookies(".x.com", ["auth_token"])
223
224 assert result is None
225
226 def test_cookies_sqlite_empty(self, mock_firefox_env):
227 """Returns None when cookies.sqlite has no rows."""
228 profiles_dir = mock_firefox_env(
229 profiles={"abc123.default-release": []}, # no cookies
230 )
231
232 with patch(
233 "lib.cookie_extract._get_firefox_profiles_dir",
234 return_value=profiles_dir,
235 ), patch(
236 "lib.cookie_extract._is_wsl",
237 return_value=False,
238 ):
239 result = extract_firefox_cookies(".x.com", ["auth_token", "ct0"])
240
241 assert result is None
242
243 def test_domain_has_no_cookies(self, mock_firefox_env):
244 """Returns None when cookies exist but not for the target domain."""
245 profiles_dir = mock_firefox_env(
246 profiles={
247 "abc123.default-release": [
248 (".example.com", "session", "sess_123"),
249 ],
250 },
251 )
252
253 with patch(
254 "lib.cookie_extract._get_firefox_profiles_dir",
255 return_value=profiles_dir,
256 ), patch(
257 "lib.cookie_extract._is_wsl",
258 return_value=False,
259 ):
260 result = extract_firefox_cookies(".x.com", ["auth_token", "ct0"])
261
262 assert result is None
263
264 def test_malformed_profiles_ini_falls_back(self, mock_firefox_env):
265 """Falls back to first profile on disk when profiles.ini is garbage."""
266 profiles_dir = mock_firefox_env(
267 profiles={
268 "zzz999.fallback": [
269 (".x.com", "auth_token", "fallback_token"),
270 ],
271 },
272 profiles_ini="this is not valid ini content\n[[[broken",
273 )
274
275 with patch(
276 "lib.cookie_extract._get_firefox_profiles_dir",
277 return_value=profiles_dir,
278 ):
279 result = extract_firefox_cookies(".x.com", ["auth_token"])
280
281 assert result is not None
282 assert result["auth_token"] == "fallback_token"
283
284 def test_utf16_profiles_ini_uses_install_default(self, mock_firefox_env):
285 """Firefox on Windows writes UTF-16 LE profiles.ini (#1067)."""
286 default_profile = "en3ndvop.default-release"
287 decoy = "aaa111.decoy"
288 profiles_dir = mock_firefox_env(
289 default_profile=default_profile,
290 profiles={
291 decoy: [
292 (".x.com", "auth_token", "decoy_token"),
293 ],
294 default_profile: [
295 (".x.com", "auth_token", "utf16_token"),
296 (".x.com", "ct0", "utf16_ct0"),
297 ],
298 },
299 )
300 ini = textwrap.dedent(f"""\
301 [Install308046B0AF4A39CB]
302 Default={default_profile}
303 Locked=1
304 """)
305 (profiles_dir / "profiles.ini").write_bytes(ini.encode("utf-16"))
306
307 with patch(
308 "lib.cookie_extract._get_firefox_profiles_dir",
309 return_value=profiles_dir,
310 ):
311 result = extract_firefox_cookies(".x.com", ["auth_token", "ct0"])
312
313 assert result is not None
314 assert result["auth_token"] == "utf16_token"
315 assert result["ct0"] == "utf16_ct0"
316
317 def test_non_default_profile_with_cookies(self, mock_firefox_env):
318 """Falls back to non-default profile when default has no X cookies."""
319 profiles_dir = mock_firefox_env(
320 profiles={
321 "aaa111.default": [
322 (".example.com", "session", "sess_other"),
323 ],
324 "bbb222.release": [
325 (".x.com", "auth_token", "tok_nondefault"),
326 (".x.com", "ct0", "ct0_nondefault"),
327 ],
328 },
329 profiles_ini=textwrap.dedent("""\
330 [General]
331 StartWithLastProfile=1
332
333 [Profile0]
334 Name=default
335 IsRelative=1
336 Path=aaa111.default
337 Default=1
338
339 [Profile1]
340 Name=release
341 IsRelative=1
342 Path=bbb222.release
343 """),
344 )
345
346 with patch(
347 "lib.cookie_extract._get_firefox_profiles_dir",
348 return_value=profiles_dir,
349 ):
350 result = extract_firefox_cookies(".x.com", ["auth_token", "ct0"])
351
352 assert result is not None
353 assert result["auth_token"] == "tok_nondefault"
354 assert result["ct0"] == "ct0_nondefault"
355
356 def test_multiple_profiles_none_have_cookies(self, mock_firefox_env):
357 """Returns None when no profile has matching cookies."""
358 profiles_dir = mock_firefox_env(
359 profiles={
360 "aaa111.default": [
361 (".example.com", "session", "sess_a"),
362 ],
363 "bbb222.release": [
364 (".other.com", "other", "val_b"),
365 ],
366 },
367 profiles_ini=textwrap.dedent("""\
368 [General]
369 StartWithLastProfile=1
370
371 [Profile0]
372 Name=default
373 IsRelative=1
374 Path=aaa111.default
375 Default=1
376
377 [Profile1]
378 Name=release
379 IsRelative=1
380 Path=bbb222.release
381 """),
382 )
383
384 with patch(
385 "lib.cookie_extract._get_firefox_profiles_dir",
386 return_value=profiles_dir,
387 ), patch(
388 "lib.cookie_extract._is_wsl",
389 return_value=False,
390 ):
391 result = extract_firefox_cookies(".x.com", ["auth_token", "ct0"])
392
393 assert result is None
394
395
396 class TestExtractCookiesAuto:
397 """Tests for extract_cookies with browser='auto'."""
398
399 def test_auto_macos_tries_chrome_then_firefox(self, mock_firefox_env):
400 """On macOS, auto tries the Chromium family first, falls back to Firefox."""
401 profiles_dir = mock_firefox_env()
402
403 # Mock every Chromium-family extractor to None so the test is hermetic
404 # regardless of which browsers are actually installed/logged-in on the
405 # machine running it (auto tries these before Firefox).
406 with (
407 patch("lib.cookie_extract.platform.system", return_value="Darwin"),
408 patch("lib.cookie_extract.extract_chrome_cookies", return_value=None),
409 patch("lib.cookie_extract.extract_brave_cookies", return_value=None),
410 patch("lib.cookie_extract.extract_edge_cookies", return_value=None),
411 patch("lib.cookie_extract.extract_vivaldi_cookies", return_value=None),
412 patch("lib.cookie_extract.extract_opera_cookies", return_value=None),
413 patch("lib.cookie_extract.extract_arc_cookies", return_value=None),
414 patch("lib.cookie_extract.extract_chromium_cookies", return_value=None),
415 patch("lib.cookie_extract.extract_safari_cookies", return_value=None),
416 patch(
417 "lib.cookie_extract._get_firefox_profiles_dir",
418 return_value=profiles_dir,
419 ),
420 ):
421 result = extract_cookies("auto", ".x.com", ["auth_token", "ct0"])
422
423 # All Chromium browsers and Safari return None, Firefox succeeds
424 assert result is not None
425 assert result["auth_token"] == "tok_abc123"
426 assert result["ct0"] == "ct0_xyz789"
427
428 def test_auto_linux_tries_firefox_only(self, mock_firefox_env):
429 """On Linux, auto only tries Firefox."""
430 profiles_dir = mock_firefox_env()
431
432 with (
433 patch("lib.cookie_extract.platform.system", return_value="Linux"),
434 patch(
435 "lib.cookie_extract._get_firefox_profiles_dir",
436 return_value=profiles_dir,
437 ),
438 ):
439 result = extract_cookies("auto", ".x.com", ["auth_token", "ct0"])
440
441 assert result is not None
442 assert result["auth_token"] == "tok_abc123"
443
444 def test_explicit_firefox(self, mock_firefox_env):
445 """Explicit browser='firefox' goes directly to Firefox."""
446 profiles_dir = mock_firefox_env()
447
448 with patch(
449 "lib.cookie_extract._get_firefox_profiles_dir",
450 return_value=profiles_dir,
451 ):
452 result = extract_cookies("firefox", ".x.com", ["auth_token"])
453
454 assert result is not None
455 assert result["auth_token"] == "tok_abc123"
456
457 def test_unknown_browser_returns_none(self):
458 """Unknown browser name returns None."""
459 result = extract_cookies("netscape", ".x.com", ["auth_token"])
460 assert result is None
461
462 def test_chrome_delegates_to_chrome_module(self):
463 """Chrome extraction delegates to chrome_cookies module."""
464 with patch(
465 "lib.cookie_extract.extract_chrome_cookies",
466 return_value={"auth_token": "chrome_tok"},
467 ):
468 result = extract_cookies("chrome", ".x.com", ["auth_token"])
469 assert result == {"auth_token": "chrome_tok"}
470
471 def test_safari_delegates_to_safari_module(self):
472 """Safari extraction delegates to safari_cookies module."""
473 with patch(
474 "lib.cookie_extract.extract_safari_cookies",
475 return_value={"auth_token": "safari_tok"},
476 ):
477 result = extract_cookies("safari", ".x.com", ["auth_token"])
478 assert result == {"auth_token": "safari_tok"}
479
479 lines PYTHON