Coverage for agentos/core/auth.py: 99%

1"""

2Production-grade authentication and authorization framework.

4Supports:

5- JWT (HS256/RS256/ES256) with key rotation

6- OAuth2 Bearer token validation

7- API Key (static + scoped)

8- Role-Based Access Control (RBAC)

9- Fine-grained Permission model

10- Token blacklisting

11- Multi-issuer trust

14"""

16from __future__ import annotations

18import hashlib

19import hmac

20import json

21import time

22import uuid

23from abc import ABC, abstractmethod

24from dataclasses import dataclass, field

25from enum import Enum

26from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union

28# ---------------------------------------------------------------------------

29# Constants & Enums

30# ---------------------------------------------------------------------------

32class AuthMethod(Enum):

33 NONE = "none"

34 JWT = "jwt"

35 API_KEY = "api_key"

36 OAUTH2_BEARER = "oauth2_bearer"

37 CUSTOM = "custom"

40class Algorithm(Enum):

41 HS256 = "HS256"

42 HS384 = "HS384"

43 HS512 = "HS512"

44 RS256 = "RS256"

45 RS384 = "RS384"

46 RS512 = "RS512"

47 ES256 = "ES256"

48 ES384 = "ES384"

49 ES512 = "ES512"

52@dataclass(frozen=True, slots=True)

53class Permission:

54 """Fine-grained permission atom."""

55 resource: str # e.g. "agent", "model", "user"

56 action: str # e.g. "read", "write", "delete", "execute"

57 scope: str = "*" # e.g. "own", "team", "org:*"

60@dataclass(frozen=True, slots=True)

61class Role:

62 """Named collection of permissions."""

63 name: str

64 permissions: Tuple[Permission, ...] = field(default_factory=tuple)

67@dataclass

68class AuthContext:

69 """Authentication result passed through the request lifecycle."""

70 authenticated: bool = False

71 method: AuthMethod = AuthMethod.NONE

72 subject: Optional[str] = None # user ID / API key ID

73 issuer: Optional[str] = None

74 roles: Set[str] = field(default_factory=set)

75 permissions: Set[Permission] = field(default_factory=set)

76 metadata: Dict[str, Any] = field(default_factory=dict)

77 expires_at: Optional[float] = None

78 token_id: Optional[str] = None # jti

81# ---------------------------------------------------------------------------

82# Token Model

83# ---------------------------------------------------------------------------

85@dataclass

86class TokenClaims:

87 """Standard JWT claims as specified in RFC 7519."""

88 sub: str

89 iat: float = field(default_factory=time.time)

90 exp: Optional[float] = None

91 iss: str = "agentos"

92 aud: Optional[Union[str, List[str]]] = None

93 jti: str = field(default_factory=lambda: uuid.uuid4().hex)

94 roles: List[str] = field(default_factory=list)

95 permissions: List[str] = field(default_factory=list)

96 extra: Dict[str, Any] = field(default_factory=dict)

98 def to_dict(self) -> Dict[str, Any]:

99 d = {

100 "sub": self.sub,

101 "iat": int(self.iat),

102 "iss": self.iss,

103 "jti": self.jti,

104 "roles": self.roles,

105 "permissions": self.permissions,

106 }

107 if self.exp is not None:

108 d["exp"] = int(self.exp)

109 if self.aud is not None:

110 d["aud"] = self.aud

111 d.update(self.extra)

112 return d

113

114 @classmethod

115 def from_dict(cls, d: Dict[str, Any]) -> "TokenClaims":

116 return cls(

117 sub=d["sub"],

118 iat=float(d.get("iat", time.time())),

119 exp=float(d["exp"]) if "exp" in d else None,

120 iss=d.get("iss", "agentos"),

121 aud=d.get("aud"),

122 jti=d.get("jti", uuid.uuid4().hex),

123 roles=d.get("roles", []),

124 permissions=d.get("permissions", []),

125 extra={k: v for k, v in d.items()

126 if k not in {"sub", "iat", "exp", "iss", "aud", "jti", "roles", "permissions"}},

127 )

