lambdagent/tests/test_authenticate.py

"""
Unit tests for `agentpaas.api.middleware.auth.authenticate()`.

audit/AUDIT_2026-06-05.md #34 baseline: the entire API-key login flow
— including the legacy-hash auto-migration that writes to
`api_keys.key_hash` — had ZERO direct test coverage. A typo in the
hash format check or the tenant-status branch would have shipped to
production with green CI.

Tests cover:
  - missing key (empty / None)
  - "Bearer " prefix stripping
  - PBKDF2 happy path (returns TenantContext, updates last_used_at)
  - rejection of invalid key with no matching prefix
  - rejection of invalid key with matching prefix but wrong hash
  - SUSPENDED tenant → 403
  - EXPIRED key → 401
  - Legacy SHA-256 hash auto-migration → key still works AND the stored
    hash is rewritten to PBKDF2 format (this is the audit's headline
    SEC-01 risk that left dormant legacy hashes plaintext forever).
"""
from __future__ import annotations

import os
os.environ.setdefault("AGENTPAAS_DATABASE_URL", "sqlite:///:memory:")

import json
import secrets

import pytest

from agentpaas.api.middleware.auth import (
    authenticate,
    AuthError,
    TenantContext,
    hash_key,
    _legacy_hash_key,
)
from agentpaas.db.models import Database, gen_id, now_utc
import agentpaas.db.session as _session_mod


# ─────────────────────────────────────────────────────────────────────────────
# Fixtures
# ─────────────────────────────────────────────────────────────────────────────

@pytest.fixture
def db():
    """Per-test in-memory SQLite. Same singleton-swap trick as
    test_api_endpoints — authenticate() calls get_db() internally."""
    prev = _session_mod._db
    _session_mod._db = Database("sqlite:///:memory:")
    try:
        yield _session_mod._db
    finally:
        _session_mod._db = prev


def _insert_tenant(db, name: str = "test-tenant", status: str = "active") -> str:
    tid = gen_id("tn_")
    db.execute(
        "INSERT INTO tenants (id, name, plan, status, created_at) "
        "VALUES (?, ?, 'free', ?, ?)",
        (tid, name, status, now_utc()),
    )
    db.commit()
    return tid


def _insert_user(db, tenant_id: str) -> str:
    uid = gen_id("usr_")
    db.execute(
        "INSERT INTO users (id, tenant_id, email, role, created_at) "
        "VALUES (?, ?, '', 'admin', ?)",
        (uid, tenant_id, now_utc()),
    )
    db.commit()
    return uid


def _insert_key(db, tenant_id: str, user_id: str,
                key_hash: str, scopes: list = None,
                status: str = "active",
                expires_at: str = "",
                prefix_override: str = "") -> tuple[str, str]:
    """Insert an api_keys row. Returns (key_id, key_prefix)."""
    key_id = gen_id("key_")
    raw = f"ap_{secrets.token_hex(16)}"
    db.execute(
        "INSERT INTO api_keys "
        "(id, tenant_id, user_id, key_hash, key_prefix, name, scopes, "
        " rate_limit, status, created_at, expires_at) "
        "VALUES (?, ?, ?, ?, ?, 'test-key', ?, 600, ?, ?, ?)",
        (key_id, tenant_id, user_id, key_hash,
         prefix_override or raw[:8],
         json.dumps(scopes or ["agents:*"]),
         status, now_utc(), expires_at),
    )
    db.commit()
    return key_id, raw


# ─────────────────────────────────────────────────────────────────────────────
# Missing / malformed input
# ─────────────────────────────────────────────────────────────────────────────

@pytest.mark.asyncio
async def test_authenticate_empty_key_raises_missing(db):
    with pytest.raises(AuthError) as exc:
        await authenticate("")
    assert exc.value.code == 401
    assert "missing" in exc.value.message.lower()


@pytest.mark.asyncio
async def test_authenticate_strips_bearer_prefix(db):
    """A real client sends `Authorization: Bearer ap_xxx`; FastAPI passes
    the whole header string. `authenticate()` must strip the prefix."""
    tid = _insert_tenant(db)
    uid = _insert_user(db, tid)
    raw_key = f"ap_{secrets.token_hex(16)}"
    _insert_key(db, tid, uid, hash_key(raw_key), prefix_override=raw_key[:8])

    # Without prefix
    ctx = await authenticate(raw_key)
    assert isinstance(ctx, TenantContext)
    assert ctx.tenant_id == tid

    # With prefix — must produce the same result
    ctx2 = await authenticate(f"Bearer {raw_key}")
    assert ctx2.tenant_id == tid


# ─────────────────────────────────────────────────────────────────────────────
# Rejection paths
# ─────────────────────────────────────────────────────────────────────────────

@pytest.mark.asyncio
async def test_authenticate_unknown_key_raises_401(db):
    with pytest.raises(AuthError) as exc:
        await authenticate("ap_unknown000000000000000000000000")
    assert exc.value.code == 401
    assert "invalid" in exc.value.message.lower()


