Coverage for agentos/tests/test_auth.py: 0%

1"""Tests for agentos.core.auth — Authenticator, RBACEngine, HS256TokenProvider.""" 

2 

3import time 

4 

5import pytest 

6 

7from agentos.core.auth import ( 

8 DEFAULT_ROLES, 

9 Algorithm, 

10 ApiKeyEntry, 

11 AuthContext, 

12 Authenticator, 

13 AuthMethod, 

14 HS256TokenProvider, 

15 InMemoryCredentialStore, 

16 InMemoryTokenBlacklist, 

17 Permission, 

18 RBACEngine, 

19 Role, 

20 TokenClaims, 

21 _b64url_decode, 

22 _b64url_encode, 

23 require_auth, 

24) 

25 

26# ============================================================================ 

27# Enums & Data classes 

28# ============================================================================ 

29 

30class TestAuthMethod: 

31 def test_values(self): 

32 assert AuthMethod.NONE.value == "none" 

33 assert AuthMethod.JWT.value == "jwt" 

34 assert AuthMethod.API_KEY.value == "api_key" 

35 

36 

37class TestAlgorithm: 

38 def test_values(self): 

39 assert Algorithm.HS256.value == "HS256" 

40 assert Algorithm.RS256.value == "RS256" 

41 

42 

43class TestPermission: 

44 def test_create(self): 

45 p = Permission("agent", "read", "own") 

46 assert p.resource == "agent" 

47 assert p.action == "read" 

48 assert p.scope == "own" 

49 

50 def test_default_scope(self): 

51 p = Permission("agent", "read") 

52 assert p.scope == "*" 

53 

54 

55class TestRole: 

56 def test_create(self): 

57 r = Role("admin", (Permission("*", "*"),)) 

58 assert r.name == "admin" 

59 assert len(r.permissions) == 1 

60 

61 

62class TestAuthContext: 

63 def test_defaults(self): 

64 ctx = AuthContext() 

65 assert ctx.authenticated is False 

66 assert ctx.method == AuthMethod.NONE 

67 assert ctx.subject is None 

68 

69 def test_authenticated(self): 

70 ctx = AuthContext(authenticated=True, method=AuthMethod.JWT, subject="u1", roles={"admin"}) 

71 assert ctx.authenticated 

72 assert "admin" in ctx.roles 

73 

74 

75# ============================================================================ 

76# TokenClaims 

77# ============================================================================ 

78 

79class TestTokenClaims: 

80 def test_create(self): 

81 tc = TokenClaims(sub="user1", roles=["admin"]) 

82 assert tc.sub == "user1" 

83 assert tc.iss == "agentos" 

84 assert len(tc.jti) == 32 

85 

86 def test_to_dict(self): 

87 tc = TokenClaims(sub="u1", exp=1000.0, roles=["r1"]) 

88 d = tc.to_dict() 

89 assert d["sub"] == "u1" 

90 assert d["exp"] == 1000 

91 assert d["roles"] == ["r1"] 

92 

93 def test_from_dict(self): 

94 d = {"sub": "u1", "roles": ["admin"], "extra_field": "x"} 

95 tc = TokenClaims.from_dict(d) 

96 assert tc.sub == "u1" 

97 assert tc.roles == ["admin"] 

98 assert tc.extra == {"extra_field": "x"} 

99 

100 def test_roundtrip(self): 

101 tc = TokenClaims(sub="u1", roles=["r1"], permissions=["agent:read"], extra={"key": "val"}) 

102 d = tc.to_dict() 

103 tc2 = TokenClaims.from_dict(d) 

104 assert tc2.sub == tc.sub 

105 assert tc2.roles == tc.roles 

106 assert tc2.extra == tc.extra 

107 

108 def test_from_dict_minimal(self): 

109 d = {"sub": "u1"} 

110 tc = TokenClaims.from_dict(d) 

111 assert tc.iat is not None 

112 

113 

114# ============================================================================ 

115# InMemoryTokenBlacklist 

