kiba-core 0.5.3.dev47__tar.gz → 0.5.3.dev49__tar.gz

{kiba_core-0.5.3.dev47 → kiba_core-0.5.3.dev49}/Dockerfile

@@ -10,4 +10,3 @@ COPY uv.lock .
 RUN make install
 
 COPY . .
-RUN make type-check

{kiba_core-0.5.3.dev47 → kiba_core-0.5.3.dev49}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kiba-core
-Version: 0.5.3.dev47
+Version: 0.5.3.dev49
 Summary: Kiba Labs' python utilities
 Project-URL: repository, https://github.com/kibalabs/core-py
 Requires-Python: <4,>=3.12

kiba_core-0.5.3.dev49/core/api/authorizer.py

@@ -0,0 +1,134 @@
+import functools
+import typing
+from collections.abc import AsyncIterator
+from typing import ParamSpec
+
+from mypy_extensions import Arg
+from pydantic import BaseModel
+
+from core import logging
+from core.api.api_request import KibaApiRequest
+from core.exceptions import ForbiddenException
+from core.exceptions import UnauthorizedException
+from core.http.basic_authentication import BasicAuthentication
+from core.http.jwt import Jwt
+
+_P = ParamSpec('_P')
+_AnyReturn = typing.Awaitable[typing.Any] | AsyncIterator[typing.Any]  # type: ignore[explicit-any]
+
+
+class Authorizer:
+    async def validate_jwt(self, jwtString: str) -> Jwt:
+        raise NotImplementedError
+
+
+class SignatureAuthorizer:
+    async def retrieve_signature_signer(self, signatureString: str) -> str:
+        raise NotImplementedError
+
+
+async def _authorize_bearer_jwt[ApiRequest: BaseModel](request: KibaApiRequest[ApiRequest], authorizer: Authorizer) -> Jwt:
+    authorization = request.headers.get('Authorization')
+    if not authorization:
+        raise ForbiddenException(message='AUTH_NOT_PROVIDED')
+    if not authorization.startswith('Bearer '):
+        raise ForbiddenException(message='AUTH_INVALID')
+    jwtString = authorization.replace('Bearer ', '')
+    try:
+        jwt = await authorizer.validate_jwt(jwtString=jwtString)
+    except BaseException:  # noqa: BLE001
+        raise ForbiddenException(message='AUTH_INVALID')
+    return jwt
+
+
+def authorize_bearer_jwt[ApiRequest: BaseModel](  # type: ignore[explicit-any]
+    authorizer: Authorizer,
+) -> typing.Callable[[typing.Callable[[Arg(KibaApiRequest[ApiRequest], 'request')], _AnyReturn]], typing.Callable[_P, typing.Any]]:
+    def decorator(func: typing.Callable[[Arg(KibaApiRequest[ApiRequest], 'request')], _AnyReturn]) -> typing.Callable[_P, typing.Any]:  # type: ignore[explicit-any]
+        @functools.wraps(func)
+        async def async_wrapper(request: KibaApiRequest[ApiRequest]) -> typing.Any:  # type: ignore[explicit-any, misc]
+            request.authJwt = await _authorize_bearer_jwt(request=request, authorizer=authorizer)
+            result = func(request=request)
+            # NOTE(krishan711): this is here to support streaming responses which return an async generator
+            if hasattr(result, '__aiter__'):
+                return result
+            return await result
+
+        # TODO(krishan711): figure out correct typing here
+        return async_wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+async def get_basic_authentication_from_authorization_signature[ApiRequest: BaseModel](request: KibaApiRequest[ApiRequest], authorizer: SignatureAuthorizer) -> BasicAuthentication:
+    authorization = request.headers.get('Authorization')
+    if not authorization:
+        raise ForbiddenException(message='AUTH_NOT_PROVIDED')
+    if not authorization.startswith('Signature '):
+        raise ForbiddenException(message='AUTH_INVALID')
+    signatureString = authorization.replace('Signature ', '', 1)
+    try:
+        signerId = await authorizer.retrieve_signature_signer(signatureString=signatureString)
+    except UnauthorizedException:
+        raise
+    except BaseException as exception:  # noqa: BLE001
+        logging.exception(exception)  # type: ignore[arg-type]
+        raise ForbiddenException(message='AUTH_INVALID')
+    return BasicAuthentication(username=signerId, password=signatureString)
+
+
+def authorize_signature[ApiRequest: BaseModel](  # type: ignore[explicit-any]
+    authorizer: SignatureAuthorizer,
+) -> typing.Callable[[typing.Callable[[Arg(KibaApiRequest[ApiRequest], 'request')], _AnyReturn]], typing.Callable[_P, typing.Any]]:
+    def decorator(func: typing.Callable[[Arg(KibaApiRequest[ApiRequest], 'request')], _AnyReturn]) -> typing.Callable[_P, typing.Any]:  # type: ignore[explicit-any]
+        @functools.wraps(func)
+        async def async_wrapper(request: KibaApiRequest[ApiRequest]) -> typing.Any:  # type: ignore[explicit-any, misc]
+            request.authBasic = await get_basic_authentication_from_authorization_signature(request=request, authorizer=authorizer)
+            result = func(request=request)
+            # NOTE(krishan711): this is here to support streaming responses which return an async generator
+            if hasattr(result, '__aiter__'):
+                return result
+            return await result
+
+        # TODO(krishan711): figure out correct typing here
+        return async_wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+class TokenAuthorizer:
+    async def validate_token(self, token: str) -> None:
+        raise NotImplementedError
+
+
+class StaticTokenAuthorizer(TokenAuthorizer):
+    def __init__(self, token: str) -> None:
+        self._token = token
+
+    async def validate_token(self, token: str) -> None:
+        if token != self._token:
+            raise ForbiddenException(message='AUTH_INVALID')
+
+
+def authorize_token[ApiRequest: BaseModel](  # type: ignore[explicit-any]
+    authorizer: TokenAuthorizer,
+) -> typing.Callable[[typing.Callable[[Arg(KibaApiRequest[ApiRequest], 'request')], _AnyReturn]], typing.Callable[_P, typing.Any]]:
+    def decorator(func: typing.Callable[[Arg(KibaApiRequest[ApiRequest], 'request')], _AnyReturn]) -> typing.Callable[_P, typing.Any]:  # type: ignore[explicit-any]
+        @functools.wraps(func)
+        async def async_wrapper(request: KibaApiRequest[ApiRequest]) -> typing.Any:  # type: ignore[explicit-any, misc]
+            authorization = request.headers.get('Authorization')
+            if not authorization:
+                raise ForbiddenException(message='AUTH_NOT_PROVIDED')
+            if not authorization.startswith('Token '):
+                raise ForbiddenException(message='AUTH_INVALID')
+            await authorizer.validate_token(authorization[6:])
+            result = func(request=request)
+            # NOTE(krishan711): this is here to support streaming responses which return an async generator
+            if hasattr(result, '__aiter__'):
+                return result
+            return await result
+
+        # TODO(krishan711): figure out correct typing here
+        return async_wrapper  # type: ignore[return-value]
+
+    return decorator

