@samitouri / QOSami_Agents_Zero-2
Code Commits Issues Pulls Sessions CI/CD
main
main
py 300 lines 9.12 KB
Raw
1 from __future__ import annotations
2
3 import json
4 import sys
5 import types
6 import time
7 from dataclasses import dataclass
8 from pathlib import Path
9 from urllib.parse import parse_qs, urlparse
10
11 import pytest
12
13 sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
14
15 @dataclass(frozen=True)
16 class PkcePair:
17 verifier: str
18 challenge: str
19
20
21 from plugins._oauth.helpers.providers.base import ProviderError, XAI_GROK_PROVIDER_ID
22 from plugins._oauth.helpers.providers import xai_grok as xai
23 from plugins._oauth.helpers.state import pop_attempt, put_attempt
24
25
26 def _discovery() -> dict[str, str]:
27 return {
28 "authorization_endpoint": "https://auth.x.ai/oauth/authorize",
29 "token_endpoint": "https://auth.x.ai/oauth/token",
30 }
31
32
33 def test_start_login_builds_xai_pkce_authorize_url_without_network(monkeypatch):
34 provider = xai.XaiGrokOAuthProvider()
35 fake_codex = types.ModuleType("plugins._oauth.helpers.codex")
36 fake_codex.PkcePair = PkcePair
37 fake_codex.generate_pkce = lambda: PkcePair("verifier", "challenge")
38 fake_codex.generate_state = lambda: "state-1"
39 monkeypatch.setitem(sys.modules, "plugins._oauth.helpers.codex", fake_codex)
40 monkeypatch.setattr(provider, "discovery", _discovery)
41
42 result = provider.start_login({})
43
44 try:
45 assert result.ok is True
46 assert result.provider_id == XAI_GROK_PROVIDER_ID
47 assert result.flow == "browser_pkce"
48 assert result.redirect_uri == "http://127.0.0.1:56121/callback"
49
50 parsed = urlparse(result.auth_url)
51 query = parse_qs(parsed.query)
52 assert result.auth_url.startswith("https://auth.x.ai/oauth/authorize?")
53 assert query["response_type"] == ["code"]
54 assert query["client_id"] == [xai.XAI_CLIENT_ID]
55 assert query["scope"] == [xai.XAI_SCOPE]
56 assert query["code_challenge"] == ["challenge"]
57 assert query["code_challenge_method"] == ["S256"]
58 assert query["state"] == ["state-1"]
59 assert query["plan"] == ["generic"]
60 assert query["referrer"] == ["agent-zero"]
61 assert query["redirect_uri"] == ["http://127.0.0.1:56121/callback"]
62 assert query["nonce"]
63 finally:
64 pop_attempt("state-1")
65
66
67 @pytest.mark.parametrize(
68 ("raw", "expected"),
69 [
70 (
71 "http://127.0.0.1:56121/callback?code=abc&state=state-1",
72 {"code": "abc", "state": "state-1", "error": None, "error_description": None},
73 ),
74 ("?code=abc&state=state-1", {"code": "abc", "state": "state-1", "error": None, "error_description": None}),
75 ("code=abc&state=state-1", {"code": "abc", "state": "state-1", "error": None, "error_description": None}),
76 ("abc", {"code": "abc", "state": None, "error": None, "error_description": None}),
77 ],
78 )
79 def test_parse_manual_callback_accepts_url_query_and_bare_code(raw, expected):
80 assert xai.parse_manual_callback(raw) == expected
81
82
83 def test_manual_callback_rejects_state_mismatch(monkeypatch):
84 provider = xai.XaiGrokOAuthProvider()
85 put_attempt(
86 "state-good",
87 "verifier",
88 xai.XAI_REDIRECT_URI,
89 provider_id=XAI_GROK_PROVIDER_ID,
90 extra={"token_endpoint": "https://auth.x.ai/oauth/token"},
91 )
92 monkeypatch.setattr(provider, "discovery", _discovery)
93
94 try:
95 result = provider.manual_callback({"callback": "code=abc&state=state-bad"})
96
97 assert result.ok is False
98 assert "state mismatch" in result.error.lower()
99 finally:
100 pop_attempt("state-good")
101 pop_attempt("state-bad")
102
103
104 def test_manual_callback_exchanges_code_and_stores_tokens(tmp_path, monkeypatch):
105 provider = xai.XaiGrokOAuthProvider()
106 auth_path = tmp_path / "auth.json"
107 put_attempt(
108 "state-1",
109 "verifier-1",
110 xai.XAI_REDIRECT_URI,
111 provider_id=XAI_GROK_PROVIDER_ID,
112 extra={"token_endpoint": "https://auth.x.ai/oauth/token", "code_challenge": "challenge-1"},
113 )
114 calls = []
115
116 monkeypatch.setattr(provider, "auth_path", lambda: auth_path)
117 monkeypatch.setattr(provider, "discovery", _discovery)
118
119 def fake_exchange(token_endpoint, code, redirect_uri, code_verifier, code_challenge):
120 calls.append((token_endpoint, code, redirect_uri, code_verifier, code_challenge))
121 return {
122 "access_token": "access-token",
123 "refresh_token": "refresh-token",
124 "expires_in": 3600,
125 "id_token": "id-token",
126 "token_type": "Bearer",
127 }
128
129 monkeypatch.setattr(provider, "exchange_code", fake_exchange)
130
131 result = provider.manual_callback(
132 {"callback": "http://127.0.0.1:56121/callback?code=code-1&state=state-1"}
133 )
134
135 assert result.ok is True
136 assert result.completed is True
137 assert result.account_label == "xAI Grok"
138 assert calls == [
139 (
140 "https://auth.x.ai/oauth/token",
141 "code-1",
142 "http://127.0.0.1:56121/callback",
143 "verifier-1",
144 "challenge-1",
145 )
146 ]
147 saved = json.loads(auth_path.read_text(encoding="utf-8"))
148 assert saved["provider"] == XAI_GROK_PROVIDER_ID
149 assert saved["access"] == "access-token"
150 assert saved["refresh"] == "refresh-token"
151
152
153 def test_models_returns_curated_list_without_network(monkeypatch):
154 provider = xai.XaiGrokOAuthProvider()
155 monkeypatch.setattr(provider, "read_auth", lambda: {})
156
157 models = provider.models()
158
159 assert models[:4] == [
160 "grok-4.3",
161 "grok-4.20-0309-reasoning",
162 "grok-4.20-0309-non-reasoning",
163 "grok-4.20-multi-agent-0309",
164 ]
165
166
167 def test_exchange_code_http_403_explains_oauth_tier_restriction(monkeypatch):
168 class FakeResponse:
169 ok = False
170 status_code = 403
171 text = "forbidden"
172
173 def json(self):
174 return {"error": "forbidden"}
175
176 class FakeRequests:
177 @staticmethod
178 def post(*args, **kwargs):
179 return FakeResponse()
180
181 monkeypatch.setitem(sys.modules, "requests", FakeRequests)
182
183 provider = xai.XaiGrokOAuthProvider()
184 with pytest.raises(ProviderError) as exc_info:
185 provider.exchange_code(
186 "https://auth.x.ai/oauth/token",
187 "code",
188 xai.XAI_REDIRECT_URI,
189 "verifier",
190 "challenge",
191 )
192
193 assert exc_info.value.status == 403
194 assert "restricted by tier" in str(exc_info.value)
195 assert "API-key `xai` provider" in str(exc_info.value)
196
197
198 def test_refresh_403_preserves_oauth_tier_guidance(monkeypatch):
199 class FakeResponse:
200 ok = False
201 status_code = 403
202
203 def json(self):
204 return {"error": "forbidden"}
205
206 class FakeRequests:
207 @staticmethod
208 def post(*args, **kwargs):
209 return FakeResponse()
210
211 monkeypatch.setitem(sys.modules, "requests", FakeRequests)
212
213 provider = xai.XaiGrokOAuthProvider()
214 monkeypatch.setattr(
215 provider,
216 "read_auth",
217 lambda: {
218 "access": "expired-access-token",
219 "refresh": "refresh-token",
220 "expires": 1,
221 "token_endpoint": "https://auth.x.ai/oauth/token",
222 },
223 )
224
225 with pytest.raises(ProviderError) as exc_info:
226 provider.ensure_fresh_auth()
227
228 assert exc_info.value.status == 403
229 assert exc_info.value.code == "oauth_tier_restricted"
230 assert "restricted by tier" in str(exc_info.value)
231 assert "API-key `xai` provider" in str(exc_info.value)
232
233
234 def test_ensure_fresh_auth_rejects_malicious_stored_token_endpoint(monkeypatch):
235 calls = []
236
237 class FakeRequests:
238 @staticmethod
239 def post(*args, **kwargs):
240 calls.append((args, kwargs))
241 raise AssertionError("malicious token endpoint must not be called")
242
243 monkeypatch.setitem(sys.modules, "requests", FakeRequests)
244
245 provider = xai.XaiGrokOAuthProvider()
246 monkeypatch.setattr(
247 provider,
248 "read_auth",
249 lambda: {
250 "access": "expired-access-token",
251 "refresh": "refresh-token",
252 "expires": 1,
253 "token_endpoint": "https://evil.example/oauth/token",
254 },
255 )
256
257 with pytest.raises(ProviderError) as exc_info:
258 provider.ensure_fresh_auth()
259
260 assert exc_info.value.code == "invalid_token_endpoint"
261 assert calls == []
262
263
264 def test_models_does_not_send_bearer_token_to_malicious_base_url(monkeypatch):
265 calls = []
266
267 class FakeResponse:
268 ok = True
269
270 def json(self):
271 return {"data": [{"id": "safe-model"}]}
272
273 class FakeRequests:
274 @staticmethod
275 def get(url, headers, timeout):
276 calls.append((url, headers, timeout))
277 return FakeResponse()
278
279 monkeypatch.setitem(sys.modules, "requests", FakeRequests)
280
281 provider = xai.XaiGrokOAuthProvider()
282 monkeypatch.setattr(
283 provider,
284 "read_auth",
285 lambda: {
286 "access": "access-token",
287 "refresh": "refresh-token",
288 "expires": int(time.time() * 1000) + 3_600_000,
289 "base_url": "https://evil.example/v1",
290 },
291 )
292
293 assert provider.models() == ["safe-model"]
294 assert calls == [
295 (
296 "https://api.x.ai/v1/models",
297 {"Accept": "application/json", "Authorization": "Bearer access-token"},
298 30,
299 )
300 ]