116# ============================================================================ 

117 

118class TestInMemoryTokenBlacklist: 

119 @pytest.mark.asyncio 

120 async def test_add_and_check(self): 

121 bl = InMemoryTokenBlacklist() 

122 await bl.add("jti1", 60) 

123 assert await bl.is_blacklisted("jti1") 

124 

125 @pytest.mark.asyncio 

126 async def test_not_blacklisted(self): 

127 bl = InMemoryTokenBlacklist() 

128 assert not await bl.is_blacklisted("missing") 

129 

130 @pytest.mark.asyncio 

131 async def test_expired_removed(self): 

132 bl = InMemoryTokenBlacklist() 

133 await bl.add("jti1", 0.001) 

134 import asyncio 

135 await asyncio.sleep(0.01) 

136 assert not await bl.is_blacklisted("jti1") 

137 

138 @pytest.mark.asyncio 

139 async def test_cleanup(self): 

140 bl = InMemoryTokenBlacklist() 

141 await bl.add("old", 0.001) 

142 await bl.add("new", 60) 

143 import asyncio 

144 await asyncio.sleep(0.01) 

145 bl._cleanup() 

146 assert "old" not in bl._store 

147 assert "new" in bl._store 

148 

149 

150# ============================================================================ 

151# InMemoryCredentialStore 

152# ============================================================================ 

153 

154class TestInMemoryCredentialStore: 

155 def test_add_and_lookup(self): 

156 store = InMemoryCredentialStore() 

157 store.add_key("raw_key_123", ApiKeyEntry( 

158 key_hash="", subject="u1", name="my key", scopes=["read"], roles=["viewer"], permissions=[] 

159 )) 

160 entry = asyncio_run(store.lookup_by_key("raw_key_123")) 

161 assert entry["subject"] == "u1" 

162 assert entry["roles"] == ["viewer"] 

163 

164 def test_lookup_missing(self): 

165 store = InMemoryCredentialStore() 

166 assert asyncio_run(store.lookup_by_key("bad")) is None 

167 

168 def test_revoke(self): 

169 store = InMemoryCredentialStore() 

170 store.add_key("k1", ApiKeyEntry(key_hash="", subject="u1", name="x", scopes=[], roles=[], permissions=[])) 

171 assert asyncio_run(store.revoke("k1")) is True 

172 assert asyncio_run(store.lookup_by_key("k1")) is None 

173 

174 def test_revoke_missing(self): 

175 store = InMemoryCredentialStore() 

176 assert asyncio_run(store.revoke("bad")) is False 

177 

178 def test_expired_key(self): 

179 store = InMemoryCredentialStore() 

180 store.add_key("k1", ApiKeyEntry( 

181 key_hash="", subject="u1", name="x", scopes=[], roles=[], permissions=[], 

182 expires_at=time.time() - 60, 

183 )) 

184 assert asyncio_run(store.lookup_by_key("k1")) is None 

185 

186 

187# ============================================================================ 

188# HS256TokenProvider 

189# ============================================================================ 

190 

191class TestHS256TokenProvider: 

192 @pytest.mark.asyncio 

193 async def test_create_and_validate(self): 

194 p = HS256TokenProvider("secret") 

195 tc = TokenClaims(sub="u1") 

196 token = await p.create_token(tc) 

197 claims = await p.validate_token(token) 

198 assert claims.sub == "u1" 

199 

200 @pytest.mark.asyncio 

201 async def test_validate_invalid_token(self): 

202 p = HS256TokenProvider("secret") 

203 assert await p.validate_token("bad.token.here") is None 

204 

205 @pytest.mark.asyncio 

206 async def test_validate_wrong_secret(self): 

207 p1 = HS256TokenProvider("s1") 

208 p2 = HS256TokenProvider("s2") 

209 token = await p1.create_token(TokenClaims(sub="u1")) 

210 assert await p2.validate_token(token) is None 

211 

212 @pytest.mark.asyncio 

