kite/backend/tests/test_auth_es256.py

"""
TDD scaffold for Phase 7.3: ES256 algorithm upgrade, startup token rotation,
and remember_me session TTL — all stubs xfail strict=True until promoted.
"""
import base64
import hashlib
import json
import time
import uuid
from datetime import datetime, timedelta, timezone

import pytest
import pytest_asyncio
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec


@pytest.fixture(autouse=True)
def es256_keys(monkeypatch):
    """Patch settings with a freshly generated P-256 key pair for each test."""
    k = ec.generate_private_key(ec.SECP256R1())
    priv = base64.b64encode(
        k.private_bytes(
            serialization.Encoding.PEM,
            serialization.PrivateFormat.PKCS8,
            serialization.NoEncryption(),
        )
    ).decode()
    pub = base64.b64encode(
        k.public_key().public_bytes(
            serialization.Encoding.PEM,
            serialization.PublicFormat.SubjectPublicKeyInfo,
        )
    ).decode()
    monkeypatch.setattr("config.settings.jwt_private_key", priv, raising=False)
    monkeypatch.setattr("config.settings.jwt_public_key", pub, raising=False)


# ── ES256 algorithm tests ─────────────────────────────────────────────────────

def test_access_token_uses_es256():
    """ES256-01: access token header must declare alg=ES256."""
    from services.auth import create_access_token
    token = create_access_token("u1", "user")
    # Decode the header (first segment of the JWT)
    header_b64 = token.split(".")[0]
    # Add padding to make it valid base64
    padding = "=" * (4 - len(header_b64) % 4)
    header = json.loads(base64.urlsafe_b64decode(header_b64 + padding))
    assert header["alg"] == "ES256"


def test_hs256_token_rejected():
    """ES256-02: an HS256 token must raise ValueError when decoded."""
    import jwt as _jwt
    from services.auth import decode_access_token
    hs256_token = _jwt.encode(
        {
            "sub": "u1",
            "typ": "access",
            "exp": int(time.time()) + 60,
            "iat": int(time.time()),
        },
        "any-hs256-secret",
        algorithm="HS256",
    )
    with pytest.raises(ValueError):
        decode_access_token(hs256_token)


def test_reset_token_uses_es256():
    """ES256-03: password-reset token header must declare alg=ES256 and round-trips."""
    from services.auth import create_password_reset_token, decode_password_reset_token
    token = create_password_reset_token("u1")
    header_b64 = token.split(".")[0]
    padding = "=" * (4 - len(header_b64) % 4)
    header = json.loads(base64.urlsafe_b64decode(header_b64 + padding))
    assert header["alg"] == "ES256"
    assert decode_password_reset_token(token) == "u1"


# ── Startup token rotation tests ──────────────────────────────────────────────

@pytest.mark.asyncio
async def test_startup_rotation_revokes_tokens(db_session):
    """ES256-04: on first run, all active refresh tokens are bulk-revoked and jwt_algorithm row is upserted."""
    import secrets as _secrets
    from main import _rotate_tokens_on_algorithm_change
    from db.models import RefreshToken, SystemSettings, User
    from sqlalchemy import select

    # Create a minimal user row (required by RefreshToken FK)
    user_id = uuid.uuid4()
    user = User(
        id=user_id,
        handle=f"testrotate_{user_id.hex[:8]}",
        email=f"testrotate_{user_id.hex[:8]}@example.com",
        password_hash="fakehash",
        role="user",
        is_active=True,
        password_must_change=False,
    )
    db_session.add(user)
    await db_session.flush()

    # Insert two active refresh tokens
    now = datetime.now(timezone.utc)
    tok1_id = uuid.uuid4()
    tok2_id = uuid.uuid4()
    tok1 = RefreshToken(
        id=tok1_id,
        user_id=user_id,
        token_hash=hashlib.sha256(_secrets.token_urlsafe(32).encode()).hexdigest(),
        expires_at=now + timedelta(days=1),
        revoked=False,
    )
    tok2 = RefreshToken(
        id=tok2_id,
        user_id=user_id,
        token_hash=hashlib.sha256(_secrets.token_urlsafe(32).encode()).hexdigest(),
        expires_at=now + timedelta(days=1),
        revoked=False,
    )
    db_session.add(tok1)
    db_session.add(tok2)
    await db_session.flush()

    # Act: no jwt_algorithm row exists yet
    await _rotate_tokens_on_algorithm_change(db_session)
    # expire_all() required: raw SQL UPDATE bypasses ORM identity map; expire_on_commit=False
    # (set in conftest.py) means SQLAlchemy won't auto-reload stale objects after commit.
    db_session.expire_all()

    # Assert: both tokens are now revoked
    result = await db_session.execute(
        select(RefreshToken).where(RefreshToken.id.in_([tok1_id, tok2_id]))
    )
    rows = result.scalars().all()
    assert len(rows) == 2
    assert all(r.revoked is True for r in rows)

    # Assert: jwt_algorithm row was created correctly
    ss_result = await db_session.execute(
        select(SystemSettings).where(SystemSettings.provider_id == "jwt_algorithm")
    )
    ss_row = ss_result.scalar_one_or_none()
    assert ss_row is not None
    assert ss_row.model_name == "ES256"
    assert ss_row.is_active is False
    assert ss_row.context_chars == 0


