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

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 os

22import time

23import uuid

24from abc import ABC, abstractmethod

25from dataclasses import dataclass, field

26from enum import Enum, auto

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

29# ---------------------------------------------------------------------------

30# Constants & Enums

31# ---------------------------------------------------------------------------

33class AuthMethod(Enum):

34 NONE = "none"

35 JWT = "jwt"

36 API_KEY = "api_key"

37 OAUTH2_BEARER = "oauth2_bearer"

38 CUSTOM = "custom"

41class Algorithm(Enum):

42 HS256 = "HS256"

43 HS384 = "HS384"

44 HS512 = "HS512"

45 RS256 = "RS256"

46 RS384 = "RS384"

47 RS512 = "RS512"

48 ES256 = "ES256"

49 ES384 = "ES384"

50 ES512 = "ES512"

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

54class Permission:

55 """Fine-grained permission atom."""

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

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

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

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

62class Role:

63 """Named collection of permissions."""

64 name: str

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

68@dataclass

69class AuthContext:

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

71 authenticated: bool = False

72 method: AuthMethod = AuthMethod.NONE

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

74 issuer: Optional[str] = None

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

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

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

78 expires_at: Optional[float] = None

79 token_id: Optional[str] = None # jti

82# ---------------------------------------------------------------------------

83# Token Model

84# ---------------------------------------------------------------------------

86@dataclass

87class TokenClaims:

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

89 sub: str

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

91 exp: Optional[float] = None

92 iss: str = "agentos"

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

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

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

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

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

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

100 d = {

101 "sub": self.sub,

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

103 "iss": self.iss,

104 "jti": self.jti,

105 "roles": self.roles,

106 "permissions": self.permissions,

107 }

108 if self.exp is not None:

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

110 if self.aud is not None:

111 d["aud"] = self.aud

112 d.update(self.extra)

113 return d

114

115 @classmethod

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

117 return cls(

118 sub=d["sub"],

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

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

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

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

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

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

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

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

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

128 )

129

130

131# ---------------------------------------------------------------------------

132# Abstract Providers

133# ---------------------------------------------------------------------------

134

135class TokenProvider(ABC):

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

137

138 @abstractmethod

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

140

141 @abstractmethod

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

143

144

145class CredentialStore(ABC):

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

147

148 @abstractmethod

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

150

151 @abstractmethod

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

153

154

155class TokenBlacklist(ABC):

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

157

158 @abstractmethod

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

160

161 @abstractmethod

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

163

164

165# ---------------------------------------------------------------------------

166# In-Memory Token Blacklist

167# ---------------------------------------------------------------------------

168

169class InMemoryTokenBlacklist(TokenBlacklist):

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

171

172 def __init__(self):

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

174

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

176 now = time.monotonic()

177 if jti in self._store:

178 if self._store[jti] > now:

179 return True

180 del self._store[jti]

181 return False

182

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

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

185

186 def _cleanup(self):

187 now = time.monotonic()

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

189

190

191# ---------------------------------------------------------------------------

192# In-Memory Credential Store

193# ---------------------------------------------------------------------------

194

195@dataclass

196class ApiKeyEntry:

197 key_hash: str

198 subject: str

199 name: str

200 scopes: List[str]

201 roles: List[str]

202 permissions: List[str]

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

204 expires_at: Optional[float] = None

205 revoked: bool = False

206

207

208class InMemoryCredentialStore(CredentialStore):

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

210

211 def __init__(self):

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

213

214 @staticmethod

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

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

217

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

219 entry.key_hash = self._hash_key(raw_key)

220 self._keys[entry.key_hash] = entry

221

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

223 key_hash = self._hash_key(api_key)

224 entry = self._keys.get(key_hash)

225 if entry is None or entry.revoked:

226 return None

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

228 return None

229 return {

230 "subject": entry.subject,

231 "name": entry.name,

232 "scopes": entry.scopes,

233 "roles": entry.roles,

234 "permissions": entry.permissions,

235 }

