diff --git a/pyproject.toml b/pyproject.toml index 83a1e5b..ef834ba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "uv_build" [project] name = "rain-api-core" -version = "0.1.0" +version = "0.1.1" description = "RAIN API Core" readme = "README.md" authors = [ @@ -12,6 +12,7 @@ authors = [ ] requires-python = "~=3.10" dependencies = [ + "boto3~=1.35", "cachetools~=5.0", "jinja2~=3.0", "netaddr~=1.0", @@ -21,7 +22,6 @@ dependencies = [ [dependency-groups] dev = [ - "boto3~=1.35", "hypothesis~=6.112", "moto~=5.0", "pytest~=8.3", diff --git a/src/rain_api_core/auth.py b/src/rain_api_core/auth.py index eeb06fd..74417dd 100644 --- a/src/rain_api_core/auth.py +++ b/src/rain_api_core/auth.py @@ -1,5 +1,7 @@ import contextlib import dataclasses +import hashlib +import json import logging from collections.abc import Mapping from http.cookies import CookieError, SimpleCookie @@ -8,6 +10,7 @@ from wsgiref.handlers import format_date_time as format_7231_date import jwt +from boto3 import client as botoclient log = logging.getLogger(__name__) @@ -56,7 +59,7 @@ def __init__( public_key: str, private_key: str, cookie_name: str, - blacklist={}, + blacklist: Mapping[str, int] = {}, session_ttl_in_hours: float = 24, ): self.algorithm = algorithm @@ -159,3 +162,96 @@ def get_header_to_set_auth_cookie( "HttpOnly; SameSite=Lax" ) } + + +class SessionManager: + def __init__( + self, + table_name: str, + column_names: Mapping[str, str] = { + "session_id": "SessionId", + "iat": "CreatedAt", + "ttl": "ExpiresAt", + }, + ): + self.client = botoclient("dynamodb") + self.table_name = table_name + self._session_id_column = column_names["session_id"] + self._iat_column = column_names["iat"] + self._ttl_column = column_names["ttl"] + + def revoke_session(self, user_profile: UserProfile): + session_id = self._get_session_id(user_profile) + exp = user_profile.exp + + assert exp is not None + + log.info( + "Revoking session for user %s with session_id: %s", + user_profile.user_id, + session_id, + ) + self.client.put_item( + TableName=self.table_name, + Item={ + self._session_id_column: { + "S": session_id, + }, + self._iat_column: { + "N": str(int(time())), + }, + self._ttl_column: { + "N": str(exp), + }, + }, + ) + + def is_session_revoked(self, user_profile: UserProfile) -> bool: + session_id = self._get_session_id(user_profile) + + iat = user_profile.iat + + assert iat is not None + + log.debug( + "Checking for revoked session for user %s with session_id: %s", + user_profile.user_id, + session_id, + ) + response = self.client.query( + TableName=self.table_name, + KeyConditionExpression=f"{self._session_id_column} = :sessionId", + FilterExpression=f"{self._ttl_column} > :now and {self._iat_column} >= :iat", + ExpressionAttributeValues={ + ":sessionId": { + "S": session_id, + }, + ":iat": { + "N": str(iat), + }, + ":now": { + "N": str(int(time())), + }, + }, + ) + + is_revoked = session_id in ( + # ruff hint + item["SessionId"]["S"] + for item in response.get("Items", ()) + ) + if is_revoked: + log.info("Session for user %s has been revoked", user_profile.user_id) + + return is_revoked + + def _get_session_id(self, user_profile: UserProfile) -> str: + # Depending on how we compute the session ID here we can determine the + # scope of the log out. If we use the entire JWT cookie payload, then + # logging out means only invalidating that cookie. If the same user has + # other cookies active as well then those will remain active. If instead + # we just use the user ID, then all cookies associated with that user + # will be invalidated. + return hashlib.sha256( + json.dumps(user_profile.to_jwt_payload()).encode(), + ).hexdigest() diff --git a/src/rain_api_core/urs_util.py b/src/rain_api_core/urs_util.py index 0ad9064..c2810ce 100644 --- a/src/rain_api_core/urs_util.py +++ b/src/rain_api_core/urs_util.py @@ -234,11 +234,11 @@ def user_in_group_urs( def user_in_group( - private_groups, - user_profile: UserProfile, - refresh_first=False, - aux_headers=None, -): + private_groups: list, + user_profile: Optional[UserProfile], + refresh_first: bool = False, + aux_headers: Optional[dict] = None, +) -> tuple[bool, Optional[UserProfile]]: aux_headers = aux_headers or {} # A safer default # If a new profile is fetched, it is assigned to this var, and returned so that a fresh jwt cookie can be set. diff --git a/tests/test_auth.py b/tests/test_auth.py index 1ccdfd5..b663ccc 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -1,14 +1,24 @@ +import dataclasses +from time import time from unittest import mock +import boto3 import jwt import jwt.utils import pytest +from moto import mock_aws -from rain_api_core.auth import JwtManager, UserProfile +from rain_api_core.auth import JwtManager, SessionManager, UserProfile MODULE = "rain_api_core.auth" +@pytest.fixture +def client_dynamodb(): + with mock_aws(): + yield boto3.client("dynamodb") + + @pytest.fixture def jwt_manager(jwt_priv_key, jwt_pub_key): return JwtManager( @@ -19,6 +29,36 @@ def jwt_manager(jwt_priv_key, jwt_pub_key): ) +@pytest.fixture +def revoked_session_table(client_dynamodb): + table_name = "UnitTestRevokedSessions" + client_dynamodb.create_table( + TableName=table_name, + AttributeDefinitions=[ + { + "AttributeName": "SessionId", + "AttributeType": "S", + }, + ], + KeySchema=[ + { + "AttributeName": "SessionId", + "KeyType": "HASH", + }, + ], + BillingMode="PAY_PER_REQUEST", + ) + client_dynamodb.update_time_to_live( + TableName=table_name, + TimeToLiveSpecification={ + "Enabled": True, + "AttributeName": "ExpiresAt", + }, + ) + + return boto3.resource("dynamodb").Table(table_name) + + def test_decode_jwt(jwt_manager, jwt_priv_key): payload = {"foo": "bar"} encoded = jwt.encode(payload, jwt_priv_key, "RS256") @@ -224,3 +264,173 @@ def test_jwt_manager_session_ttl_sub_hour(mock_time): "iat": 1, "exp": 1801, } + + +def test_session_manager(revoked_session_table): + session_manager = SessionManager(revoked_session_table.name) + + profile = UserProfile( + user_id="test_user", + token="test_token", + groups=["test_group1", "test_group2"], + first_name="test_first_name", + last_name="test_last_name", + email="test_email", + iat=int(time()), + exp=int(time()) + 1000, + ) + + assert session_manager.is_session_revoked(profile) is False + session_manager.revoke_session(profile) + assert session_manager.is_session_revoked(profile) is True + + +def test_session_manager_revoke_session(revoked_session_table): + session_manager = SessionManager(revoked_session_table.name) + + profile = UserProfile( + user_id="test_user", + token="test_token", + groups=["test_group1", "test_group2"], + first_name="test_first_name", + last_name="test_last_name", + email="test_email", + iat=0, + exp=1000, + ) + + session_manager.revoke_session(profile) + items = revoked_session_table.scan()["Items"] + + session_id = items[0]["SessionId"] + assert items == [ + { + "SessionId": session_id, + "CreatedAt": mock.ANY, + "ExpiresAt": 1000, + }, + ] + + # Calling it twice does not create new items. The creation time may change. + session_manager.revoke_session(profile) + items = revoked_session_table.scan()["Items"] + assert items == [ + { + "SessionId": session_id, + "CreatedAt": mock.ANY, + "ExpiresAt": 1000, + }, + ] + + +def test_session_manager_is_session_revoked(revoked_session_table): + session_manager = SessionManager(revoked_session_table.name) + + profile = UserProfile( + user_id="test_user", + token="test_token", + groups=["test_group1", "test_group2"], + first_name="test_first_name", + last_name="test_last_name", + email="test_email", + iat=int(time()) - 1000, + exp=int(time()) + 1000, + ) + + session_id = session_manager._get_session_id(profile) + + revoked_session_table.put_item( + Item={ + "SessionId": session_id, + "CreatedAt": int(time()), + "ExpiresAt": int(time()) + 1000, + }, + ) + assert session_manager.is_session_revoked(profile) is True + + # Created at time is before token minting time + revoked_session_table.put_item( + Item={ + "SessionId": session_id, + "CreatedAt": int(time()) - 2000, + "ExpiresAt": int(time()) + 1000, + }, + ) + assert session_manager.is_session_revoked(profile) is False + + # Revocation has expired + revoked_session_table.put_item( + Item={ + "SessionId": session_id, + "CreatedAt": int(time()), + "ExpiresAt": int(time()) - 1000, + }, + ) + assert session_manager.is_session_revoked(profile) is False + + +def test_session_manager_get_session_id(revoked_session_table): + session_manager = SessionManager(revoked_session_table.name) + + profile = UserProfile( + user_id="test_user", + token="test_token", + groups=["test_group1", "test_group2"], + first_name="test_first_name", + last_name="test_last_name", + email="test_email", + iat=0, + exp=1000, + ) + + original_session_id = session_manager._get_session_id(profile) + + # Changing an attribute changes the ID + assert ( + session_manager._get_session_id( + dataclasses.replace(profile, user_id="replaced"), + ) + != original_session_id + ) + assert ( + session_manager._get_session_id( + dataclasses.replace(profile, token="replaced"), + ) + != original_session_id + ) + assert ( + session_manager._get_session_id( + dataclasses.replace(profile, groups=["replaced"]), + ) + != original_session_id + ) + assert ( + session_manager._get_session_id( + dataclasses.replace(profile, first_name="replaced"), + ) + != original_session_id + ) + assert ( + session_manager._get_session_id( + dataclasses.replace(profile, last_name="replaced"), + ) + != original_session_id + ) + assert ( + session_manager._get_session_id( + dataclasses.replace(profile, email="replaced"), + ) + != original_session_id + ) + assert ( + session_manager._get_session_id( + dataclasses.replace(profile, iat=50), + ) + != original_session_id + ) + assert ( + session_manager._get_session_id( + dataclasses.replace(profile, exp=5000), + ) + != original_session_id + )