Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 76 additions & 18 deletions boss_cli/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -299,6 +299,26 @@ def _extract_cookies_from_jar(jar: Any, source: str = "unknown") -> dict[str, st
return None


def _prefer_cookie_candidate(
best_cred: Credential | None,
best_source: str | None,
cookies: dict[str, str] | None,
source: str,
) -> tuple[Credential | None, str | None]:
"""Prefer complete credentials, otherwise keep the richer partial set."""
if not cookies:
return best_cred, best_source

candidate = Credential(cookies=cookies)
if candidate.has_required_cookies:
return candidate, source

if best_cred is None or (not best_cred.has_required_cookies and len(candidate.cookies) > len(best_cred.cookies)):
return candidate, source

return best_cred, best_source


def _extract_in_process(cookie_source: str | None = None) -> tuple[Credential | None, list[str]]:
"""Extract cookies in the main process.

Expand Down Expand Up @@ -328,6 +348,8 @@ def _extract_in_process(cookie_source: str | None = None) -> tuple[Credential |

diagnostics: list[str] = []
attempts: list[str] = []
best_cred: Credential | None = None
best_source: str | None = None

for name in _get_browser_order(cookie_source):
fn = browser_fns.get(name)
Expand All @@ -347,10 +369,15 @@ def _extract_in_process(cookie_source: str | None = None) -> tuple[Credential |
diagnostics.append(f"{name}: {e}")
continue
cookies = _extract_cookies_from_jar(jar, source=f"{name}(in-process)")
if cookies:
cred = Credential(cookies=cookies)
logger.info("Found cookies in %s (in-process, default)", name)
return cred, diagnostics
best_cred, best_source = _prefer_cookie_candidate(
best_cred,
best_source,
cookies,
f"{name}(in-process, default)",
)
if best_cred and best_cred.has_required_cookies:
logger.info("Found complete cookies in %s", best_source)
return best_cred, diagnostics
attempts.append(f"{name}=no-cookies")
continue

Expand All @@ -364,10 +391,15 @@ def _extract_in_process(cookie_source: str | None = None) -> tuple[Credential |
diagnostics.append(f"{name}[{profile_name}]: {e}")
continue
cookies = _extract_cookies_from_jar(jar, source=f"{name}[{profile_name}](in-process)")
if cookies:
cred = Credential(cookies=cookies)
logger.info("Found cookies in %s profile '%s' (in-process)", name, profile_name)
return cred, diagnostics
best_cred, best_source = _prefer_cookie_candidate(
best_cred,
best_source,
cookies,
f"{name}[{profile_name}](in-process)",
)
if best_cred and best_cred.has_required_cookies:
logger.info("Found complete cookies in %s", best_source)
return best_cred, diagnostics
attempts.append(f"{name}[{profile_name}]=no-cookies")
else:
# Non-Chromium (Firefox): use default behavior
Expand All @@ -379,14 +411,26 @@ def _extract_in_process(cookie_source: str | None = None) -> tuple[Credential |
diagnostics.append(f"{name}: {e}")
continue
cookies = _extract_cookies_from_jar(jar, source=f"{name}(in-process)")
if cookies:
cred = Credential(cookies=cookies)
logger.info("Found cookies in %s (in-process)", name)
return cred, diagnostics
best_cred, best_source = _prefer_cookie_candidate(
best_cred,
best_source,
cookies,
f"{name}(in-process)",
)
if best_cred and best_cred.has_required_cookies:
logger.info("Found complete cookies in %s", best_source)
return best_cred, diagnostics
attempts.append(f"{name}=no-cookies")

if attempts:
logger.debug("In-process extraction attempts: %s", ", ".join(attempts))
if best_cred:
logger.info(
"Using partial cookies from %s (in-process, %d cookies)",
best_source,
len(best_cred.cookies),
)
return best_cred, diagnostics
return None, diagnostics


Expand Down Expand Up @@ -458,6 +502,7 @@ def iter_cookie_files(browser_name):
sys.exit(0)

attempts = []
best = None
for name, loader in browsers:
if name in CHROMIUM_BASE_DIRS:
cookie_files = iter_cookie_files(name)
Expand All @@ -466,8 +511,11 @@ def iter_cookie_files(browser_name):
cj = loader(domain_name=".zhipin.com")
cookies = {c.name: c.value for c in cj if "zhipin.com" in (c.domain or "")}
if cookies:
print(json.dumps({"browser": name, "cookies": cookies}))
sys.exit(0)
if REQUIRED_COOKIES.issubset(set(cookies)):
print(json.dumps({"browser": name, "cookies": cookies}))
sys.exit(0)
if best is None or len(cookies) > len(best["cookies"]):
best = {"browser": name, "cookies": cookies}
attempts.append(f"{name}=no-cookies")
except Exception as exc:
attempts.append(f"{name}={type(exc).__name__}: {exc}")
Expand All @@ -478,8 +526,11 @@ def iter_cookie_files(browser_name):
cj = loader(cookie_file=cf, domain_name=".zhipin.com")
cookies = {c.name: c.value for c in cj if "zhipin.com" in (c.domain or "")}
if cookies:
print(json.dumps({"browser": name, "cookies": cookies}))
sys.exit(0)
if REQUIRED_COOKIES.issubset(set(cookies)):
print(json.dumps({"browser": name, "cookies": cookies}))
sys.exit(0)
if best is None or len(cookies) > len(best["cookies"]):
best = {"browser": f"{name}[{pname}]", "cookies": cookies}
attempts.append(f"{name}[{pname}]=no-cookies")
except Exception as exc:
attempts.append(f"{name}[{pname}]={type(exc).__name__}: {exc}")
Expand All @@ -488,12 +539,19 @@ def iter_cookie_files(browser_name):
cj = loader(domain_name=".zhipin.com")
cookies = {c.name: c.value for c in cj if "zhipin.com" in (c.domain or "")}
if cookies:
print(json.dumps({"browser": name, "cookies": cookies}))
sys.exit(0)
if REQUIRED_COOKIES.issubset(set(cookies)):
print(json.dumps({"browser": name, "cookies": cookies}))
sys.exit(0)
if best is None or len(cookies) > len(best["cookies"]):
best = {"browser": name, "cookies": cookies}
attempts.append(f"{name}=no-cookies")
except Exception as exc:
attempts.append(f"{name}={type(exc).__name__}: {exc}")

if best:
print(json.dumps(best))
sys.exit(0)

print(json.dumps({"error": "no_cookies", "attempts": attempts}))
'''

Expand Down
2 changes: 1 addition & 1 deletion boss_cli/commands/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ def _finalize_login(cred, *, from_qr: bool = False) -> None:
# Show diagnostics hint if available
hint = _diagnose_extraction_issues(diagnostics)
if hint:
console.print(f"[yellow]⚠️ Cookie 提取诊断:[/yellow]")
console.print("[yellow]⚠️ Cookie 提取诊断:[/yellow]")
for line in hint.splitlines():
console.print(f" [dim]{line}[/dim]")
console.print()
Expand Down
43 changes: 40 additions & 3 deletions tests/test_auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,6 @@
import os
from unittest.mock import MagicMock, patch

import pytest


# ── Diagnostics ─────────────────────────────────────────────────────

Expand Down Expand Up @@ -187,7 +185,6 @@ def test_returns_none_when_bc3_not_installed(self):
from boss_cli.auth import _extract_in_process

with patch.dict("sys.modules", {"browser_cookie3": None}):
import importlib
cred, diag = _extract_in_process()
# We can't easily remove from sys.modules in a test,
# but at minimum verify the function returns a tuple
Expand Down Expand Up @@ -215,4 +212,44 @@ def __init__(self, domain, name, value):

assert cred is not None
assert cred.cookies["wt2"] == "test_val"

def test_prefers_complete_profile_over_partial_default(self):
from boss_cli.auth import _extract_in_process

class FakeCookie:
def __init__(self, domain, name, value):
self.domain = domain
self.name = name
self.value = value

def chrome_loader(*, cookie_file=None, domain_name=None):
if cookie_file and "Profile 1" in cookie_file:
return [
FakeCookie(".zhipin.com", "__zp_stoken__", "s"),
FakeCookie(".zhipin.com", "wt2", "1"),
FakeCookie(".zhipin.com", "wbg", "2"),
FakeCookie(".zhipin.com", "zp_at", "3"),
FakeCookie(".zhipin.com", "__c", "4"),
]
return [FakeCookie(".zhipin.com", "__a", "partial")]

mock_bc3 = MagicMock()
mock_bc3.chrome.side_effect = chrome_loader
mock_bc3.firefox.return_value = []
mock_bc3.edge.return_value = []
mock_bc3.brave.return_value = []

with patch.dict("sys.modules", {"browser_cookie3": mock_bc3}), \
patch(
"boss_cli.auth._iter_chrome_cookie_files",
return_value=[
"/tmp/Chrome/Default/Cookies",
"/tmp/Chrome/Profile 1/Cookies",
],
):
cred, _ = _extract_in_process("chrome")

assert cred is not None
assert cred.has_required_cookies is True
assert cred.cookies["__zp_stoken__"] == "s"
"""Tests for boss_cli.auth — diagnostics, env fallback, and extraction."""