stackpilot/backend/tests/test_token_revocation.py

"""Token revocation.

Before 0.46.0 a token was valid until it expired, full stop. Resetting a
compromised account's password changed nothing for whoever held its tokens —
up to 30 days for a refresh token — and demoting or disabling an account only
took effect once the same clock ran out.

Every token now carries the ``token_version`` it was minted with, and every
request compares that against the live user. Bumping the version is therefore
the revoke switch, and it is pulled on the three changes that alter what an
account may do: password, role, active flag.
"""
from __future__ import annotations

import pytest
from sqlmodel import Session, select


@pytest.fixture
def victim(db):
    """A throwaway account whose tokens the tests revoke."""
    import auth as auth_mod
    from database import engine
    from models.user import User

    with Session(engine) as session:
        existing = session.exec(select(User).where(User.username == "revoke-me")).first()
        if existing:
            session.delete(existing)
            session.commit()
        user = User(
            username="revoke-me",
            hashed_password=auth_mod.hash_password("original-password"),
            role="user",
        )
        session.add(user)
        session.commit()
        session.refresh(user)
        session.expunge(user)
    return user


def _token_for(username: str) -> str:
    import auth as auth_mod
    from database import engine
    from models.user import User

    with Session(engine) as session:
        user = session.exec(select(User).where(User.username == username)).one()
        return auth_mod.create_access_token(user)


def _still_works(client, token: str) -> bool:
    response = client.get("/api/auth/me", headers={"Authorization": f"Bearer {token}"})
    return response.status_code == 200


# --------------------------------------------------------------------------- #
# The claim itself
# --------------------------------------------------------------------------- #


def test_tokens_carry_the_version_they_were_minted_with(victim):
    import auth as auth_mod

    payload = auth_mod.decode_token(auth_mod.create_access_token(victim), "access")
    assert payload["ver"] == 1


def test_a_null_version_from_an_older_schema_reads_as_one(victim):
    """Tolerated so a half-migrated database does not lock everyone out."""
    import auth as auth_mod

    victim.token_version = None
    assert auth_mod.token_version_of(victim) == 1


# --------------------------------------------------------------------------- #
# What revokes
# --------------------------------------------------------------------------- #


@pytest.mark.parametrize(
    "change",
    [
        pytest.param({"password": "brand-new-password"}, id="password-reset"),
        pytest.param({"role": "admin"}, id="role-change"),
        pytest.param({"is_active": False}, id="account-disabled"),
    ],
)
def test_changing_what_an_account_may_do_kills_its_tokens(
    as_admin, client, victim, change
):
    token = _token_for("revoke-me")
    assert _still_works(client, token)

    assert as_admin.request(
        "PATCH", f"/api/auth/users/{victim.id}", json=change
    ).status_code == 200

    assert not _still_works(client, token), f"{change} left the old token usable"


def test_a_cosmetic_update_does_not_sign_the_user_out(as_admin, client, victim):
    """Re-saving the same role must not invalidate a working session."""
    token = _token_for("revoke-me")
    assert as_admin.request(
        "PATCH", f"/api/auth/users/{victim.id}", json={"role": victim.role}
    ).status_code == 200
    assert _still_works(client, token)


def test_a_fresh_token_works_after_a_revoke(as_admin, client, victim):
    old = _token_for("revoke-me")
    as_admin.request(
        "PATCH", f"/api/auth/users/{victim.id}", json={"password": "another-one"}
    )
    assert not _still_works(client, old)
    assert _still_works(client, _token_for("revoke-me"))


def test_refresh_rejects_a_revoked_token(as_admin, client, victim):
    """The long-lived token is the one that mattered — 30 days of access."""
    import auth as auth_mod
    from database import engine
    from models.user import User

    with Session(engine) as session:
        user = session.exec(select(User).where(User.username == "revoke-me")).one()
        refresh = auth_mod.create_refresh_token(user)

    as_admin.request(
        "PATCH", f"/api/auth/users/{victim.id}", json={"password": "changed-again"}
    )
    response = client.post("/api/auth/refresh", json={"refresh_token": refresh})
    assert response.status_code == 401