@pytest.mark.asyncio
async def test_authenticate_matching_prefix_wrong_hash_raises_401(db):
    """Tenant A has ap_aaa.., tenant B uses ap_aaa.. too (same prefix
    by coincidence). B's key shouldn't unlock A's tenant just because
    the prefix lookup pulls A's row out — the verify_key_hash call must
    detect the mismatch and fall through. This is the silent-bypass
    smell SEC-01 originally flagged."""
    tid = _insert_tenant(db)
    uid = _insert_user(db, tid)
    real_key = "ap_aaaa11112222333344445555aaaaaaaa"
    _insert_key(db, tid, uid, hash_key(real_key), prefix_override="ap_aaaa")

    forged_key = "ap_aaaadeadbeefdeadbeefdeadbeefdead"  # same prefix
    with pytest.raises(AuthError) as exc:
        await authenticate(forged_key)
    assert exc.value.code == 401


# ─────────────────────────────────────────────────────────────────────────────
# Happy path: PBKDF2 key → TenantContext
# ─────────────────────────────────────────────────────────────────────────────

@pytest.mark.asyncio
async def test_authenticate_pbkdf2_happy_path(db):
    tid = _insert_tenant(db, name="acme")
    uid = _insert_user(db, tid)
    raw_key = f"ap_{secrets.token_hex(16)}"
    scopes = ["agents:read", "agents:write", "keys:*"]
    _insert_key(db, tid, uid, hash_key(raw_key), scopes=scopes,
                prefix_override=raw_key[:8])

    ctx = await authenticate(raw_key)
    assert ctx.tenant_id == tid
    assert ctx.tenant_name == "acme"
    assert ctx.plan == "free"
    assert ctx.user_id == uid
    # Scopes survived the round trip through JSON.
    assert set(ctx.scopes) == set(scopes)


# ─────────────────────────────────────────────────────────────────────────────
# Tenant / key status
# ─────────────────────────────────────────────────────────────────────────────

@pytest.mark.asyncio
async def test_authenticate_suspended_tenant_raises_403(db):
    tid = _insert_tenant(db, status="suspended")
    uid = _insert_user(db, tid)
    raw_key = f"ap_{secrets.token_hex(16)}"
    _insert_key(db, tid, uid, hash_key(raw_key), prefix_override=raw_key[:8])

    with pytest.raises(AuthError) as exc:
        await authenticate(raw_key)
    assert exc.value.code == 403
    assert "suspended" in exc.value.message.lower()


@pytest.mark.asyncio
async def test_authenticate_expired_key_raises_401(db):
    tid = _insert_tenant(db)
    uid = _insert_user(db, tid)
    raw_key = f"ap_{secrets.token_hex(16)}"
    # Expired at 2020-01-01 — well in the past regardless of test clock.
    _insert_key(db, tid, uid, hash_key(raw_key),
                expires_at="2020-01-01T00:00:00+00:00",
                prefix_override=raw_key[:8])

    with pytest.raises(AuthError) as exc:
        await authenticate(raw_key)
    assert exc.value.code == 401
    assert "expired" in exc.value.message.lower()


# ─────────────────────────────────────────────────────────────────────────────
# audit SEC-01 headline: legacy hash auto-migration
# ─────────────────────────────────────────────────────────────────────────────

@pytest.mark.asyncio
async def test_authenticate_legacy_hash_auto_migrates_to_pbkdf2(db):
    """The single most important authenticate() branch:
    a key stored under the legacy fixed-salt SHA-256 hash must (1) still
    authenticate and (2) be transparently rewritten to PBKDF2 on success.

    Without this test the SEC-01 'partial fix' could silently regress —
    a future refactor of verify_key_hash could break the legacy fallback
    and CI would still be green."""
    tid = _insert_tenant(db)
    uid = _insert_user(db, tid)
    raw_key = f"ap_{secrets.token_hex(16)}"

    legacy_hash = _legacy_hash_key(raw_key)
    key_id, _ = _insert_key(db, tid, uid, legacy_hash,
                            prefix_override=raw_key[:8])

    # Sanity: directly inserting a SHA-256 hex (not pbkdf2-prefixed) is
    # what we wanted.
    pre = db.fetchone("SELECT key_hash FROM api_keys WHERE id = ?", (key_id,))
    assert not pre["key_hash"].startswith("pbkdf2:")

    ctx = await authenticate(raw_key)
    assert ctx.tenant_id == tid

    # Auto-migration: hash was rewritten in-place.
    post = db.fetchone("SELECT key_hash FROM api_keys WHERE id = ?", (key_id,))
    assert post["key_hash"].startswith("pbkdf2:"), (
        f"legacy hash was not auto-migrated; still {post['key_hash'][:20]}"
    )

    # Migrated row should also still authenticate on the next call (the
    # entire point of the upgrade).
    ctx2 = await authenticate(raw_key)
    assert ctx2.tenant_id == tid