128

129

130# ---------------------------------------------------------------------------

131# Abstract Providers

132# ---------------------------------------------------------------------------

133

134class TokenProvider(ABC):

135 """Abstract interface for creating and validating tokens."""

136

137 @abstractmethod

138 async def create_token(self, claims: TokenClaims) -> str: ...

139

140 @abstractmethod

141 async def validate_token(self, token: str) -> Optional[TokenClaims]: ...

142

143

144class CredentialStore(ABC):

145 """Abstract store for API keys and secrets."""

146

147 @abstractmethod

148 async def lookup_by_key(self, api_key: str) -> Optional[Dict[str, Any]]: ...

149

150 @abstractmethod

151 async def revoke(self, api_key: str) -> bool: ...

152

153

154class TokenBlacklist(ABC):

155 """Abstract token blacklist (jti-based revocation)."""

156

157 @abstractmethod

158 async def is_blacklisted(self, jti: str) -> bool: ...

159

160 @abstractmethod

161 async def add(self, jti: str, ttl: float) -> None: ...

162

163

164# ---------------------------------------------------------------------------

165# In-Memory Token Blacklist

166# ---------------------------------------------------------------------------

167

168class InMemoryTokenBlacklist(TokenBlacklist):

169 """Simple in-memory blacklist with TTL-based eviction."""

170

171 def __init__(self):

172 self._store: Dict[str, float] = {} # jti → expiry_time

173

174 async def is_blacklisted(self, jti: str) -> bool:

175 now = time.monotonic()

176 if jti in self._store:

177 if self._store[jti] > now:

178 return True

179 del self._store[jti]

180 return False

181

182 async def add(self, jti: str, ttl: float) -> None:

183 self._store[jti] = time.monotonic() + ttl

184

185 def _cleanup(self):

186 now = time.monotonic()

187 self._store = {k: v for k, v in self._store.items() if v > now}

188

189

190# ---------------------------------------------------------------------------

191# In-Memory Credential Store

192# ---------------------------------------------------------------------------

193

194@dataclass

195class ApiKeyEntry:

196 key_hash: str

197 subject: str

198 name: str

199 scopes: List[str]

200 roles: List[str]

201 permissions: List[str]

202 created_at: float = field(default_factory=time.time)

203 expires_at: Optional[float] = None

204 revoked: bool = False

205

206

207class InMemoryCredentialStore(CredentialStore):

208 """In-memory API key store with SHA-256 hashing."""

209

210 def __init__(self):

211 self._keys: Dict[str, ApiKeyEntry] = {}

212

213 @staticmethod

214 def _hash_key(raw_key: str) -> str:

215 return hashlib.sha256(raw_key.encode()).hexdigest()

216

217 def add_key(self, raw_key: str, entry: ApiKeyEntry) -> None:

218 entry.key_hash = self._hash_key(raw_key)

219 self._keys[entry.key_hash] = entry

220

221 async def lookup_by_key(self, api_key: str) -> Optional[Dict[str, Any]]:

222 key_hash = self._hash_key(api_key)

223 entry = self._keys.get(key_hash)

224 if entry is None or entry.revoked:

225 return None

226 if entry.expires_at is not None and entry.expires_at < time.time():

227 return None

228 return {

229 "subject": entry.subject,

230 "name": entry.name,

231 "scopes": entry.scopes,

232 "roles": entry.roles,

233 "permissions": entry.permissions,

234 }

235

236 async def revoke(self, api_key: str) -> bool:

237 key_hash = self._hash_key(api_key)

238 entry = self._keys.get(key_hash)

239 if entry is not None:

240 entry.revoked = True

241 return True

242 return False

243

244

245# ---------------------------------------------------------------------------

246# JWT Provider (HS256-only; RS/ES require cryptography)

247# ---------------------------------------------------------------------------

248

249class HS256TokenProvider(TokenProvider):

250 """HS256 HMAC-based JWT provider with key rotation support."""

251

252 def __init__(self, secret: str, issuer: str = "agentos",

253 default_ttl: float = 3600.0):

