210 lines
9.0 KiB
Python
210 lines
9.0 KiB
Python
import asyncio
|
|
from urllib.parse import quote, urlsplit
|
|
from uuid import uuid4
|
|
|
|
from fastapi.testclient import TestClient
|
|
from httpx import Response
|
|
from pwdlib import PasswordHash
|
|
from pydantic import SecretStr
|
|
from sqlalchemy import delete, select
|
|
|
|
from ai_platform.bootstrap.app import create_app
|
|
from ai_platform.infrastructure.config.settings import Settings
|
|
from ai_platform.infrastructure.persistence.database import make_session_factory
|
|
from ai_platform.infrastructure.persistence.models import (
|
|
AdminSession,
|
|
AdminUser,
|
|
TenantKey,
|
|
TenantKeyAuditEvent,
|
|
)
|
|
from ai_platform.platform.security import digest
|
|
|
|
|
|
def expect_status(response: Response, expected: int) -> None:
|
|
assert response.status_code == expected
|
|
|
|
|
|
def assert_secret_absent(secret: str, body: str) -> None:
|
|
if secret in body:
|
|
raise AssertionError("tenant key appeared in list response")
|
|
|
|
|
|
def test_admin_and_provisioner_tenant_key_lifecycle():
|
|
loaded = Settings.load()
|
|
database_name = urlsplit(loaded.database_url).path.rstrip("/").rsplit("/", 1)[-1]
|
|
if database_name != "ai_platform_test":
|
|
raise RuntimeError("Tenant key API tests require /ai_platform_test")
|
|
settings = Settings(
|
|
database_url=loaded.database_url,
|
|
environment="development",
|
|
tenant_provisioning_secret=SecretStr("test-only-provisioner-secret"),
|
|
)
|
|
sessions = make_session_factory(settings.database_url)
|
|
tenant_id, admin_tenant = f"tenant/{uuid4().hex}", f"tenant-admin-{uuid4().hex}"
|
|
tenant_route_id = quote(tenant_id, safe="")
|
|
tenant_key_path = f"/v1/admin/tenant-keys/{tenant_route_id}"
|
|
admin_email, password = f"tenant-admin-{uuid4().hex}@example.invalid", "test-only-password"
|
|
app = create_app(settings)
|
|
|
|
async def create_admin() -> None:
|
|
async with sessions() as db:
|
|
db.add(
|
|
AdminUser(
|
|
email=admin_email,
|
|
password_hash=PasswordHash.recommended().hash(password),
|
|
)
|
|
)
|
|
await db.commit()
|
|
await sessions.kw["bind"].dispose()
|
|
|
|
async def cleanup() -> None:
|
|
async with sessions() as db:
|
|
tenants = [tenant_id, admin_tenant]
|
|
await db.execute(
|
|
delete(TenantKeyAuditEvent).where(TenantKeyAuditEvent.tenant_id.in_(tenants))
|
|
)
|
|
await db.execute(delete(AdminSession).where(AdminSession.admin_email == admin_email))
|
|
await db.execute(delete(AdminUser).where(AdminUser.email == admin_email))
|
|
await db.execute(delete(TenantKey).where(TenantKey.tenant_id.in_(tenants)))
|
|
await db.commit()
|
|
await sessions.kw["bind"].dispose()
|
|
|
|
asyncio.run(create_admin())
|
|
try:
|
|
with TestClient(app) as client:
|
|
expect_status(client.get("/v1/admin/tenant-keys"), 401)
|
|
expect_status(client.post("/v1/admin/tenant-keys", json={"tenant_id": tenant_id}), 401)
|
|
provision_path = "/v1/platform/tenant-keys"
|
|
expect_status(
|
|
client.post(
|
|
provision_path,
|
|
headers={"Authorization": "Bearer incorrect"},
|
|
json={"tenant_id": tenant_id},
|
|
),
|
|
401,
|
|
)
|
|
provision_headers = {"Authorization": "Bearer test-only-provisioner-secret"}
|
|
expect_status(
|
|
client.post(provision_path, headers=provision_headers, json={"tenant_id": " "}),
|
|
422,
|
|
)
|
|
expect_status(
|
|
client.post(
|
|
provision_path,
|
|
headers=provision_headers,
|
|
json={"tenant_id": tenant_id, "unexpected": True},
|
|
),
|
|
422,
|
|
)
|
|
created = client.post(
|
|
provision_path, headers=provision_headers, json={"tenant_id": tenant_id}
|
|
)
|
|
expect_status(created, 201)
|
|
assert created.headers["cache-control"] == "no-store"
|
|
initial_key = created.json()["tenant_key"]
|
|
expect_status(
|
|
client.post(
|
|
provision_path, headers=provision_headers, json={"tenant_id": tenant_id}
|
|
),
|
|
409,
|
|
)
|
|
|
|
login = client.post(
|
|
"/v1/admin/login", json={"email": admin_email, "password": password}
|
|
)
|
|
expect_status(login, 200)
|
|
csrf = {"x-csrf-token": login.json()["csrf_token"]}
|
|
listed = client.get("/v1/admin/tenant-keys")
|
|
expect_status(listed, 200)
|
|
item = next(x for x in listed.json()["tenants"] if x["tenant_id"] == tenant_id)
|
|
assert set(item) == {"tenant_id", "active", "created_at", "updated_at"}
|
|
assert item["active"] and item["created_at"] and item["updated_at"]
|
|
assert_secret_absent(initial_key, listed.text)
|
|
assert "verifier" not in listed.text
|
|
expect_status(client.post("/v1/admin/tenant-keys", json={"tenant_id": "no-csrf"}), 403)
|
|
expect_status(client.post(f"{tenant_key_path}/rotate"), 403)
|
|
expect_status(client.delete(tenant_key_path), 403)
|
|
expect_status(
|
|
client.post("/v1/admin/tenant-keys", headers=csrf, json={"tenant_id": tenant_id}),
|
|
409,
|
|
)
|
|
|
|
run_url = f"/v1/lead-scoring/runs/{uuid4()}"
|
|
|
|
def tenant_auth(key: str) -> dict[str, str]:
|
|
return {"Authorization": f"Bearer {key}"}
|
|
|
|
expect_status(client.get(run_url, headers=tenant_auth(initial_key)), 404)
|
|
admin_created = client.post(
|
|
"/v1/admin/tenant-keys", headers=csrf, json={"tenant_id": admin_tenant}
|
|
)
|
|
expect_status(admin_created, 201)
|
|
assert admin_created.headers["cache-control"] == "no-store"
|
|
admin_key = admin_created.json()["tenant_key"]
|
|
assert_secret_absent(admin_key, client.get("/v1/admin/tenant-keys").text)
|
|
|
|
rotated = client.post(f"{tenant_key_path}/rotate", headers=csrf)
|
|
expect_status(rotated, 200)
|
|
assert rotated.headers["cache-control"] == "no-store"
|
|
rotated_key = rotated.json()["tenant_key"]
|
|
if rotated_key == initial_key:
|
|
raise AssertionError("rotation reused the previous key")
|
|
expect_status(client.get(run_url, headers=tenant_auth(initial_key)), 401)
|
|
expect_status(client.get(run_url, headers=tenant_auth(rotated_key)), 404)
|
|
assert_secret_absent(rotated_key, client.get("/v1/admin/tenant-keys").text)
|
|
|
|
revoked = client.delete(tenant_key_path, headers=csrf)
|
|
expect_status(revoked, 200)
|
|
assert revoked.json()["active"] is False
|
|
expect_status(client.get(run_url, headers=tenant_auth(rotated_key)), 401)
|
|
reactivated = client.post(f"{tenant_key_path}/rotate", headers=csrf)
|
|
expect_status(reactivated, 200)
|
|
assert reactivated.headers["cache-control"] == "no-store"
|
|
active_key = reactivated.json()["tenant_key"]
|
|
if active_key == rotated_key:
|
|
raise AssertionError("rotation reused a revoked key")
|
|
expect_status(client.get(run_url, headers=tenant_auth(active_key)), 404)
|
|
|
|
async def assert_audit_and_verifier() -> None:
|
|
try:
|
|
async with sessions() as db:
|
|
row = await db.get(TenantKey, tenant_id)
|
|
assert row is not None and row.active and row.verifier == digest(active_key)
|
|
events = (
|
|
await db.scalars(
|
|
select(TenantKeyAuditEvent)
|
|
.where(TenantKeyAuditEvent.tenant_id == tenant_id)
|
|
.order_by(TenantKeyAuditEvent.id)
|
|
)
|
|
).all()
|
|
assert [(e.actor, e.action) for e in events] == [
|
|
("provisioner", "create"),
|
|
(admin_email, "rotate"),
|
|
(admin_email, "revoke"),
|
|
(admin_email, "rotate"),
|
|
]
|
|
assert all(e.occurred_at is not None for e in events)
|
|
created_by_admin = await db.scalar(
|
|
select(TenantKeyAuditEvent).where(
|
|
TenantKeyAuditEvent.tenant_id == admin_tenant
|
|
)
|
|
)
|
|
assert created_by_admin is not None
|
|
assert (created_by_admin.actor, created_by_admin.action) == (
|
|
admin_email,
|
|
"create",
|
|
)
|
|
assert {c.name for c in TenantKeyAuditEvent.__table__.columns} == {
|
|
"id",
|
|
"tenant_id",
|
|
"actor",
|
|
"action",
|
|
"occurred_at",
|
|
}
|
|
finally:
|
|
await sessions.kw["bind"].dispose()
|
|
|
|
asyncio.run(assert_audit_and_verifier())
|
|
finally:
|
|
asyncio.run(cleanup())
|