236

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

238 key_hash = self._hash_key(api_key)

239 entry = self._keys.get(key_hash)

240 if entry is not None:

241 entry.revoked = True

242 return True

243 return False

244

245

246# ---------------------------------------------------------------------------

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

248# ---------------------------------------------------------------------------

249

250class HS256TokenProvider(TokenProvider):

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

252

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

254 default_ttl: float = 3600.0):

255 self._current_secret = secret.encode()

256 self._previous_secret: Optional[bytes] = None

257 self._issuer = issuer

258 self._default_ttl = default_ttl

259

260 def rotate_secret(self, new_secret: str):

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

262 self._previous_secret = self._current_secret

263 self._current_secret = new_secret.encode()

264

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

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

267 segments = [

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

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

270 ]

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

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

273 segments.append(_b64url_encode(signature))

274 return ".".join(segments)

275

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

277 parts = token.split(".")

278 if len(parts) != 3:

279 return None

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

281 if secret is None:

282 continue

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

284 expected_sig = hmac.new(

285 secret, signing_input, hashlib.sha256

286 ).digest()

287 actual_sig = _b64url_decode(parts[2])

288 if hmac.compare_digest(expected_sig, actual_sig):

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

290 return TokenClaims.from_dict(payload)

291 return None

292

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

294 if claims.exp is None:

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

296 if claims.iss != self._issuer:

297 claims.iss = self._issuer

298 return self._encode(claims)

299

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

301 claims = self._decode(token)

302 if claims is None:

303 return None

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

305 return None

306 return claims

307

308

309# ---------------------------------------------------------------------------

310# RBAC Engine

311# ---------------------------------------------------------------------------

312

313class RBACEngine:

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

315

316 def __init__(self):

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

318

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

320 self._roles[role.name] = role

321

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

323 for role in roles:

324 self.register_role(role)

325

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

327 result: Set[Permission] = set()

328 for name in role_names:

329 role = self._roles.get(name)

330 if role is not None:

331 result.update(role.permissions)

332 return result

333

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

335 for perm in self.get_permissions(role_names):

336 if self._match(perm, required):

337 return True

338 return False

339

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

341 perms = self.get_permissions(role_names)

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

343

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

345 perms = self.get_permissions(role_names)

346 return all(

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

348 for r in required

349 )

350

351 @staticmethod

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

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

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

355 return True

356 if ":" in need:

357 # hierarchical: "org:engineering"

358 parts_need = need.split(":")

359 parts_have = have.split(":")

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

361 return False

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

363 return have == need

364 return (

365 _seg_match(perm.resource, required.resource)

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

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

368 )

369

370

371# ---------------------------------------------------------------------------

372# Authenticator

373# ---------------------------------------------------------------------------

374

375@dataclass

376class AuthenticatorConfig:

377 """Authenticator configuration."""

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

379 AuthMethod.JWT, AuthMethod.API_KEY, AuthMethod.OAUTH2_BEARER

380 )

381 default_issuer: str = "agentos"

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

383 api_key_query_param: str = "api_key"

384 require_auth: bool = True

385 token_strict_expiry: bool = True

386 clock_skew: float = 30.0 # seconds

387 max_token_lifetime: float = 86400.0 # 24h

388

389

390class Authenticator:

391 """Main authentication orchestrator.

392

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

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

395

396 Usage:

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

398 blacklist=bl, rbac=rbac)

399 ctx = await auth.authenticate(request_headers)

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

401 ...

402 """

403

404 def __init__(

405 self,

406 *,

407 token_provider: Optional[TokenProvider] = None,

408 credential_store: Optional[CredentialStore] = None,

409 blacklist: Optional[TokenBlacklist] = None,

410 rbac: Optional[RBACEngine] = None,

411 config: Optional[AuthenticatorConfig] = None,

412 ):

413 self._token_provider = token_provider