254 self._current_secret = secret.encode()

255 self._previous_secret: Optional[bytes] = None

256 self._issuer = issuer

257 self._default_ttl = default_ttl

258

259 def rotate_secret(self, new_secret: str):

260 """Rotate to new secret; old secret retained for validation grace period."""

261 self._previous_secret = self._current_secret

262 self._current_secret = new_secret.encode()

263

264 def _encode(self, claims: TokenClaims) -> str:

265 header = {"alg": "HS256", "typ": "JWT"}

266 segments = [

267 _b64url_encode(json.dumps(header, separators=(",", ":")).encode()),

268 _b64url_encode(json.dumps(claims.to_dict(), separators=(",", ":")).encode()),

269 ]

270 signing_input = f"{segments[0]}.{segments[1]}".encode()

271 signature = hmac.new(self._current_secret, signing_input, hashlib.sha256).digest()

272 segments.append(_b64url_encode(signature))

273 return ".".join(segments)

274

275 def _decode(self, token: str) -> Optional[TokenClaims]:

276 parts = token.split(".")

277 if len(parts) != 3:

278 return None

279 for secret in (self._current_secret, self._previous_secret):

280 if secret is None:

281 continue

282 signing_input = f"{parts[0]}.{parts[1]}".encode()

283 expected_sig = hmac.new(

284 secret, signing_input, hashlib.sha256

285 ).digest()

286 actual_sig = _b64url_decode(parts[2])

287 if hmac.compare_digest(expected_sig, actual_sig):

288 payload = json.loads(_b64url_decode(parts[1]))

289 return TokenClaims.from_dict(payload)

290 return None

291

292 async def create_token(self, claims: TokenClaims) -> str:

293 if claims.exp is None:

294 claims.exp = time.time() + self._default_ttl

295 if claims.iss != self._issuer:

296 claims.iss = self._issuer

297 return self._encode(claims)

298

299 async def validate_token(self, token: str) -> Optional[TokenClaims]:

300 claims = self._decode(token)

301 if claims is None:

302 return None

303 if claims.exp is not None and claims.exp < time.time():

304 return None

305 return claims

306

307

308# ---------------------------------------------------------------------------

309# RBAC Engine

310# ---------------------------------------------------------------------------

311

312class RBACEngine:

313 """Role-based access control engine with role-permission resolution."""

314

315 def __init__(self):

316 self._roles: Dict[str, Role] = {}

317

318 def register_role(self, role: Role) -> None:

319 self._roles[role.name] = role

320

321 def register_roles(self, roles: List[Role]) -> None:

322 for role in roles:

323 self.register_role(role)

324

325 def get_permissions(self, role_names: Set[str]) -> Set[Permission]:

326 result: Set[Permission] = set()

327 for name in role_names:

328 role = self._roles.get(name)

329 if role is not None:

330 result.update(role.permissions)

331 return result

332

333 def check(self, role_names: Set[str], required: Permission) -> bool:

334 for perm in self.get_permissions(role_names):

335 if self._match(perm, required):

336 return True

337 return False

338

339 def check_any(self, role_names: Set[str], required: List[Permission]) -> bool:

340 perms = self.get_permissions(role_names)

341 return any(self._match(p, r) for r in required for p in perms)

342

343 def check_all(self, role_names: Set[str], required: List[Permission]) -> bool:

344 perms = self.get_permissions(role_names)

345 return all(

346 any(self._match(p, r) for p in perms)

347 for r in required

348 )

349

350 @staticmethod

351 def _match(perm: Permission, required: Permission) -> bool:

352 def _seg_match(have: str, need: str) -> bool:

353 if have == "*" or need == "*":

354 return True

355 if ":" in need:

356 # hierarchical: "org:engineering"

357 parts_need = need.split(":")

358 parts_have = have.split(":")

359 if len(parts_have) < len(parts_need):

360 return False

361 return all(h == n for h, n in zip(parts_have[:len(parts_need)], parts_need))

362 return have == need

363 return (

364 _seg_match(perm.resource, required.resource)

365 and _seg_match(perm.action, required.action)

366 and _seg_match(perm.scope, required.scope)

367 )