def test_logout_everywhere_revokes_the_callers_own_tokens(client, victim):
    import auth as auth_mod
    from database import engine
    from models.user import User

    with Session(engine) as session:
        user = session.exec(select(User).where(User.username == "revoke-me")).one()
        token = auth_mod.create_access_token(user)

    assert client.post(
        "/api/auth/logout-everywhere", headers={"Authorization": f"Bearer {token}"}
    ).status_code == 200
    assert not _still_works(client, token)


# --------------------------------------------------------------------------- #
# The refresh cookie
# --------------------------------------------------------------------------- #


def test_login_puts_the_refresh_token_in_an_httponly_cookie(client, victim):
    response = client.post(
        "/api/auth/login", json={"username": "revoke-me", "password": "original-password"}
    )
    assert response.status_code == 200

    # Not in the body — that is the whole point.
    assert response.json().get("refresh_token") is None
    assert response.json()["access_token"]

    cookie = response.headers["set-cookie"]
    assert "stackpilot_refresh=" in cookie
    assert "HttpOnly" in cookie
    assert "Path=/api/auth" in cookie
    assert "SameSite=lax" in cookie.lower().replace("samesite=lax", "SameSite=lax")


def test_refresh_works_off_the_cookie_alone(client, victim):
    client.post(
        "/api/auth/login", json={"username": "revoke-me", "password": "original-password"}
    )
    # TestClient keeps the cookie jar, so no body is sent here.
    response = client.post("/api/auth/refresh", json={})
    assert response.status_code == 200
    assert response.json()["access_token"]


def test_scripted_clients_can_still_ask_for_the_token_in_the_body(client, victim):
    response = client.post(
        "/api/auth/login?in_body=true",
        json={"username": "revoke-me", "password": "original-password"},
    )
    assert response.json()["refresh_token"]


def test_logout_clears_the_cookie(client, victim):
    client.post(
        "/api/auth/login", json={"username": "revoke-me", "password": "original-password"}
    )
    response = client.post("/api/auth/logout")
    assert response.status_code == 200
    assert 'stackpilot_refresh=""' in response.headers["set-cookie"] or (
        "stackpilot_refresh=;" in response.headers["set-cookie"]
    )


def test_the_cookie_is_not_marked_secure_over_plain_http(client, victim):
    """A homelab on plain HTTP must still be able to stay signed in.

    Marking the cookie Secure unconditionally would make the browser drop it and
    nobody could hold a session. Over HTTPS the flag is set — see _issue().
    """
    response = client.post(
        "/api/auth/login", json={"username": "revoke-me", "password": "original-password"}
    )
    assert "secure" not in response.headers["set-cookie"].lower()


# --------------------------------------------------------------------------- #
# WebSockets go through the same check (F9)
# --------------------------------------------------------------------------- #


def test_websocket_auth_resolves_against_the_database(victim):
    """A socket must not trust the role baked into the token.

    ``/ws/exec`` is root-equivalent on the host, and a token minted while the
    account was an admin stays syntactically valid after a demotion — so the
    role has to come from the database, not the claim.
    """
    import auth as auth_mod
    from database import engine
    from models.user import User
    from routers.ws import _user_for

    with Session(engine) as session:
        user = session.exec(select(User).where(User.username == "revoke-me")).one()
        user.role = "admin"
        auth_mod.bump_token_version(user)
        session.add(user)
        session.commit()
        session.refresh(user)
        admin_token = auth_mod.create_access_token(user)

    assert _user_for(admin_token).role == "admin"

    # Demote. The token still decodes, but must no longer resolve.
    with Session(engine) as session:
        user = session.exec(select(User).where(User.username == "revoke-me")).one()
        user.role = "user"
        auth_mod.bump_token_version(user)
        session.add(user)
        session.commit()

    assert auth_mod.decode_token(admin_token, "access")["role"] == "admin"
    assert _user_for(admin_token) is None


def test_websocket_auth_rejects_a_disabled_account(victim):
    import auth as auth_mod
    from database import engine
    from models.user import User
    from routers.ws import _user_for

    with Session(engine) as session:
        user = session.exec(select(User).where(User.username == "revoke-me")).one()
        token = auth_mod.create_access_token(user)
        assert _user_for(token) is not None
        user.is_active = False
        session.add(user)
        session.commit()

    assert _user_for(token) is None


def test_websocket_auth_rejects_garbage(victim):
    from routers.ws import _user_for

    assert _user_for(None) is None
    assert _user_for("") is None
    assert _user_for("not-a-jwt") is None