reworked api definition

src/testdata/testdata.py

@@ -1,10 +1,14 @@
 import os
 import json
 import asyncio
+import inspect
+import functools
+import random
 
-from typing_extensions import Annotated
 import uvicorn
-from fastapi import FastAPI, Request, status, HTTPException
+from typing_extensions import Annotated
+from fastapi import FastAPI, Request, Security, status, HTTPException
+from fastapi.security import APIKeyHeader, APIKeyQuery
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, ConfigDict, Field, BeforeValidator, ValidationError
 
@@ -54,72 +58,147 @@ class Testdata:
 
     def __init__(self, config: Config):
         self._config = config
-        self._api = FastAPI(docs_url=None, redoc_url=None)
         self._logger = logger.getLogger('testdata')
+        self._api = self._setup_api()
 
         # Store internal state
         self._state = {'data-used': 0}
 
-        @self._api.get('/zeros')
+    def _setup_api(self) -> FastAPI:
-        async def zeros(api_key: str, size: int | str, request: Request) -> StreamingResponse:
+        api = FastAPI(docs_url='/', redoc_url=None)
-            try:
-                extra = {'api_key': api_key, 'ip': request.client.host if request.client is not None else None, 'size': size}
-                self._logger.debug('Initiated request.', extra=extra)
 
-                if api_key not in config.authorized_keys:
+        # Security
-                    raise HTTPException(
+        def get_api_key(
-                        status_code=status.HTTP_401_UNAUTHORIZED,
+            api_key_query: str = Security(APIKeyQuery(name="api_key", auto_error=False)),
-                        detail='Invalid API Key.'
+            api_key_header: str = Security(APIKeyHeader(name="x-api-key", auto_error=False))
-                    )
+        ) -> str:
-                try:
+            # https://joshdimella.com/blog/adding-api-key-auth-to-fast-api
-                    size = convert_to_bytes(size)
-                except ValueError as err:
-                    self._logger.warning('Invalid format for size.', extra=extra)
-                    raise HTTPException(
-                        status_code=status.HTTP_400_BAD_REQUEST,
-                        detail='Invalid format for size.'
-                    ) from err
 
-                if size < 0:
+            if api_key_query in self._config.authorized_keys:
-                    raise MinSizePerRequestError
+                return api_key_query
-                if config.max_size < size:
+            if api_key_header in self._config.authorized_keys:
-                    raise MaxSizePerRequestError
+                return api_key_header
 
-                # update internal state
+            raise HTTPException(
-                if config.max_data < self._state['data-used'] + size:
+                status_code=status.HTTP_401_UNAUTHORIZED,
-                    raise HTTPException(
+                detail='Invalid or missing API Key'
-                        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            )
-                        detail='Service not available.'
-                    )
-                self._state['data-used'] += size
 