368

369

370# ---------------------------------------------------------------------------

371# Authenticator

372# ---------------------------------------------------------------------------

373

374@dataclass

375class AuthenticatorConfig:

376 """Authenticator configuration."""

377 allowed_methods: Tuple[AuthMethod, ...] = (

378 AuthMethod.JWT, AuthMethod.API_KEY, AuthMethod.OAUTH2_BEARER

379 )

380 default_issuer: str = "agentos"

381 api_key_header: str = "X-API-Key"

382 api_key_query_param: str = "api_key"

383 require_auth: bool = True

384 token_strict_expiry: bool = True

385 clock_skew: float = 30.0 # seconds

386 max_token_lifetime: float = 86400.0 # 24h

387

388

389class Authenticator:

390 """Main authentication orchestrator.

391

392 Coordinates JWT validation, API key lookup, OAuth2 introspection,

393 and RBAC resolution into a single `authenticate` entry point.

394

395 Usage:

396 auth = Authenticator(token_provider=jwt, credential_store=keys,

397 blacklist=bl, rbac=rbac)

398 ctx = await auth.authenticate(request_headers)

399 if ctx.authenticated and rbac.check(ctx.roles, Permission("agent", "execute")):

400 ...

401 """

402

403 def __init__(

404 self,

405 *,

406 token_provider: Optional[TokenProvider] = None,

407 credential_store: Optional[CredentialStore] = None,

408 blacklist: Optional[TokenBlacklist] = None,

409 rbac: Optional[RBACEngine] = None,

410 config: Optional[AuthenticatorConfig] = None,

411 ):

412 self._token_provider = token_provider

413 self._credential_store = credential_store

414 self._blacklist = blacklist

415 self._rbac = rbac or RBACEngine()

416 self._config = config or AuthenticatorConfig()

417

418 @property

419 def rbac(self) -> RBACEngine:

420 return self._rbac

421

422 async def authenticate(

423 self,

424 headers: Dict[str, str],

425 query_params: Optional[Dict[str, str]] = None,

426 ) -> AuthContext:

427 """Authenticate a request from headers and query parameters.

428 Returns AuthContext with authenticated=False if auth fails.

429 """

430 # Try JWT Bearer

431 auth_header = headers.get("authorization", headers.get("Authorization", ""))

432 if auth_header.startswith("Bearer ") and AuthMethod.JWT in self._config.allowed_methods:

433 token = auth_header[7:]

434 ctx = await self._authenticate_jwt(token)

435 if ctx.authenticated:

436 return ctx

437

438 # Try OAuth2 Bearer

439 if auth_header.startswith("Bearer ") and AuthMethod.OAUTH2_BEARER in self._config.allowed_methods:

440 token = auth_header[7:]

441 ctx = await self._authenticate_oauth2_bearer(token)

442 if ctx.authenticated:

443 return ctx

444

445 # Try API Key

446 if AuthMethod.API_KEY in self._config.allowed_methods:

447 api_key = headers.get(self._config.api_key_header, "")

448 if not api_key and query_params:

449 api_key = query_params.get(self._config.api_key_query_param, "")

450 if api_key:

451 ctx = await self._authenticate_api_key(api_key)

452 if ctx.authenticated:

453 return ctx

454

455 return AuthContext()

456

457 async def _authenticate_jwt(self, token: str) -> AuthContext:

458 if self._token_provider is None:

459 return AuthContext()

460 claims = await self._token_provider.validate_token(token)

461 if claims is None:

462 return AuthContext()

463 if self._blacklist and claims.jti:

464 if await self._blacklist.is_blacklisted(claims.jti):

465 return AuthContext()

466 expires_at = claims.exp

467 if expires_at and self._config.token_strict_expiry:

468 if expires_at < time.time() + self._config.clock_skew:

469 return AuthContext()

470 perm_set = {self._parse_permission_str(p) for p in claims.permissions if p}