@pytest.mark.asyncio
async def test_startup_rotation_idempotent(db_session):
    """ES256-05: when jwt_algorithm row already has model_name=ES256, no tokens are revoked."""
    import secrets as _secrets
    from main import _rotate_tokens_on_algorithm_change
    from db.models import RefreshToken, SystemSettings, User
    from sqlalchemy import select, func

    # Create a user
    user_id = uuid.uuid4()
    user = User(
        id=user_id,
        handle=f"testidempotent_{user_id.hex[:8]}",
        email=f"testidempotent_{user_id.hex[:8]}@example.com",
        password_hash="fakehash",
        role="user",
        is_active=True,
        password_must_change=False,
    )
    db_session.add(user)
    await db_session.flush()

    # Seed the jwt_algorithm row as already migrated
    existing_marker = SystemSettings(
        id=uuid.uuid4(),
        provider_id="jwt_algorithm",
        model_name="ES256",
        context_chars=0,
        is_active=False,
        api_key_enc=None,
        base_url=None,
    )
    db_session.add(existing_marker)

    # Insert one fresh active token
    now = datetime.now(timezone.utc)
    tok_id = uuid.uuid4()
    tok = RefreshToken(
        id=tok_id,
        user_id=user_id,
        token_hash=hashlib.sha256(_secrets.token_urlsafe(32).encode()).hexdigest(),
        expires_at=now + timedelta(days=1),
        revoked=False,
    )
    db_session.add(tok)
    await db_session.flush()

    # Act
    await _rotate_tokens_on_algorithm_change(db_session)
    # expire_all() to clear identity map cache (same reason as test_startup_rotation_revokes_tokens)
    db_session.expire_all()

    # Assert: the fresh token was NOT revoked (bulk update did not fire)
    result = await db_session.execute(
        select(RefreshToken).where(RefreshToken.id == tok_id)
    )
    fresh_tok = result.scalar_one_or_none()
    assert fresh_tok is not None
    assert fresh_tok.revoked is False

    # Assert: still exactly one jwt_algorithm row (no duplicate upsert)
    count_result = await db_session.execute(
        select(func.count()).where(SystemSettings.provider_id == "jwt_algorithm")
    )
    count = count_result.scalar_one()
    assert count == 1


# ── Remember-me TTL stubs ─────────────────────────────────────────────────────

@pytest.mark.xfail(strict=True, reason="RM-01: not yet implemented")
@pytest.mark.asyncio
async def test_default_ttl_16_hours(async_client, db_session, auth_user):
    pytest.xfail("not yet implemented")


@pytest.mark.xfail(strict=True, reason="RM-02: not yet implemented")
@pytest.mark.asyncio
async def test_remember_me_ttl_30_days(async_client, db_session, auth_user):
    pytest.xfail("not yet implemented")


@pytest.mark.xfail(strict=True, reason="RM-03: not yet implemented")
@pytest.mark.asyncio
async def test_remember_me_cookie_max_age(async_client, auth_user):
    pytest.xfail("not yet implemented")


# ── Config field tests ────────────────────────────────────────────────────────

def test_settings_has_jwt_keys():
    """CFG-01: Settings must expose jwt_private_key, jwt_public_key, and refresh_token_expire_hours=16."""
    from config import settings
    assert hasattr(settings, "jwt_private_key")
    assert hasattr(settings, "jwt_public_key")
    assert hasattr(settings, "refresh_token_expire_hours")
    assert settings.refresh_token_expire_hours == 16