213 async def test_key_rotation_new_token(self): 

214 p = HS256TokenProvider("old_secret") 

215 p.rotate_secret("new_secret") 

216 tc = TokenClaims(sub="u1") 

217 token = await p.create_token(tc) 

218 claims = await p.validate_token(token) 

219 assert claims.sub == "u1" 

220 

221 @pytest.mark.asyncio 

222 async def test_expired_token(self): 

223 p = HS256TokenProvider("s", default_ttl=0.001) 

224 tc = TokenClaims(sub="u1") 

225 tc.exp = time.time() - 10 

226 token = await p.create_token(tc) 

227 import asyncio 

228 await asyncio.sleep(0.01) 

229 assert await p.validate_token(token) is None 

230 

231 @pytest.mark.asyncio 

232 async def test_default_expiry_set(self): 

233 p = HS256TokenProvider("s", default_ttl=3600) 

234 tc = TokenClaims(sub="u1", exp=None) 

235 token = await p.create_token(tc) 

236 claims = await p.validate_token(token) 

237 assert claims.exp is not None 

238 

239 @pytest.mark.asyncio 

240 async def test_key_rotation_grace_period(self): 

241 p = HS256TokenProvider("secret") 

242 token_old = await p.create_token(TokenClaims(sub="u1")) 

243 p.rotate_secret("new_secret") 

244 assert await p.validate_token(token_old) is not None 

245 

246 @pytest.mark.asyncio 

247 async def test_validate_two_part_token(self): 

248 p = HS256TokenProvider("s") 

249 assert await p.validate_token("a.b") is None 

250 

251 

252# ============================================================================ 

253# RBACEngine 

254# ============================================================================ 

255 

256class TestRBACEngine: 

257 def test_register_role(self): 

258 rbac = RBACEngine() 

259 rbac.register_role(Role("admin", (Permission("*", "*"),))) 

260 perms = rbac.get_permissions({"admin"}) 

261 assert len(perms) == 1 

262 

263 def test_register_roles(self): 

264 rbac = RBACEngine() 

265 rbac.register_roles([Role("a", (Permission("x", "y"),)), Role("b", (Permission("z", "w"),))]) 

266 assert len(rbac.get_permissions({"a", "b"})) == 2 

267 

268 def test_check_wildcard(self): 

269 rbac = RBACEngine() 

270 rbac.register_role(Role("admin", (Permission("*", "*"),))) 

271 assert rbac.check({"admin"}, Permission("agent", "execute")) 

272 

273 def test_check_exact(self): 

274 rbac = RBACEngine() 

275 rbac.register_role(Role("dev", (Permission("agent", "read"),))) 

276 assert rbac.check({"dev"}, Permission("agent", "read")) 

277 assert not rbac.check({"dev"}, Permission("agent", "write")) 

278 

279 def test_check_any(self): 

280 rbac = RBACEngine() 

281 rbac.register_role(Role("dev", (Permission("agent", "read"),))) 

282 assert rbac.check_any({"dev"}, [Permission("agent", "write"), Permission("agent", "read")]) 

283 

284 def test_check_all(self): 

285 rbac = RBACEngine() 

286 rbac.register_role(Role("dev", (Permission("agent", "read"), Permission("agent", "execute")))) 

287 assert rbac.check_all({"dev"}, [Permission("agent", "read"), Permission("agent", "execute")]) 

288 assert not rbac.check_all({"dev"}, [Permission("agent", "read"), Permission("agent", "delete")]) 

289 

290 def test_hierarchical_scope(self): 

291 rbac = RBACEngine() 

292 rbac.register_role(Role("admin", (Permission("*", "*", "org:engineering"),))) 

293 assert rbac.check({"admin"}, Permission("agent", "read", "org:engineering")) 

294 # Narrower doesn't match broader 

295 assert not rbac.check({"admin"}, Permission("agent", "read", "org")) 

296 

297 def test_get_permissions_missing_role(self): 

298 rbac = RBACEngine() 

