| 1 | from __future__ import annotations |
| 2 | |
| 3 | import base64 |
| 4 | import json |
| 5 | import sys |
| 6 | import time |
| 7 | import types |
| 8 | from dataclasses import dataclass |
| 9 | from pathlib import Path |
| 10 | from urllib.parse import parse_qs, urlparse |
| 11 | |
| 12 | import pytest |
| 13 | |
| 14 | sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| 15 | |
| 16 | |
| 17 | @dataclass(frozen=True) |
| 18 | class PkcePair: |
| 19 | verifier: str |
| 20 | challenge: str |
| 21 | |
| 22 | |
| 23 | from plugins._oauth.helpers.providers import gemini_api as gemini |
| 24 | from plugins._oauth.helpers.providers.base import GEMINI_API_PROVIDER_ID, ProviderError |
| 25 | from plugins._oauth.helpers.state import pop_attempt, put_attempt |
| 26 | |
| 27 | |
| 28 | def _config() -> dict: |
| 29 | return { |
| 30 | "enabled": True, |
| 31 | "client_id": "client-id.apps.googleusercontent.com", |
| 32 | "client_secret": "client-secret", |
| 33 | "scopes": [ |
| 34 | "openid", |
| 35 | "email", |
| 36 | "profile", |
| 37 | "https://www.googleapis.com/auth/cloud-platform", |
| 38 | "https://www.googleapis.com/auth/generative-language.retriever", |
| 39 | ], |
| 40 | "quota_project_id": "quota-project", |
| 41 | "api_base_url": gemini.GEMINI_OPENAI_API_BASE, |
| 42 | "proxy_base_path": "/oauth/gemini-api", |
| 43 | "callback_path": "/oauth/gemini-api/callback", |
| 44 | } |
| 45 | |
| 46 | |
| 47 | class FakeRequest: |
| 48 | headers = {} |
| 49 | url_root = "http://localhost:50001/" |
| 50 | |
| 51 | |
| 52 | def _fake_codex_helper(): |
| 53 | fake_codex = types.SimpleNamespace() |
| 54 | fake_codex.generate_pkce = lambda: PkcePair("verifier", "challenge") |
| 55 | fake_codex.generate_state = lambda: "state-1" |
| 56 | return fake_codex |
| 57 | |
| 58 | |
| 59 | def _jwt(claims: dict) -> str: |
| 60 | header = _b64({"alg": "none", "typ": "JWT"}) |
| 61 | payload = _b64(claims) |
| 62 | return f"{header}.{payload}.signature" |
| 63 | |
| 64 | |
| 65 | def _b64(payload: dict) -> str: |
| 66 | raw = json.dumps(payload, separators=(",", ":")).encode("utf-8") |
| 67 | return base64.urlsafe_b64encode(raw).decode("utf-8").rstrip("=") |
| 68 | |
| 69 | |
| 70 | def test_start_login_builds_google_pkce_authorize_url_without_network(monkeypatch): |
| 71 | provider = gemini.GeminiApiOAuthProvider() |
| 72 | monkeypatch.setattr(gemini, "_gemini_api_config", _config) |
| 73 | monkeypatch.setattr(gemini, "_codex_helper", _fake_codex_helper) |
| 74 | |
| 75 | result = provider.start_login({}, FakeRequest()) |
| 76 | |
| 77 | try: |
| 78 | assert result.ok is True |
| 79 | assert result.provider_id == GEMINI_API_PROVIDER_ID |
| 80 | assert result.flow == "browser_pkce" |
| 81 | assert result.redirect_uri == "http://localhost:50001/oauth/gemini-api/callback" |
| 82 | |
| 83 | parsed = urlparse(result.auth_url) |
| 84 | query = parse_qs(parsed.query) |
| 85 | assert result.auth_url.startswith("https://accounts.google.com/o/oauth2/v2/auth?") |
| 86 | assert query["response_type"] == ["code"] |
| 87 | assert query["client_id"] == ["client-id.apps.googleusercontent.com"] |
| 88 | assert query["scope"] == [" ".join(_config()["scopes"])] |
| 89 | assert query["code_challenge"] == ["challenge"] |
| 90 | assert query["code_challenge_method"] == ["S256"] |
| 91 | assert query["state"] == ["state-1"] |
| 92 | assert query["access_type"] == ["offline"] |
| 93 | assert query["prompt"] == ["consent"] |
| 94 | assert query["redirect_uri"] == ["http://localhost:50001/oauth/gemini-api/callback"] |
| 95 | finally: |
| 96 | pop_attempt("state-1") |
| 97 | |
| 98 | |
| 99 | def test_start_login_requires_user_supplied_oauth_client(monkeypatch): |
| 100 | provider = gemini.GeminiApiOAuthProvider() |
| 101 | cfg = _config() |
| 102 | cfg["client_id"] = "" |
| 103 | cfg["client_secret"] = "" |
| 104 | monkeypatch.setattr(gemini, "_gemini_api_config", lambda: cfg) |
| 105 | monkeypatch.setattr(gemini, "_codex_helper", _fake_codex_helper) |
| 106 | |
| 107 | result = provider.start_login({}, FakeRequest()) |
| 108 | |
| 109 | assert result.ok is False |
| 110 | assert result.provider_id == GEMINI_API_PROVIDER_ID |
| 111 | assert "OAuth client ID and client secret" in result.error |
| 112 | |
| 113 | |
| 114 | @pytest.mark.parametrize( |
| 115 | ("raw", "expected"), |
| 116 | [ |
| 117 | ( |
| 118 | "http://localhost:50001/oauth/gemini-api/callback?code=abc&state=state-1", |
| 119 | {"code": "abc", "state": "state-1", "error": None, "error_description": None}, |
| 120 | ), |
| 121 | ("?code=abc&state=state-1", {"code": "abc", "state": "state-1", "error": None, "error_description": None}), |
| 122 | ("code=abc&state=state-1", {"code": "abc", "state": "state-1", "error": None, "error_description": None}), |
| 123 | ("abc", {"code": "abc", "state": None, "error": None, "error_description": None}), |
| 124 | ], |
| 125 | ) |
| 126 | def test_parse_manual_callback_accepts_url_query_and_bare_code(raw, expected): |
| 127 | assert gemini.parse_manual_callback(raw) == expected |
| 128 | |
| 129 | |
| 130 | def test_manual_callback_exchanges_code_and_stores_tokens(tmp_path, monkeypatch): |
| 131 | provider = gemini.GeminiApiOAuthProvider() |
| 132 | auth_path = tmp_path / "auth.json" |
| 133 | put_attempt( |
| 134 | "state-1", |
| 135 | "verifier-1", |
| 136 | "http://localhost:50001/oauth/gemini-api/callback", |
| 137 | provider_id=GEMINI_API_PROVIDER_ID, |
| 138 | extra={ |
| 139 | "client_id": "client-id.apps.googleusercontent.com", |
| 140 | "client_secret": "client-secret", |
| 141 | "quota_project_id": "quota-project", |
| 142 | "token_endpoint": gemini.GOOGLE_TOKEN_ENDPOINT, |
| 143 | "api_base_url": gemini.GEMINI_OPENAI_API_BASE, |
| 144 | }, |
| 145 | ) |
| 146 | calls = [] |
| 147 | |
| 148 | monkeypatch.setattr(provider, "auth_path", lambda: auth_path) |
| 149 | |
| 150 | def fake_exchange(token_endpoint, code, redirect_uri, code_verifier, client_id, client_secret): |
| 151 | calls.append((token_endpoint, code, redirect_uri, code_verifier, client_id, client_secret)) |
| 152 | return { |
| 153 | "access_token": "access-token", |
| 154 | "refresh_token": "refresh-token", |
| 155 | "expires_in": 3600, |
| 156 | "id_token": _jwt({"email": "user@example.com"}), |
| 157 | "token_type": "Bearer", |
| 158 | } |
| 159 | |
| 160 | monkeypatch.setattr(provider, "exchange_code", fake_exchange) |
| 161 | |
| 162 | result = provider.manual_callback( |
| 163 | {"callback": "http://localhost:50001/oauth/gemini-api/callback?code=code-1&state=state-1"} |
| 164 | ) |
| 165 | |
| 166 | assert result.ok is True |
| 167 | assert result.completed is True |
| 168 | assert result.account_label == "user@example.com" |
| 169 | assert calls == [ |
| 170 | ( |
| 171 | gemini.GOOGLE_TOKEN_ENDPOINT, |
| 172 | "code-1", |
| 173 | "http://localhost:50001/oauth/gemini-api/callback", |
| 174 | "verifier-1", |
| 175 | "client-id.apps.googleusercontent.com", |
| 176 | "client-secret", |
| 177 | ) |
| 178 | ] |
| 179 | saved = json.loads(auth_path.read_text(encoding="utf-8")) |
| 180 | assert saved["provider"] == GEMINI_API_PROVIDER_ID |
| 181 | assert saved["access"] == "access-token" |
| 182 | assert saved["refresh"] == "refresh-token" |
| 183 | assert saved["quota_project_id"] == "quota-project" |
| 184 | assert saved["client_id"] == "client-id.apps.googleusercontent.com" |
| 185 | |
| 186 | |
| 187 | def test_ensure_fresh_auth_rejects_malicious_stored_token_endpoint(monkeypatch): |
| 188 | provider = gemini.GeminiApiOAuthProvider() |
| 189 | monkeypatch.setattr( |
| 190 | provider, |
| 191 | "read_auth", |
| 192 | lambda: { |
| 193 | "access": "expired-access-token", |
| 194 | "refresh": "refresh-token", |
| 195 | "expires": 1, |
| 196 | "client_id": "client-id.apps.googleusercontent.com", |
| 197 | "client_secret": "client-secret", |
| 198 | "token_endpoint": "https://evil.example/oauth/token", |
| 199 | }, |
| 200 | ) |
| 201 | |
| 202 | with pytest.raises(ProviderError) as exc_info: |
| 203 | provider.ensure_fresh_auth() |
| 204 | |
| 205 | assert exc_info.value.code == "invalid_token_endpoint" |
| 206 | |
| 207 | |
| 208 | def test_models_returns_curated_list_without_network(monkeypatch): |
| 209 | provider = gemini.GeminiApiOAuthProvider() |
| 210 | monkeypatch.setattr(provider, "read_auth", lambda: {}) |
| 211 | |
| 212 | models = provider.models() |
| 213 | |
| 214 | assert models[:3] == [ |
| 215 | "gemini-3.5-flash", |
| 216 | "gemini-3.1-pro-preview", |
| 217 | "gemini-3-flash-preview", |
| 218 | ] |
| 219 | |
| 220 | |
| 221 | def test_models_does_not_send_bearer_token_to_malicious_base_url(monkeypatch): |
| 222 | calls = [] |
| 223 | |
| 224 | class FakeResponse: |
| 225 | ok = True |
| 226 | |
| 227 | def json(self): |
| 228 | return {"data": [{"id": "models/gemini-3.5-flash"}]} |
| 229 | |
| 230 | class FakeRequests: |
| 231 | @staticmethod |
| 232 | def get(url, headers, timeout): |
| 233 | calls.append((url, headers, timeout)) |
| 234 | return FakeResponse() |
| 235 | |
| 236 | monkeypatch.setitem(sys.modules, "requests", FakeRequests) |
| 237 | |
| 238 | provider = gemini.GeminiApiOAuthProvider() |
| 239 | monkeypatch.setattr( |
| 240 | provider, |
| 241 | "read_auth", |
| 242 | lambda: { |
| 243 | "access": "access-token", |
| 244 | "refresh": "refresh-token", |
| 245 | "expires": int(time.time() * 1000) + 3_600_000, |
| 246 | "base_url": "https://evil.example/v1", |
| 247 | "quota_project_id": "quota-project", |
| 248 | }, |
| 249 | ) |
| 250 | |
| 251 | assert provider.models() == ["gemini-3.5-flash"] |
| 252 | assert calls == [ |
| 253 | ( |
| 254 | "https://generativelanguage.googleapis.com/v1beta/openai/models", |
| 255 | { |
| 256 | "Accept": "application/json", |
| 257 | "Authorization": "Bearer access-token", |
| 258 | "x-goog-user-project": "quota-project", |
| 259 | }, |
| 260 | 30, |
| 261 | ) |
| 262 | ] |