| from __future__ import annotations |
|
|
| import base64 |
| import json |
| import unittest |
| from typing import Any |
| from unittest.mock import Mock, patch |
|
|
| import services.openai_backend_api as backend_module |
| from services.account_service import ( |
| AccountService, |
| DEFAULT_CHATGPT_USER_AGENT, |
| build_chatgpt_client_fingerprint, |
| sanitize_registered_proxy, |
| ) |
| from services.openai_backend_api import OpenAIBackendAPI |
|
|
|
|
| class MemoryStorage: |
| def __init__(self, accounts: list[dict[str, Any]] | None = None) -> None: |
| self.accounts = [dict(item) for item in (accounts or [])] |
|
|
| def load_accounts(self) -> list[dict[str, Any]]: |
| return [dict(item) for item in self.accounts] |
|
|
| def save_accounts(self, accounts: list[dict[str, Any]]) -> None: |
| self.accounts = [dict(item) for item in accounts] |
|
|
|
|
| def make_jwt(payload: dict[str, Any]) -> str: |
| def encode(value: dict[str, Any]) -> str: |
| raw = json.dumps(value, separators=(",", ":")).encode("utf-8") |
| return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=") |
|
|
| return f'{encode({"alg": "none"})}.{encode(payload)}.sig' |
|
|
|
|
| class FakeCookies: |
| def __init__(self) -> None: |
| self.values: list[tuple[str, str, str]] = [] |
|
|
| def set(self, name: str, value: str, *, domain: str) -> None: |
| self.values.append((name, value, domain)) |
|
|
|
|
| class FakeSession: |
| def __init__(self) -> None: |
| self.headers: dict[str, str] = {} |
| self.cookies = FakeCookies() |
| self.closed = False |
|
|
| def close(self) -> None: |
| self.closed = True |
|
|
|
|
| class AccountClientContextTests(unittest.TestCase): |
| def test_account_context_is_migrated_and_stable_across_reload(self) -> None: |
| access_token = make_jwt( |
| {"https://api.openai.com/auth": {"chatgpt_account_id": "acct_stable"}} |
| ) |
| storage = MemoryStorage([{"access_token": access_token}]) |
|
|
| first_service = AccountService(storage) |
| first = first_service.get_account(access_token) |
| second_service = AccountService(storage) |
| second = second_service.get_account(access_token) |
|
|
| self.assertEqual(first["account_id"], "acct_stable") |
| self.assertTrue(first["account_key"]) |
| self.assertEqual(first["account_key"], second["account_key"]) |
| self.assertEqual(first["fp"], second["fp"]) |
| self.assertEqual(first["fp"]["impersonate"], "chrome145") |
| self.assertIn("Chrome/145.", first["fp"]["user-agent"]) |
|
|
| rotated = build_chatgpt_client_fingerprint( |
| "rotated-access-token", |
| account_key=first["account_key"], |
| ) |
| self.assertEqual(rotated["oai-device-id"], first["fp"]["oai-device-id"]) |
| self.assertEqual(rotated["oai-session-id"], first["fp"]["oai-session-id"]) |
|
|
| def test_legacy_mismatched_edge_profile_is_replaced_by_consistent_profile(self) -> None: |
| fp = build_chatgpt_client_fingerprint( |
| "token", |
| { |
| "user-agent": ( |
| "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 " |
| "Chrome/143.0.0.0 Safari/537.36 Edg/143.0.0.0" |
| ), |
| "impersonate": "edge101", |
| "oai-device-id": "device-stable", |
| "oai-session-id": "session-stable", |
| }, |
| ) |
|
|
| self.assertEqual(fp["impersonate"], "chrome145") |
| self.assertIn("Chrome/145.", fp["user-agent"]) |
| self.assertEqual(fp["oai-device-id"], "device-stable") |
| self.assertEqual(fp["oai-session-id"], "session-stable") |
|
|
| def test_backend_reuses_context_proxy_and_account_header_for_rotated_token(self) -> None: |
| storage = MemoryStorage() |
| service = AccountService(storage) |
| fp = build_chatgpt_client_fingerprint( |
| raw_fp={"user-agent": DEFAULT_CHATGPT_USER_AGENT, "impersonate": "chrome145"}, |
| device_id="login-device", |
| session_id="login-session", |
| ) |
| session = FakeSession() |
| context = { |
| "account_id": "acct_context", |
| "fp": fp, |
| "registered_proxy": "http://127.0.0.1:7890", |
| } |
|
|
| with ( |
| patch.object(backend_module, "account_service", service), |
| patch.object(backend_module.requests, "Session", return_value=session) as factory, |
| ): |
| backend = OpenAIBackendAPI("rotated-token", account_context=context) |
|
|
| self.assertEqual( |
| factory.call_args.kwargs, |
| {"impersonate": "chrome110", "verify": True, "proxy": "http://127.0.0.1:7890"}, |
| ) |
| self.assertEqual(backend.device_id, "login-device") |
| self.assertEqual(backend.session_id, "login-session") |
| self.assertEqual(session.headers["ChatGPT-Account-Id"], "acct_context") |
| self.assertEqual(session.headers["OAI-Client-Version"], "prod-a194cd50d4416d3c0b47c740f206b12ce60f5887") |
| self.assertEqual(session.headers["OAI-Client-Build-Number"], "6708908") |
| self.assertEqual(session.headers["Sec-Ch-Ua-Platform-Version"], '"19.0.0"') |
| self.assertEqual(backend._headers("/backend-api/me")["ChatGPT-Account-Id"], "acct_context") |
| backend.close() |
| self.assertTrue(session.closed) |
|
|
| def test_discovered_account_id_is_applied_before_followup_requests_and_saved(self) -> None: |
| storage = MemoryStorage([{"access_token": "opaque-token"}]) |
| service = AccountService(storage) |
| session = FakeSession() |
|
|
| with ( |
| patch.object(backend_module, "account_service", service), |
| patch.object(backend_module.requests, "Session", return_value=session), |
| ): |
| backend = OpenAIBackendAPI("opaque-token") |
| backend._get_default_account = Mock(return_value={"account_id": "acct_discovered", "plan_type": "free"}) |
|
|
| def get_me() -> dict[str, str]: |
| self.assertEqual(session.headers["ChatGPT-Account-Id"], "acct_discovered") |
| return {"email": "user@example.com", "id": "user-1"} |
|
|
| backend._get_me = Mock(side_effect=get_me) |
| backend._get_conversation_init = Mock(return_value={"limits_progress": []}) |
| result = backend.get_user_info() |
|
|
| self.assertEqual(result["account_id"], "acct_discovered") |
| self.assertEqual(service.get_account("opaque-token")["account_id"], "acct_discovered") |
|
|
| def test_authenticated_proxy_is_not_persisted_in_account_data(self) -> None: |
| self.assertEqual(sanitize_registered_proxy("http://user:secret@proxy.example:8080"), "") |
| self.assertEqual( |
| sanitize_registered_proxy("http://127.0.0.1:7890/path?token=secret"), |
| "http://127.0.0.1:7890", |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