299 assert rbac.get_permissions({"missing"}) == set() 

300 

301 def test_wildcard_action_in_perm(self): 

302 rbac = RBACEngine() 

303 rbac.register_role(Role("op", (Permission("agent", "*"),))) 

304 assert rbac.check({"op"}, Permission("agent", "execute")) 

305 

306 

307# ============================================================================ 

308# Authenticator 

309# ============================================================================ 

310 

311class TestAuthenticator: 

312 @pytest.mark.asyncio 

313 async def test_no_auth_header(self): 

314 auth = Authenticator() 

315 ctx = await auth.authenticate({}) 

316 assert not ctx.authenticated 

317 

318 @pytest.mark.asyncio 

319 async def test_jwt_authentication(self): 

320 jwt = HS256TokenProvider("secret") 

321 auth = Authenticator(token_provider=jwt) 

322 token = await jwt.create_token(TokenClaims(sub="u1", roles=["admin"])) 

323 ctx = await auth.authenticate({"Authorization": f"Bearer {token}"}) 

324 assert ctx.authenticated 

325 assert ctx.method == AuthMethod.JWT 

326 assert ctx.subject == "u1" 

327 assert "admin" in ctx.roles 

328 

329 @pytest.mark.asyncio 

330 async def test_jwt_invalid_token(self): 

331 jwt = HS256TokenProvider("secret") 

332 auth = Authenticator(token_provider=jwt) 

333 ctx = await auth.authenticate({"Authorization": "Bearer bad.token.here"}) 

334 assert not ctx.authenticated 

335 

336 @pytest.mark.asyncio 

337 async def test_api_key_authentication(self): 

338 store = InMemoryCredentialStore() 

339 store.add_key("secret_key", ApiKeyEntry( 

340 key_hash="", subject="u1", name="test", scopes=[], roles=["viewer"], permissions=["agent:read"] 

341 )) 

342 auth = Authenticator(credential_store=store) 

343 ctx = await auth.authenticate({"X-API-Key": "secret_key"}) 

344 assert ctx.authenticated 

345 assert ctx.method == AuthMethod.API_KEY 

346 assert ctx.subject == "u1" 

347 assert "viewer" in ctx.roles 

348 

349 @pytest.mark.asyncio 

350 async def test_api_key_bad(self): 

351 store = InMemoryCredentialStore() 

352 auth = Authenticator(credential_store=store) 

353 ctx = await auth.authenticate({"X-API-Key": "bad"}) 

354 assert not ctx.authenticated 

355 

356 @pytest.mark.asyncio 

357 async def test_jwt_blacklisted(self): 

358 jwt = HS256TokenProvider("secret") 

359 bl = InMemoryTokenBlacklist() 

360 token = await jwt.create_token(TokenClaims(sub="u1")) 

361 tc = await jwt.validate_token(token) 

362 await bl.add(tc.jti, 60) 

363 auth = Authenticator(token_provider=jwt, blacklist=bl) 

364 ctx = await auth.authenticate({"Authorization": f"Bearer {token}"}) 

365 assert not ctx.authenticated 

366 

367 @pytest.mark.asyncio 

368 async def test_jwt_with_permissions(self): 

369 jwt = HS256TokenProvider("secret") 

370 auth = Authenticator(token_provider=jwt) 

371 token = await jwt.create_token(TokenClaims(sub="u1", permissions=["agent:read:own"])) 

372 ctx = await auth.authenticate({"Authorization": f"Bearer {token}"}) 

373 assert ctx.authenticated 

374 assert len(ctx.permissions) == 1 

375 

376 @pytest.mark.asyncio 

377 async def test_oauth2_stub(self): 

378 auth = Authenticator() 

379 ctx = await auth.authenticate({"Authorization": "Bearer oauth_token"}) 

380 assert not ctx.authenticated # Stub returns unauthenticated 

381 

382 @pytest.mark.asyncio 

383 async def test_api_key_query_param(self): 