-                self._logger.debug('Successfully processed request.', extra=extra)
+        # A wrapper to set the function signature to accept the api key dependency
-                return StreamingResponse(
+        def secure(func):
-                    status_code=status.HTTP_200_OK,
+            # Get old signature
-                    content=generate_data(size, config.buffer_size),
+            positional_only, positional_or_keyword, variadic_positional, keyword_only, variadic_keyword = [], [], [], [], []
-                    media_type='application/octet-stream',
+            for value in inspect.signature(func).parameters.values():
-                    headers={
+                if value.kind == inspect.Parameter.POSITIONAL_ONLY:
-                        'Content-Length': str(size)
+                    positional_only.append(value)
-                    }
+                elif value.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD:
+                    positional_or_keyword.append(value)
+                elif value.kind == inspect.Parameter.VAR_POSITIONAL:
+                    variadic_positional.append(value)
+                elif value.kind == inspect.Parameter.KEYWORD_ONLY:
+                    keyword_only.append(value)
+                elif value.kind == inspect.Parameter.VAR_KEYWORD:
+                    variadic_keyword.append(value)
+
+            # Avoid passing an unrecognized keyword
+            if inspect.iscoroutinefunction(func):
+                async def wrapper(*args, **kwargs):
+                    if len(variadic_keyword) == 0:
+                        if 'api_key' in kwargs:
+                            del kwargs['api_key']
+
+                    return await func(*args, **kwargs)
+            else:
+                def wrapper(*args, **kwargs):
+                    if len(variadic_keyword) == 0:
+                        if 'api_key' in kwargs:
+                            del kwargs['api_key']
+
+                    return func(*args, **kwargs)
+
+            # Override signature
+            wrapper.__signature__ = inspect.signature(func).replace(
+                parameters=(
+                    *positional_only,
+                    *positional_or_keyword,
+                    *variadic_positional,
+                    *keyword_only,
+                    inspect.Parameter('api_key', inspect.Parameter.POSITIONAL_OR_KEYWORD, default=Security(get_api_key)),
+                    *variadic_keyword
                 )
+            )
 
-            except MinSizePerRequestError as err:
+            return functools.wraps(func)(wrapper)
-                self._logger.warning('Size if negative.', extra=extra)
+
-                raise HTTPException(
+        # Routes
-                    status_code=status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE,
+        api.get('/zeros')(secure(self._zeros))
-                    detail='Size has to be non-negative.'
+
-                ) from err
+        return api
-            except MaxSizePerRequestError as err:
+
-                self._logger.warning('Exceeded max size per request.', extra=extra)
+    async def _zeros(self, size: int | str, request: Request, filename: str = 'zeros.bin') -> StreamingResponse:
-                raise HTTPException(
+        try:
-                    status_code=status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE,
+            extra = {'id': f'{random.randint(0, 2 ** 32 - 1):08X}'}
-                    detail=f'Exceeded max size per request of {config.max_size} Bytes.'
+            self._logger.debug(
-                ) from err
+                'Initiated request.',
-            except Exception as err:
+                extra=extra | {
-                self._logger.exception(err)
+                    'ip': request.client.host if request.client is not None else None,
-                raise err
+                    'query-params': dict(request.query_params),
+                    'headers': dict(request.headers)
+                }
+            )
+
+            try:
+                size = convert_to_bytes(size)
+            except ValueError as err:
+                self._logger.warning('Invalid format for size.', extra=extra)
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail='Invalid format for size.'
+                ) from err
+
+            if size < 0:
+                raise MinSizePerRequestError
+            if self._config.max_size < size:
+                raise MaxSizePerRequestError
+
+            # update internal state
+            if self._config.max_data < self._state['data-used'] + size:
+                raise HTTPException(
+                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                    detail='Service not available.'
+                )
+            self._state['data-used'] += size
+
+            self._logger.debug('Successfully processed request.', extra=extra)
+            return StreamingResponse(
+                status_code=status.HTTP_200_OK,
+                content=generate_data(size, self._config.buffer_size),
+                media_type='application/octet-stream',
+                headers={
+                    'Content-Length': str(size),
+                    'Content-Disposition': f'attachment; filename="{filename}"'
+                }
+            )
+
+        except MinSizePerRequestError as err:
+            self._logger.warning('Size if negative.', extra=extra)
+            raise HTTPException(
+                status_code=status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE,
+                detail='Size has to be non-negative.'
+            ) from err
+        except MaxSizePerRequestError as err:
+            self._logger.warning('Exceeded max size per request.', extra=extra)
+            raise HTTPException(
+                status_code=status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE,
+                detail=f'Exceeded max size per request of {self._config.max_size} Bytes.'
+            ) from err
+        except Exception as err:
+            self._logger.exception(err)
+            raise err
+
+    async def _update_state(self) -> None:
+        assert self._config.database is not None
 
-    async def _update_state(self):
         mode = 'r+' if os.path.exists(self._config.database) else 'w+'
 
         with open(self._config.database, mode, encoding='utf-8') as file:
@@ -142,7 +221,7 @@ class Testdata:
             self._logger = logger.getLogger('testdata')
             self._logger.info('Server started.')
 
-            coroutines = [asyncio.create_task(uvicorn.Server(uvicorn.Config(self._api, host, port)).serve())]
+            coroutines = [uvicorn.Server(uvicorn.Config(self._api, host, port)).serve()]
             if self._config.database is not None:
                 coroutines.append(self._update_state())