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 time

22import uuid

23from abc import ABC, abstractmethod

24from collections.abc import Callable

25from dataclasses import dataclass, field

26from enum import Enum

27from typing import Any

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

30# Constants & Enums

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

34class AuthMethod(Enum):

35 NONE = "none"

36 JWT = "jwt"

37 API_KEY = "api_key"

38 OAUTH2_BEARER = "oauth2_bearer"

39 CUSTOM = "custom"

42class Algorithm(Enum):

43 HS256 = "HS256"

44 HS384 = "HS384"

45 HS512 = "HS512"

46 RS256 = "RS256"

47 RS384 = "RS384"

48 RS512 = "RS512"

49 ES256 = "ES256"

50 ES384 = "ES384"

51 ES512 = "ES512"

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

55class Permission:

56 """Fine-grained permission atom."""

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

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

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

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

64class Role:

65 """Named collection of permissions."""

67 name: str

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

71@dataclass

72class AuthContext:

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

75 authenticated: bool = False

76 method: AuthMethod = AuthMethod.NONE

77 subject: str | None = None # user ID / API key ID

78 issuer: str | None = None

79 roles: set[str] = field(default_factory=set)

80 permissions: set[Permission] = field(default_factory=set)

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

82 expires_at: float | None = None

83 token_id: str | None = None # jti

86# ---------------------------------------------------------------------------

87# Token Model

88# ---------------------------------------------------------------------------

91@dataclass

92class TokenClaims:

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

95 sub: str

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

97 exp: float | None = None

98 iss: str = "agentos"

99 aud: str | list[str] | None = None

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

101 roles: list[str] = field(default_factory=list)

102 permissions: list[str] = field(default_factory=list)

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

104

105 def to_dict(self) -> dict[str, Any]:

106 d = {

107 "sub": self.sub,

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

109 "iss": self.iss,

110 "jti": self.jti,

111 "roles": self.roles,

112 "permissions": self.permissions,

113 }

114 if self.exp is not None:

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

116 if self.aud is not None:

117 d["aud"] = self.aud

118 d.update(self.extra)

119 return d

120

121 @classmethod

122 def from_dict(cls, d: dict[str, Any]) -> TokenClaims:

123 return cls(

124 sub=d["sub"],

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

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

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

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

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

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

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

132 extra={

133 k: v

134 for k, v in d.items()

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

136 },

137 )

138

139

140# ---------------------------------------------------------------------------

141# Abstract Providers

142# ---------------------------------------------------------------------------

143

144

145class TokenProvider(ABC):

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

147

148 @abstractmethod

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

150

151 @abstractmethod

152 async def validate_token(self, token: str) -> TokenClaims | None: ...

153

154

155class CredentialStore(ABC):

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

157

158 @abstractmethod

159 async def lookup_by_key(self, api_key: str) -> dict[str, Any] | None: ...

160

161 @abstractmethod

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

163

164

165class TokenBlacklist(ABC):

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

167

168 @abstractmethod

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

170

171 @abstractmethod

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

173

174

175# ---------------------------------------------------------------------------

176# In-Memory Token Blacklist

177# ---------------------------------------------------------------------------

178

179

180class InMemoryTokenBlacklist(TokenBlacklist):

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

182

183 def __init__(self):

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

185

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

187 now = time.monotonic()

188 if jti in self._store:

189 if self._store[jti] > now:

190 return True

191 del self._store[jti]

192 return False

193

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

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

196

197 def _cleanup(self):

198 now = time.monotonic()

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

200

201

202# ---------------------------------------------------------------------------

203# In-Memory Credential Store

204# ---------------------------------------------------------------------------

205

206

207@dataclass

208class ApiKeyEntry:

209 key_hash: str

210 subject: str

211 name: str

212 scopes: list[str]

213 roles: list[str]

214 permissions: list[str]

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

216 expires_at: float | None = None

217 revoked: bool = False

218

219

220class InMemoryCredentialStore(CredentialStore):

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

222

223 def __init__(self):

224 self._keys: dict[str, ApiKeyEntry] = {}

225

226 @staticmethod

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

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

229

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

231 entry.key_hash = self._hash_key(raw_key)

232 self._keys[entry.key_hash] = entry

233

234 async def lookup_by_key(self, api_key: str) -> dict[str, Any] | None:

235 key_hash = self._hash_key(api_key)

236 entry = self._keys.get(key_hash)

237 if entry is None or entry.revoked:

238 return None

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

240 return None

241 return {

242 "subject": entry.subject,

243 "name": entry.name,

244 "scopes": entry.scopes,

245 "roles": entry.roles,

246 "permissions": entry.permissions,

247 }

248

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

250 key_hash = self._hash_key(api_key)

251 entry = self._keys.get(key_hash)

252 if entry is not None:

253 entry.revoked = True

254 return True

255 return False

256

257

258# ---------------------------------------------------------------------------

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

260# ---------------------------------------------------------------------------

261

262

263class HS256TokenProvider(TokenProvider):

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

265

266 def __init__(self, secret: str, issuer: str = "agentos", default_ttl: float = 3600.0):

267 self._current_secret = secret.encode()

268 self._previous_secret: bytes | None = None

269 self._issuer = issuer

270 self._default_ttl = default_ttl

271

272 def rotate_secret(self, new_secret: str):

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

274 self._previous_secret = self._current_secret

275 self._current_secret = new_secret.encode()

276

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

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

279 segments = [

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

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

282 ]

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

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

285 segments.append(_b64url_encode(signature))

286 return ".".join(segments)

287

288 def _decode(self, token: str) -> TokenClaims | None:

289 parts = token.split(".")

290 if len(parts) != 3:

291 return None

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

293 if secret is None:

294 continue

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

296 expected_sig = hmac.new(secret, signing_input, hashlib.sha256).digest()

297 actual_sig = _b64url_decode(parts[2])

298 if hmac.compare_digest(expected_sig, actual_sig):

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

300 return TokenClaims.from_dict(payload)

301 return None

302

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

304 if claims.exp is None:

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

306 if claims.iss != self._issuer:

307 claims.iss = self._issuer

308 return self._encode(claims)

309

310 async def validate_token(self, token: str) -> TokenClaims | None:

311 claims = self._decode(token)

312 if claims is None:

313 return None

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

315 return None

316 return claims

317

318

319# ---------------------------------------------------------------------------

320# RBAC Engine

321# ---------------------------------------------------------------------------

322

323

324class RBACEngine:

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

326

327 def __init__(self):

328 self._roles: dict[str, Role] = {}

329

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

331 self._roles[role.name] = role

332

333 def register_roles(self, roles: list[Role]) -> None:

334 for role in roles:

335 self.register_role(role)

336

337 def get_permissions(self, role_names: set[str]) -> set[Permission]:

338 result: set[Permission] = set()

339 for name in role_names:

340 role = self._roles.get(name)

341 if role is not None:

342 result.update(role.permissions)

343 return result

344

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

346 for perm in self.get_permissions(role_names):

347 if self._match(perm, required):

348 return True

349 return False

350

351 def check_any(self, role_names: set[str], required: list[Permission]) -> bool:

352 perms = self.get_permissions(role_names)

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

354

355 def check_all(self, role_names: set[str], required: list[Permission]) -> bool:

356 perms = self.get_permissions(role_names)

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

358

359 @staticmethod

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

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

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

363 return True

364 if ":" in need:

365 # hierarchical: "org:engineering"

366 parts_need = need.split(":")

367 parts_have = have.split(":")

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

369 return False

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

371 return have == need

372

373 return (

374 _seg_match(perm.resource, required.resource)

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

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

377 )

378

379

380# ---------------------------------------------------------------------------

381# Authenticator

382# ---------------------------------------------------------------------------

383

384

385@dataclass

386class AuthenticatorConfig:

387 """Authenticator configuration."""

388

389 allowed_methods: tuple[AuthMethod, ...] = (

390 AuthMethod.JWT,

391 AuthMethod.API_KEY,

392 AuthMethod.OAUTH2_BEARER,

393 )

394 default_issuer: str = "agentos"

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

396 api_key_query_param: str = "api_key"

397 require_auth: bool = True

398 token_strict_expiry: bool = True

399 clock_skew: float = 30.0 # seconds

400 max_token_lifetime: float = 86400.0 # 24h

401

402

403class Authenticator:

404 """Main authentication orchestrator.

405

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

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

408

409 Usage:

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

411 blacklist=bl, rbac=rbac)

412 ctx = await auth.authenticate(request_headers)

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

414 ...

415 """

416

417 def __init__(

418 self,

419 *,

420 token_provider: TokenProvider | None = None,

421 credential_store: CredentialStore | None = None,

422 blacklist: TokenBlacklist | None = None,

423 rbac: RBACEngine | None = None,

424 config: AuthenticatorConfig | None = None,

425 ):