384 store = InMemoryCredentialStore() 

385 store.add_key("k1", ApiKeyEntry(key_hash="", subject="u1", name="x", scopes=[], roles=[], permissions=[])) 

386 auth = Authenticator(credential_store=store) 

387 ctx = await auth.authenticate({}, {"api_key": "k1"}) 

388 assert ctx.authenticated 

389 

390 

391# ============================================================================ 

392# require_auth decorator 

393# ============================================================================ 

394 

395class TestRequireAuth: 

396 @pytest.mark.asyncio 

397 async def test_missing_auth_context(self): 

398 @require_auth() 

399 async def handler(*args, **kwargs): 

400 return "ok" 

401 

402 with pytest.raises(PermissionError, match="auth_context"): 

403 await handler() 

404 

405 @pytest.mark.asyncio 

406 async def test_not_authenticated(self): 

407 @require_auth() 

408 async def handler(*args, **kwargs): 

409 return "ok" 

410 

411 with pytest.raises(PermissionError, match="authentication"): 

412 await handler(auth_context=AuthContext()) 

413 

414 @pytest.mark.asyncio 

415 async def test_authenticated_no_permission(self): 

416 @require_auth() 

417 async def handler(*args, **kwargs): 

418 return "ok" 

419 

420 ctx = AuthContext(authenticated=True, method=AuthMethod.JWT, subject="u1") 

421 result = await handler(auth_context=ctx) 

422 assert result == "ok" 

423 

424 @pytest.mark.asyncio 

425 async def test_permission_check_pass(self): 

426 rbac = RBACEngine() 

427 rbac.register_role(Role("admin", (Permission("*", "*"),))) 

428 

429 @require_auth(permission=Permission("agent", "execute")) 

430 async def handler(*args, **kwargs): 

431 return "ok" 

432 

433 ctx = AuthContext(authenticated=True, roles={"admin"}) 

434 result = await handler(auth_context=ctx, _rbac=rbac) 

435 assert result == "ok" 

436 

437 @pytest.mark.asyncio 

438 async def test_permission_check_fail(self): 

439 rbac = RBACEngine() 

440 rbac.register_role(Role("viewer", (Permission("agent", "read"),))) 

441 

442 @require_auth(permission=Permission("agent", "delete")) 

443 async def handler(*args, **kwargs): 

444 return "ok" 

445 

446 ctx = AuthContext(authenticated=True, roles={"viewer"}) 

447 with pytest.raises(PermissionError, match="missing permissions"): 

448 await handler(auth_context=ctx, _rbac=rbac) 

449 

450 

451# ============================================================================ 

452# _b64url helpers 

453# ============================================================================ 

454 

455class TestB64: 

456 def test_roundtrip(self): 

457 data = b"hello world" 

458 assert _b64url_decode(_b64url_encode(data)) == data 

459 

460 def test_encode_no_padding(self): 

461 enc = _b64url_encode(b"test") 

462 assert "=" not in enc 

463 

464 

465# ============================================================================ 

466# DEFAULT_ROLES 

467# ============================================================================ 

468 

469class TestDefaultRoles: 

470 def test_four_roles(self): 

471 assert len(DEFAULT_ROLES) == 4 

472 names = {r.name for r in DEFAULT_ROLES} 

473 assert names == {"admin", "developer", "viewer", "operator"} 

474 

475 def test_admin_has_all(self): 

476 admin = next(r for r in DEFAULT_ROLES if r.name == "admin") 

477 assert any(p.resource == "*" and p.action == "*" for p in admin.permissions) 

478 

479 

480# ============================================================================ 

481# helper 

482# ============================================================================ 

483 

484def asyncio_run(coro): 

485 import asyncio 

486 try: 

487 _ = asyncio.get_running_loop() 

488 except RuntimeError: 

489 return asyncio.run(coro) 

490 import concurrent.futures 

491 with concurrent.futures.ThreadPoolExecutor() as ex: 

492 return ex.submit(asyncio.run, coro).result()