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.
"""
from __future__ import annotations

import base64
import hashlib
import json
import time
import uuid
from datetime import datetime, timedelta, timezone

import pytest
import pytest_asyncio


# ── ES256 key fixture ────────────────────────────────────────────────────────


@pytest.fixture(autouse=True)
def es256_keys(monkeypatch):
    """Generate a throw-away P-256 keypair and monkeypatch it into settings.

    raising=False: settings fields do not exist until Plan 02 adds them;
    this fixture must not error before that.
    """
    from cryptography.hazmat.primitives.asymmetric import ec
    from cryptography.hazmat.primitives import serialization

    private_key = ec.generate_private_key(ec.SECP256R1())
    private_pem = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )
    public_pem = private_key.public_key().public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )
    private_b64 = base64.b64encode(private_pem).decode()
    public_b64 = base64.b64encode(public_pem).decode()

    import config
    monkeypatch.setattr(config.settings, "jwt_private_key", private_b64, raising=False)
    monkeypatch.setattr(config.settings, "jwt_public_key", public_b64, raising=False)


# ── ES256-01: access token algorithm ─────────────────────────────────────────


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


# ── ES256-02: HS256 token rejected ───────────────────────────────────────────


def test_hs256_token_rejected():
    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)


# ── ES256-03: password-reset token algorithm ──────────────────────────────────


def test_reset_token_uses_es256():
    from services.auth import create_password_reset_token, decode_password_reset_token
    token = create_password_reset_token("u1")
    segment = token.split(".")[0]
    segment += "=" * (4 - len(segment) % 4)
    header = json.loads(base64.urlsafe_b64decode(segment))
    assert header["alg"] == "ES256"
    assert decode_password_reset_token(token) == "u1"


# ── ES256-04: startup rotation revokes tokens ─────────────────────────────────


@pytest.mark.asyncio
async def test_startup_rotation_revokes_tokens(db_session):
    from datetime import datetime, timezone, timedelta
    import hashlib
    import secrets
    import uuid as _uuid
    from sqlalchemy import select
    from db.models import RefreshToken, SystemSettings, User, Quota
    from main import _rotate_tokens_on_algorithm_change
    from services.auth import hash_password

    # Set up: create a user + two RefreshToken rows with revoked=False
    user_id = _uuid.uuid4()
    user = User(
        id=user_id,
        handle=f"rottest_{user_id.hex[:8]}",
        email=f"rottest_{user_id.hex[:8]}@example.com",
        password_hash=hash_password("Testpassword123!"),
        role="user",
        is_active=True,
        password_must_change=False,
    )
    quota = Quota(user_id=user_id, limit_bytes=104857600, used_bytes=0)
    db_session.add(user)
    db_session.add(quota)
    now = datetime.now(timezone.utc)
    rt1_id = _uuid.uuid4()
    rt2_id = _uuid.uuid4()
    rt1 = RefreshToken(
        id=rt1_id,
        user_id=user_id,
        token_hash=hashlib.sha256(secrets.token_urlsafe(16).encode()).hexdigest(),
        expires_at=now + timedelta(days=1),
        revoked=False,
    )
    rt2 = RefreshToken(
        id=rt2_id,
        user_id=user_id,
        token_hash=hashlib.sha256(secrets.token_urlsafe(16).encode()).hexdigest(),
        expires_at=now + timedelta(days=1),
        revoked=False,
    )
    db_session.add(rt1)
    db_session.add(rt2)
    await db_session.flush()
    # No jwt_algorithm row in system_settings

    # Act
    await _rotate_tokens_on_algorithm_change(db_session)
    # After the helper's commit, expire the identity map so we read fresh DB state
    db_session.expire_all()

    # Assert: both tokens revoked
    result = await db_session.execute(
        select(RefreshToken).where(RefreshToken.id.in_([rt1_id, rt2_id]))
    )
    rows = result.scalars().all()
    assert len(rows) == 2
    assert all(r.revoked for r in rows)

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


# ── ES256-05: startup rotation is idempotent ─────────────────────────────────


@pytest.mark.asyncio
async def test_startup_rotation_idempotent(db_session):
    import uuid as _uuid
    from sqlalchemy import select, func
    from db.models import RefreshToken, SystemSettings, User, Quota
    from main import _rotate_tokens_on_algorithm_change
    from services.auth import hash_password
    from datetime import datetime, timezone, timedelta
    import hashlib
    import secrets

    # Set up: pre-existing jwt_algorithm marker row with model_name='ES256'
    marker_id = _uuid.uuid4()
    marker = SystemSettings(
        id=marker_id,
        provider_id="jwt_algorithm",
        model_name="ES256",
        context_chars=0,
        is_active=False,
        api_key_enc=None,
        base_url=None,
    )
    db_session.add(marker)

    # Create a user + one fresh RefreshToken with revoked=False
    user_id = _uuid.uuid4()
    user = User(
        id=user_id,
        handle=f"idem_{user_id.hex[:8]}",
        email=f"idem_{user_id.hex[:8]}@example.com",
        password_hash=hash_password("Testpassword123!"),
        role="user",
        is_active=True,
        password_must_change=False,
    )
    quota = Quota(user_id=user_id, limit_bytes=104857600, used_bytes=0)
    db_session.add(user)
    db_session.add(quota)
    now = datetime.now(timezone.utc)
    rt_id = _uuid.uuid4()
    rt = RefreshToken(
        id=rt_id,
        user_id=user_id,
        token_hash=hashlib.sha256(secrets.token_urlsafe(16).encode()).hexdigest(),
        expires_at=now + timedelta(days=1),
        revoked=False,
    )
    db_session.add(rt)
    await db_session.flush()

    # Act: run rotation — should be no-op because model_name is already 'ES256'
    await _rotate_tokens_on_algorithm_change(db_session)

    # Assert: RefreshToken still not revoked
    result = await db_session.execute(
        select(RefreshToken).where(RefreshToken.id == rt_id)
    )
    token_row = result.scalar_one_or_none()
    assert token_row is not None
    assert not token_row.revoked

    # Assert: exactly one jwt_algorithm row (no duplicates)
    result2 = await db_session.execute(
        select(SystemSettings).where(SystemSettings.provider_id == "jwt_algorithm")
    )
    markers = result2.scalars().all()
    assert len(markers) == 1


# ── RM-01: default TTL is 16 hours ───────────────────────────────────────────


@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")


# ── RM-02: remember_me TTL is 30 days ────────────────────────────────────────


@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")


# ── RM-03: cookie Max-Age values ─────────────────────────────────────────────


@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")


# ── CFG-01 satellite: settings has jwt key fields ────────────────────────────


def test_settings_has_jwt_keys():
    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