426 self._token_provider = token_provider

427 self._credential_store = credential_store

428 self._blacklist = blacklist

429 self._rbac = rbac or RBACEngine()

430 self._config = config or AuthenticatorConfig()

431

432 @property

433 def rbac(self) -> RBACEngine:

434 return self._rbac

435

436 async def authenticate(

437 self,

438 headers: dict[str, str],

439 query_params: dict[str, str] | None = None,

440 ) -> AuthContext:

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

442 Returns AuthContext with authenticated=False if auth fails.

443 """

444 # Try JWT Bearer

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

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

447 token = auth_header[7:]

448 ctx = await self._authenticate_jwt(token)

449 if ctx.authenticated:

450 return ctx

451

452 # Try OAuth2 Bearer

453 if (

454 auth_header.startswith("Bearer ")

455 and AuthMethod.OAUTH2_BEARER in self._config.allowed_methods

456 ):

457 token = auth_header[7:]

458 ctx = await self._authenticate_oauth2_bearer(token)

459 if ctx.authenticated:

460 return ctx

461

462 # Try API Key

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

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

465 if not api_key and query_params:

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

467 if api_key:

468 ctx = await self._authenticate_api_key(api_key)

469 if ctx.authenticated:

470 return ctx

471

472 return AuthContext()

473

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

475 if self._token_provider is None:

476 return AuthContext()

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

478 if claims is None:

479 return AuthContext()

480 if self._blacklist and claims.jti:

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

482 return AuthContext()

483 expires_at = claims.exp

484 if expires_at and self._config.token_strict_expiry:

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

486 return AuthContext()

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

488 return AuthContext(

489 authenticated=True,

490 method=AuthMethod.JWT,

491 subject=claims.sub,

492 issuer=claims.iss,

493 roles=set(claims.roles),

494 permissions=perm_set,

495 metadata=claims.extra,

496 expires_at=expires_at,

497 token_id=claims.jti,

498 )

499

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

501 if self._credential_store is None:

502 return AuthContext()

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

504 if entry is None:

505 return AuthContext()

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

507 return AuthContext(

508 authenticated=True,

509 method=AuthMethod.API_KEY,

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

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

512 permissions=perm_set,

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

514 )

515

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

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

518 return AuthContext()

519

520 @staticmethod

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

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

523 if len(parts) == 1:

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

525 elif len(parts) == 2:

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

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

528

529

530# ---------------------------------------------------------------------------

531# Decorators

532# ---------------------------------------------------------------------------

533

534

535def require_auth(permission: Permission | None = None, permissions: list[Permission] | None = None):

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

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

538 """

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

540

541 def decorator(fn: Callable):

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

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

544 if ctx is None:

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

546 if not ctx.authenticated:

547 raise PermissionError("authentication required")

548 if required:

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

550 if rbac is None:

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

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

553 raise PermissionError(

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

555 )

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

557

558 wrapper.__name__ = fn.__name__

559 wrapper.__doc__ = fn.__doc__

560 return wrapper

561

562 return decorator

563

564

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

566# Helpers

567# ---------------------------------------------------------------------------

568

569

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

571 import base64

572

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

574

575

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

577 import base64

578

579 padding = 4 - len(data) % 4

580 if padding != 4:

581 data += "=" * padding

582 return base64.urlsafe_b64decode(data)

583

584

585# ---------------------------------------------------------------------------

586# Default Roles

587# ---------------------------------------------------------------------------

588

589DEFAULT_ROLES = [

590 Role("admin", (Permission("*", "*", "*"),)),

591 Role(

592 "developer",

593 (

594 Permission("agent", "*"),

595 Permission("model", "read"),

596 Permission("model", "execute"),

597 Permission("tool", "*"),

598 Permission("log", "read"),

599 ),

600 ),

601 Role(

602 "viewer",

603 (

604 Permission("agent", "read"),

605 Permission("model", "read"),

606 Permission("log", "read"),

607 Permission("metric", "read"),

608 ),

609 ),

610 Role(

611 "operator",

612 (

613 Permission("agent", "read"),

614 Permission("agent", "execute"),

615 Permission("model", "read"),

616 Permission("model", "execute"),

617 Permission("tool", "execute"),

618 Permission("log", "read"),

619 Permission("metric", "read"),

620 ),

621 ),

622]