kiba_core-0.5.3.dev49/core/api/middleware/database_connection_middleware.py

@@ -0,0 +1,21 @@
+from starlette.types import ASGIApp
+from starlette.types import Receive
+from starlette.types import Scope
+from starlette.types import Send
+
+from core.store.database import Database
+
+
+class DatabaseConnectionMiddleware:
+    def __init__(self, app: ASGIApp, database: Database) -> None:
+        self.app = app
+        self.database = database
+
+    # NOTE(krishan711): see note in database.py about why this can cause problems with concurrent operations
+    # NOTE(krishan711): raw ASGI (not BaseHTTPMiddleware) so the DB connection stays open across streaming body
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope['type'] != 'http':
+            await self.app(scope, receive, send)
+            return
+        async with self.database.create_context_connection():
+            await self.app(scope, receive, send)

{kiba_core-0.5.3.dev47 → kiba_core-0.5.3.dev49}/kiba_core.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kiba-core
-Version: 0.5.3.dev47
+Version: 0.5.3.dev49
 Summary: Kiba Labs' python utilities
 Project-URL: repository, https://github.com/kibalabs/core-py
 Requires-Python: <4,>=3.12

{kiba_core-0.5.3.dev47 → kiba_core-0.5.3.dev49}/kiba_core.egg-info/SOURCES.txt

@@ -76,6 +76,7 @@ kiba_core.egg-info/dependency_links.txt
 kiba_core.egg-info/requires.txt
 kiba_core.egg-info/top_level.txt
 tests/__init__.py
+tests/api/test_authorizer.py
 tests/api/test_json_route.py
 tests/api/test_streaming_json_route.py
 tests/caching/test_dict_cache.py

{kiba_core-0.5.3.dev47 → kiba_core-0.5.3.dev49}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kiba-core"
-version = "0.5.3.dev47"
+version = "0.5.3.dev49"
 description = "Kiba Labs' python utilities"
 readme = "README.md"
 requires-python = ">=3.12, <4"

