返回 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_non_default_profile_with_cookies(self, mock_firefox_env):
285 """Falls back to non-default profile when default has no X cookies."""
286 profiles_dir = mock_firefox_env(
287 profiles={
288 "aaa111.default": [
289 (".example.com", "session", "sess_other"),
290 ],
291 "bbb222.release": [
292 (".x.com", "auth_token", "tok_nondefault"),
293 (".x.com", "ct0", "ct0_nondefault"),
294 ],
295 },
296 profiles_ini=textwrap.dedent("""\
297 [General]
298 StartWithLastProfile=1
299
300 [Profile0]
301 Name=default
302 IsRelative=1
303 Path=aaa111.default
304 Default=1
305
306 [Profile1]
307 Name=release
308 IsRelative=1
309 Path=bbb222.release
310 """),
311 )
312
313 with patch(
314 "lib.cookie_extract._get_firefox_profiles_dir",
315 return_value=profiles_dir,
316 ):
317 result = extract_firefox_cookies(".x.com", ["auth_token", "ct0"])
318
319 assert result is not None
320 assert result["auth_token"] == "tok_nondefault"
321 assert result["ct0"] == "ct0_nondefault"
322
323 def test_multiple_profiles_none_have_cookies(self, mock_firefox_env):
324 """Returns None when no profile has matching cookies."""
325 profiles_dir = mock_firefox_env(
326 profiles={
327 "aaa111.default": [
328 (".example.com", "session", "sess_a"),
329 ],
330 "bbb222.release": [
331 (".other.com", "other", "val_b"),
332 ],
333 },
334 profiles_ini=textwrap.dedent("""\
335 [General]
336 StartWithLastProfile=1
337
338 [Profile0]
339 Name=default
340 IsRelative=1
341 Path=aaa111.default
342 Default=1
343
344 [Profile1]
345 Name=release
346 IsRelative=1
347 Path=bbb222.release
348 """),
349 )
350
351 with patch(
352 "lib.cookie_extract._get_firefox_profiles_dir",
353 return_value=profiles_dir,
354 ), patch(
355 "lib.cookie_extract._is_wsl",
356 return_value=False,
357 ):
358 result = extract_firefox_cookies(".x.com", ["auth_token", "ct0"])
359
360 assert result is None
361
362
363 class TestExtractCookiesAuto:
364 """Tests for extract_cookies with browser='auto'."""
365
366 def test_auto_macos_tries_chrome_then_firefox(self, mock_firefox_env):
367 """On macOS, auto tries the Chromium family first, falls back to Firefox."""
368 profiles_dir = mock_firefox_env()
369
370 # Mock every Chromium-family extractor to None so the test is hermetic
371 # regardless of which browsers are actually installed/logged-in on the
372 # machine running it (auto tries these before Firefox).
373 with (
374 patch("lib.cookie_extract.platform.system", return_value="Darwin"),
375 patch("lib.cookie_extract.extract_chrome_cookies", return_value=None),
376 patch("lib.cookie_extract.extract_brave_cookies", return_value=None),
377 patch("lib.cookie_extract.extract_edge_cookies", return_value=None),
378 patch("lib.cookie_extract.extract_vivaldi_cookies", return_value=None),
379 patch("lib.cookie_extract.extract_opera_cookies", return_value=None),
380 patch("lib.cookie_extract.extract_arc_cookies", return_value=None),
381 patch("lib.cookie_extract.extract_chromium_cookies", return_value=None),
382 patch("lib.cookie_extract.extract_safari_cookies", return_value=None),
383 patch(
384 "lib.cookie_extract._get_firefox_profiles_dir",
385 return_value=profiles_dir,
386 ),
387 ):
388 result = extract_cookies("auto", ".x.com", ["auth_token", "ct0"])
389
390 # All Chromium browsers and Safari return None, Firefox succeeds
391 assert result is not None
392 assert result["auth_token"] == "tok_abc123"
393 assert result["ct0"] == "ct0_xyz789"
394
395 def test_auto_linux_tries_firefox_only(self, mock_firefox_env):
396 """On Linux, auto only tries Firefox."""
397 profiles_dir = mock_firefox_env()
398
399 with (
400 patch("lib.cookie_extract.platform.system", return_value="Linux"),
401 patch(
402 "lib.cookie_extract._get_firefox_profiles_dir",
403 return_value=profiles_dir,
404 ),
405 ):
406 result = extract_cookies("auto", ".x.com", ["auth_token", "ct0"])
407
408 assert result is not None
409 assert result["auth_token"] == "tok_abc123"
410
411 def test_explicit_firefox(self, mock_firefox_env):
412 """Explicit browser='firefox' goes directly to Firefox."""
413 profiles_dir = mock_firefox_env()
414
415 with patch(
416 "lib.cookie_extract._get_firefox_profiles_dir",
417 return_value=profiles_dir,
418 ):
419 result = extract_cookies("firefox", ".x.com", ["auth_token"])
420
421 assert result is not None
422 assert result["auth_token"] == "tok_abc123"
423
424 def test_unknown_browser_returns_none(self):
425 """Unknown browser name returns None."""
426 result = extract_cookies("netscape", ".x.com", ["auth_token"])
427 assert result is None
428
429 def test_chrome_delegates_to_chrome_module(self):
430 """Chrome extraction delegates to chrome_cookies module."""
431 with patch(
432 "lib.cookie_extract.extract_chrome_cookies",
433 return_value={"auth_token": "chrome_tok"},
434 ):
435 result = extract_cookies("chrome", ".x.com", ["auth_token"])
436 assert result == {"auth_token": "chrome_tok"}
437
438 def test_safari_delegates_to_safari_module(self):
439 """Safari extraction delegates to safari_cookies module."""
440 with patch(
441 "lib.cookie_extract.extract_safari_cookies",
442 return_value={"auth_token": "safari_tok"},
443 ):
444 result = extract_cookies("safari", ".x.com", ["auth_token"])
445 assert result == {"auth_token": "safari_tok"}
446
446 lines PYTHON