414 self._credential_store = credential_store

415 self._blacklist = blacklist

416 self._rbac = rbac or RBACEngine()

417 self._config = config or AuthenticatorConfig()

418

419 @property

420 def rbac(self) -> RBACEngine:

421 return self._rbac

422

423 async def authenticate(

424 self,

425 headers: Dict[str, str],

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

427 ) -> AuthContext:

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

429 Returns AuthContext with authenticated=False if auth fails.

430 """

431 # Try JWT Bearer

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

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

434 token = auth_header[7:]

435 ctx = await self._authenticate_jwt(token)

436 if ctx.authenticated:

437 return ctx

438

439 # Try OAuth2 Bearer

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

441 token = auth_header[7:]

442 ctx = await self._authenticate_oauth2_bearer(token)

443 if ctx.authenticated:

444 return ctx

445

446 # Try API Key

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

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

449 if not api_key and query_params:

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

451 if api_key:

452 ctx = await self._authenticate_api_key(api_key)

453 if ctx.authenticated:

454 return ctx

455

456 return AuthContext()

457

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

459 if self._token_provider is None:

460 return AuthContext()

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

462 if claims is None:

463 return AuthContext()

464 if self._blacklist and claims.jti:

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

466 return AuthContext()

467 expires_at = claims.exp

468 if expires_at and self._config.token_strict_expiry:

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

470 return AuthContext()

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

472 return AuthContext(

473 authenticated=True,

474 method=AuthMethod.JWT,

475 subject=claims.sub,

476 issuer=claims.iss,

477 roles=set(claims.roles),

478 permissions=perm_set,

479 metadata=claims.extra,

480 expires_at=expires_at,

481 token_id=claims.jti,

482 )

483

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

485 if self._credential_store is None:

486 return AuthContext()

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

488 if entry is None:

489 return AuthContext()

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

491 return AuthContext(

492 authenticated=True,

493 method=AuthMethod.API_KEY,

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

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

496 permissions=perm_set,

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

498 )

499

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

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

502 return AuthContext()

503

504 @staticmethod

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

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

507 if len(parts) == 1:

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

509 elif len(parts) == 2:

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

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

512

513

514# ---------------------------------------------------------------------------

515# Decorators

516# ---------------------------------------------------------------------------

517

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

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

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

521 """

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

523

524 def decorator(fn: Callable):

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

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

527 if ctx is None:

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

529 if not ctx.authenticated:

530 raise PermissionError("authentication required")

531 if required:

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

533 if rbac is None:

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

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

536 raise PermissionError(

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

538 )

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

540 wrapper.__name__ = fn.__name__

541 wrapper.__doc__ = fn.__doc__

542 return wrapper

543 return decorator

544

545

546# ---------------------------------------------------------------------------

547# Helpers

548# ---------------------------------------------------------------------------

549

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

551 import base64

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

553

554

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

556 import base64

557 padding = 4 - len(data) % 4

558 if padding != 4:

559 data += "=" * padding

560 return base64.urlsafe_b64decode(data)

561

562

563# ---------------------------------------------------------------------------

564# Default Roles

565# ---------------------------------------------------------------------------

566

567DEFAULT_ROLES = [

568 Role("admin", (

569 Permission("*", "*", "*"),

570 )),

571 Role("developer", (

572 Permission("agent", "*"),

573 Permission("model", "read"),

574 Permission("model", "execute"),

575 Permission("tool", "*"),

576 Permission("log", "read"),

577 )),

578 Role("viewer", (

579 Permission("agent", "read"),

580 Permission("model", "read"),

581 Permission("log", "read"),

582 Permission("metric", "read"),

583 )),

584 Role("operator", (

585 Permission("agent", "read"),

586 Permission("agent", "execute"),

587 Permission("model", "read"),

588 Permission("model", "execute"),

589 Permission("tool", "execute"),

590 Permission("log", "read"),

591 Permission("metric", "read"),

592 )),

593]