kiba_core-0.5.3.dev49/tests/api/test_authorizer.py

@@ -0,0 +1,265 @@
+import json
+import pytest
+from collections.abc import AsyncIterator
+from pydantic import BaseModel
+from starlette.applications import Starlette
+from starlette.routing import Route
+from starlette.testclient import TestClient
+
+from core.api.api_request import KibaApiRequest
+from core.api.authorizer import Authorizer, SignatureAuthorizer, StaticTokenAuthorizer, authorize_bearer_jwt, authorize_signature, authorize_token
+from core.api.json_route import json_route
+from core.api.middleware.exception_handling_middleware import ExceptionHandlingMiddleware
+from core.api.streaming_json_route import streaming_json_route
+from core.exceptions import ForbiddenException
+from core.http.jwt import Jwt
+
+
+VALID_JWT_TOKEN = 'valid-token'
+VALID_SIGNATURE = 'valid-sig'
+VALID_STATIC_TOKEN = 'secret-token'
+VALID_USER_ID = 'user-123'
+
+
+class MockJwtAuthorizer(Authorizer):
+    async def validate_jwt(self, jwtString: str) -> Jwt:
+        if jwtString != VALID_JWT_TOKEN:
+            raise ForbiddenException('Invalid token')
+        jwt = Jwt(payloadDict={'sub': VALID_USER_ID})
+        jwt.userId = VALID_USER_ID  # type: ignore[attr-defined]
+        return jwt
+
+
+class MockSignatureAuthorizer(SignatureAuthorizer):
+    async def retrieve_signature_signer(self, signatureString: str) -> str:
+        if signatureString != VALID_SIGNATURE:
+            raise ForbiddenException('Invalid signature')
+        return VALID_USER_ID
+
+
+class SimpleRequest(BaseModel):
+    value: str
+
+
+class SimpleResponse(BaseModel):
+    result: str
+    user_id: str | None = None
+
+
+jwt_authorizer = MockJwtAuthorizer()
+sig_authorizer = MockSignatureAuthorizer()
+
+
+# --- authorize_bearer_jwt fixtures ---
+
+@pytest.fixture
+def jwt_json_client():
+    @json_route(requestType=SimpleRequest, responseType=SimpleResponse)
+    @authorize_bearer_jwt(authorizer=jwt_authorizer)
+    async def protected_endpoint(request: KibaApiRequest[SimpleRequest]) -> SimpleResponse:
+        jwt = request.authJwt
+        return SimpleResponse(result=request.data.value, user_id=getattr(jwt, 'userId', None))
+
+    app = Starlette(routes=[Route('/protected', protected_endpoint, methods=['POST'])])
+    app.add_middleware(ExceptionHandlingMiddleware)
+    return TestClient(app, raise_server_exceptions=False)
+
+
+@pytest.fixture
+def jwt_streaming_client():
+    @streaming_json_route(requestType=SimpleRequest, responseType=SimpleResponse)
+    @authorize_bearer_jwt(authorizer=jwt_authorizer)
+    async def protected_streaming_endpoint(request: KibaApiRequest[SimpleRequest]) -> AsyncIterator[SimpleResponse]:
+        jwt = request.authJwt
+        yield SimpleResponse(result=request.data.value, user_id=getattr(jwt, 'userId', None))
+
+    app = Starlette(routes=[Route('/protected-stream', protected_streaming_endpoint, methods=['POST'])])
+    app.add_middleware(ExceptionHandlingMiddleware)
+    return TestClient(app, raise_server_exceptions=False)
+
+
+# --- authorize_signature fixtures ---
+
+@pytest.fixture
+def sig_json_client():
+    @json_route(requestType=SimpleRequest, responseType=SimpleResponse)
+    @authorize_signature(authorizer=sig_authorizer)
+    async def protected_endpoint(request: KibaApiRequest[SimpleRequest]) -> SimpleResponse:
+        return SimpleResponse(result=request.data.value, user_id=request.authBasic.username if request.authBasic else None)
+
+    app = Starlette(routes=[Route('/protected', protected_endpoint, methods=['POST'])])
+    app.add_middleware(ExceptionHandlingMiddleware)
+    return TestClient(app, raise_server_exceptions=False)
+
+
+@pytest.fixture
+def sig_streaming_client():
+    @streaming_json_route(requestType=SimpleRequest, responseType=SimpleResponse)
+    @authorize_signature(authorizer=sig_authorizer)
+    async def protected_streaming_endpoint(request: KibaApiRequest[SimpleRequest]) -> AsyncIterator[SimpleResponse]:
+        yield SimpleResponse(result=request.data.value, user_id=request.authBasic.username if request.authBasic else None)
+
+    app = Starlette(routes=[Route('/protected-stream', protected_streaming_endpoint, methods=['POST'])])
+    app.add_middleware(ExceptionHandlingMiddleware)
+    return TestClient(app, raise_server_exceptions=False)
+
+
+# --- authorize_token fixtures ---
+
+token_authorizer = StaticTokenAuthorizer(token=VALID_STATIC_TOKEN)
+
+
+@pytest.fixture
+def token_json_client():
+    @json_route(requestType=SimpleRequest, responseType=SimpleResponse)
+    @authorize_token(authorizer=token_authorizer)
+    async def protected_endpoint(request: KibaApiRequest[SimpleRequest]) -> SimpleResponse:
+        return SimpleResponse(result=request.data.value)
+
+    app = Starlette(routes=[Route('/protected', protected_endpoint, methods=['POST'])])
+    app.add_middleware(ExceptionHandlingMiddleware)
+    return TestClient(app, raise_server_exceptions=False)
+
+
+@pytest.fixture
+def token_streaming_client():
+    @streaming_json_route(requestType=SimpleRequest, responseType=SimpleResponse)
+    @authorize_token(authorizer=token_authorizer)
+    async def protected_streaming_endpoint(request: KibaApiRequest[SimpleRequest]) -> AsyncIterator[SimpleResponse]:
+        yield SimpleResponse(result=request.data.value)
+
+    app = Starlette(routes=[Route('/protected-stream', protected_streaming_endpoint, methods=['POST'])])
+    app.add_middleware(ExceptionHandlingMiddleware)
+    return TestClient(app, raise_server_exceptions=False)
+
+
+# --- authorize_bearer_jwt + json_route ---
+
+def test_jwt_json_no_auth_header_returns_403(jwt_json_client):
+    response = jwt_json_client.post('/protected', json={'value': 'hello'})
+    assert response.status_code == 403
+
+def test_jwt_json_wrong_scheme_returns_403(jwt_json_client):
+    response = jwt_json_client.post('/protected', json={'value': 'hello'}, headers={'Authorization': 'Basic creds'})
+    assert response.status_code == 403
+
+def test_jwt_json_invalid_token_returns_403(jwt_json_client):
+    response = jwt_json_client.post('/protected', json={'value': 'hello'}, headers={'Authorization': 'Bearer bad-token'})
+    assert response.status_code == 403
+
+def test_jwt_json_valid_token_returns_200(jwt_json_client):
+    response = jwt_json_client.post('/protected', json={'value': 'hello'}, headers={'Authorization': f'Bearer {VALID_JWT_TOKEN}'})
+    assert response.status_code == 200
+    assert response.json()['result'] == 'hello'
+
+def test_jwt_json_sets_auth_jwt_on_request(jwt_json_client):
+    response = jwt_json_client.post('/protected', json={'value': 'hello'}, headers={'Authorization': f'Bearer {VALID_JWT_TOKEN}'})
+    assert response.status_code == 200
+    assert response.json()['user_id'] == VALID_USER_ID
+
+
+# --- authorize_bearer_jwt + streaming_json_route ---
+
+def test_jwt_streaming_no_auth_header_returns_403(jwt_streaming_client):
+    response = jwt_streaming_client.post('/protected-stream', json={'value': 'hello'})
+    assert response.status_code == 403
+
+def test_jwt_streaming_invalid_token_returns_403(jwt_streaming_client):
+    response = jwt_streaming_client.post('/protected-stream', json={'value': 'hello'}, headers={'Authorization': 'Bearer bad-token'})
+    assert response.status_code == 403
+
+def test_jwt_streaming_valid_token_returns_200(jwt_streaming_client):
+    response = jwt_streaming_client.post('/protected-stream', json={'value': 'hello'}, headers={'Authorization': f'Bearer {VALID_JWT_TOKEN}'})
+    assert response.status_code == 200
+
+def test_jwt_streaming_valid_token_streams_data(jwt_streaming_client):
+    response = jwt_streaming_client.post('/protected-stream', json={'value': 'hello'}, headers={'Authorization': f'Bearer {VALID_JWT_TOKEN}'})
+    data = json.loads(response.content.decode().strip())
+    assert data['result'] == 'hello'
+    assert data['user_id'] == VALID_USER_ID
+
+
+# --- authorize_signature + json_route ---
+
+def test_sig_json_no_auth_header_returns_403(sig_json_client):
+    response = sig_json_client.post('/protected', json={'value': 'hello'})
+    assert response.status_code == 403
+
+def test_sig_json_wrong_scheme_returns_403(sig_json_client):
+    response = sig_json_client.post('/protected', json={'value': 'hello'}, headers={'Authorization': 'Bearer something'})
+    assert response.status_code == 403
+
+def test_sig_json_invalid_signature_returns_403(sig_json_client):
+    response = sig_json_client.post('/protected', json={'value': 'hello'}, headers={'Authorization': 'Signature bad-sig'})
+    assert response.status_code == 403
+
+def test_sig_json_valid_signature_returns_200(sig_json_client):
+    response = sig_json_client.post('/protected', json={'value': 'hello'}, headers={'Authorization': f'Signature {VALID_SIGNATURE}'})
+    assert response.status_code == 200
+    assert response.json()['result'] == 'hello'
+
+def test_sig_json_sets_auth_basic_on_request(sig_json_client):
+    response = sig_json_client.post('/protected', json={'value': 'hello'}, headers={'Authorization': f'Signature {VALID_SIGNATURE}'})
+    assert response.status_code == 200
+    assert response.json()['user_id'] == VALID_USER_ID
+
+
+# --- authorize_signature + streaming_json_route ---
+
+def test_sig_streaming_no_auth_header_returns_403(sig_streaming_client):
+    response = sig_streaming_client.post('/protected-stream', json={'value': 'hello'})
+    assert response.status_code == 403
+
+def test_sig_streaming_invalid_signature_returns_403(sig_streaming_client):
+    response = sig_streaming_client.post('/protected-stream', json={'value': 'hello'}, headers={'Authorization': 'Signature bad-sig'})
+    assert response.status_code == 403
+
+def test_sig_streaming_valid_signature_returns_200(sig_streaming_client):
+    response = sig_streaming_client.post('/protected-stream', json={'value': 'hello'}, headers={'Authorization': f'Signature {VALID_SIGNATURE}'})
+    assert response.status_code == 200
+
+def test_sig_streaming_valid_signature_streams_data(sig_streaming_client):
+    response = sig_streaming_client.post('/protected-stream', json={'value': 'hello'}, headers={'Authorization': f'Signature {VALID_SIGNATURE}'})
+    data = json.loads(response.content.decode().strip())
+    assert data['result'] == 'hello'
+    assert data['user_id'] == VALID_USER_ID
+
+
+# --- authorize_token + json_route ---
+
+def test_token_json_no_auth_header_returns_403(token_json_client):
+    response = token_json_client.post('/protected', json={'value': 'hello'})
+    assert response.status_code == 403
+
+def test_token_json_wrong_scheme_returns_403(token_json_client):
+    response = token_json_client.post('/protected', json={'value': 'hello'}, headers={'Authorization': 'Bearer something'})
+    assert response.status_code == 403
+
+def test_token_json_invalid_token_returns_403(token_json_client):
+    response = token_json_client.post('/protected', json={'value': 'hello'}, headers={'Authorization': 'Token wrong-token'})
+    assert response.status_code == 403
+
+def test_token_json_valid_token_returns_200(token_json_client):
+    response = token_json_client.post('/protected', json={'value': 'hello'}, headers={'Authorization': f'Token {VALID_STATIC_TOKEN}'})
+    assert response.status_code == 200
+    assert response.json()['result'] == 'hello'
+
+
+# --- authorize_token + streaming_json_route ---
+
+def test_token_streaming_no_auth_header_returns_403(token_streaming_client):
+    response = token_streaming_client.post('/protected-stream', json={'value': 'hello'})
+    assert response.status_code == 403
+
+def test_token_streaming_invalid_token_returns_403(token_streaming_client):
+    response = token_streaming_client.post('/protected-stream', json={'value': 'hello'}, headers={'Authorization': 'Token wrong-token'})
+    assert response.status_code == 403
+
+def test_token_streaming_valid_token_returns_200(token_streaming_client):
+    response = token_streaming_client.post('/protected-stream', json={'value': 'hello'}, headers={'Authorization': f'Token {VALID_STATIC_TOKEN}'})
+    assert response.status_code == 200
+
+def test_token_streaming_valid_token_streams_data(token_streaming_client):
+    response = token_streaming_client.post('/protected-stream', json={'value': 'hello'}, headers={'Authorization': f'Token {VALID_STATIC_TOKEN}'})
+    data = json.loads(response.content.decode().strip())
+    assert data['result'] == 'hello'