471 return AuthContext(

472 authenticated=True,

473 method=AuthMethod.JWT,

474 subject=claims.sub,

475 issuer=claims.iss,

476 roles=set(claims.roles),

477 permissions=perm_set,

478 metadata=claims.extra,

479 expires_at=expires_at,

480 token_id=claims.jti,

481 )

482

483 async def _authenticate_api_key(self, api_key: str) -> AuthContext:

484 if self._credential_store is None:

485 return AuthContext()

486 entry = await self._credential_store.lookup_by_key(api_key)

487 if entry is None:

488 return AuthContext()

489 perm_set = {self._parse_permission_str(p) for p in entry.get("permissions", []) if p}

490 return AuthContext(

491 authenticated=True,

492 method=AuthMethod.API_KEY,

493 subject=entry.get("subject", ""),

494 roles=set(entry.get("roles", [])),

495 permissions=perm_set,

496 metadata={"name": entry.get("name", "")},

497 )

498

499 async def _authenticate_oauth2_bearer(self, token: str) -> AuthContext:

500 """OAuth2 introspection stub — implement via introspection endpoint."""

501 return AuthContext()

502

503 @staticmethod

504 def _parse_permission_str(raw: str) -> Permission:

505 parts = raw.split(":", 2)

506 if len(parts) == 1:

507 return Permission(resource=parts[0], action="*")

508 elif len(parts) == 2:

509 return Permission(resource=parts[0], action=parts[1])

510 return Permission(resource=parts[0], action=parts[1], scope=parts[2])

511

512

513# ---------------------------------------------------------------------------

514# Decorators

515# ---------------------------------------------------------------------------

516

517def require_auth(permission: Optional[Permission] = None, permissions: Optional[List[Permission]] = None):

518 """Decorator to require authentication and optional permissions.

519 To be used with a framework that provides `auth_context` in the call scope.

520 """

521 required = permissions or ([permission] if permission else [])

522

523 def decorator(fn: Callable):

524 async def wrapper(*args, **kwargs):

525 ctx: AuthContext = kwargs.pop("auth_context", None)

526 if ctx is None:

527 raise PermissionError("auth_context required but not provided")

528 if not ctx.authenticated:

529 raise PermissionError("authentication required")

530 if required:

531 rbac = kwargs.pop("_rbac", None)

532 if rbac is None:

533 raise PermissionError("RBAC engine required for permission check")

534 if not rbac.check_all(ctx.roles, required):

535 raise PermissionError(

536 f"missing permissions: {[f'{p.resource}:{p.action}' for p in required]}"

537 )

538 return await fn(*args, **kwargs)

539 wrapper.__name__ = fn.__name__

540 wrapper.__doc__ = fn.__doc__

541 return wrapper

542 return decorator

543

544

545# ---------------------------------------------------------------------------

546# Helpers

547# ---------------------------------------------------------------------------

548

549def _b64url_encode(data: bytes) -> str:

550 import base64

551 return base64.urlsafe_b64encode(data).rstrip(b"=").decode()

552

553

554def _b64url_decode(data: str) -> bytes:

555 import base64

556 padding = 4 - len(data) % 4

557 if padding != 4:

558 data += "=" * padding

559 return base64.urlsafe_b64decode(data)

560

561

562# ---------------------------------------------------------------------------

563# Default Roles

564# ---------------------------------------------------------------------------

565

566DEFAULT_ROLES = [

567 Role("admin", (

568 Permission("*", "*", "*"),

569 )),

570 Role("developer", (

571 Permission("agent", "*"),

572 Permission("model", "read"),

573 Permission("model", "execute"),

574 Permission("tool", "*"),

575 Permission("log", "read"),

576 )),

577 Role("viewer", (

578 Permission("agent", "read"),

579 Permission("model", "read"),

580 Permission("log", "read"),

581 Permission("metric", "read"),

582 )),

583 Role("operator", (

584 Permission("agent", "read"),

585 Permission("agent", "execute"),

586 Permission("model", "read"),

587 Permission("model", "execute"),

588 Permission("tool", "execute"),

589 Permission("log", "read"),

590 Permission("metric", "read"),

591 )),

592]