Compare commits

...

5 Commits
v1.1.0 ... main

8 changed files with 261 additions and 85 deletions

28
Dockerfile Normal file
View File

@ -0,0 +1,28 @@
# Nix builder
FROM nixos/nix:latest AS builder
# Copy our source and setup our working dir.
COPY . /tmp/build
WORKDIR /tmp/build
# Build our Nix environment
RUN nix \
--extra-experimental-features "nix-command flakes" \
--option filter-syscalls false \
build
# Copy the Nix store closure into a directory. The Nix store closure is the
# entire set of Nix store values that we need for our build.
RUN mkdir /tmp/nix-store-closure
RUN cp -r $(nix-store -qR result/) /tmp/nix-store-closure
# Final image is based on scratch. We copy a bunch of Nix dependencies
# but they're fully self-contained so we don't need Nix anymore.
FROM scratch
WORKDIR /app
# Copy /nix/store
COPY --from=builder /tmp/nix-store-closure /nix/store
COPY --from=builder /tmp/build/result /app
CMD ["/app/bin/testdata"]

15
docker-compose.yaml Normal file
View File

@ -0,0 +1,15 @@
services:
testdata:
image: result/latest
build:
dockerfile: ./Dockerfile
environment:
TESTDATA_HOST: 0.0.0.0
TESTDATA_PORT: 1234
TESTDATA_CONFIG: ./config.json
volumes:
- ./config.json:/app/config.json
- ./db.json:/app/db.json
- ./log.jsonl:/app/log.jsonl

View File

@ -61,17 +61,17 @@
in '' in ''
${ ${
if builtins.elem "pytest" dev && !skipCheck if builtins.elem "pytest" dev && !skipCheck
then "pytest src tests" then "pytest tests"
else "" else ""
} }
${ ${
if builtins.elem "mypy" dev && !skipCheck if builtins.elem "mypy" dev && !skipCheck
then "mypy src tests" then "mypy src"
else "" else ""
} }
${ ${
if builtins.elem "pylint" dev && !skipCheck if builtins.elem "pylint" dev && !skipCheck
then "pylint src tests" then "pylint src"
else "" else ""
} }
''; '';

View File

@ -1,27 +1,27 @@
[project] [project]
name = "testdata" name = "testdata"
version = "1.1.0" version = "1.2.1"
requires-python = "~=3.12, <4" requires-python = "~=3.12, <4"
dependencies = [ dependencies = [
"fastapi~=0.115.3", "fastapi~=0.115",
"uvicorn~=0.32.0", "uvicorn~=0.32",
"pydantic~=2.9.2", "pydantic~=2.9",
] ]
[project.optional-dependencies] [project.optional-dependencies]
dev = [ dev = [
"pytest~=8.3.4", "pytest~=8.3",
"mypy~=1.13.0", "mypy~=1.13",
"pylint~=3.3.3", "pylint~=3.3",
"requests~=2.32.3", "requests~=2.32",
"types-requests~=2.32.0" "types-requests~=2.32"
] ]
[project.scripts] [project.scripts]
testdata = "testdata.main:main" testdata = "testdata.main:main"
[build-system] [build-system]
requires = ["setuptools~=75.1.1"] requires = ["setuptools~=75.1"]
build-backend = "setuptools.build_meta" build-backend = "setuptools.build_meta"
[tool.setuptools.packages.find] [tool.setuptools.packages.find]

View File

@ -132,7 +132,7 @@ def generate_log_config(log_path: str | None = None) -> dict:
'class': logging.handlers.RotatingFileHandler, 'class': logging.handlers.RotatingFileHandler,
'level': 'DEBUG', 'level': 'DEBUG',
'formatter': 'json', 'formatter': 'json',
'filename': 'log.jsonl', 'filename': log_path,
'maxBytes': 1024 * 1024 * 10, # 10 MiB 'maxBytes': 1024 * 1024 * 10, # 10 MiB
'backupCount': 3 'backupCount': 3
}} if log_path is not None else {}), }} if log_path is not None else {}),

26
src/testdata/main.py vendored
View File

@ -1,14 +1,32 @@
import os
import sys import sys
import argparse import argparse
import asyncio import asyncio
import shutil
from .testdata import Testdata from .testdata import Testdata
def parse_args(args: list[str]): def parse_args(args: list[str]):
parser = argparse.ArgumentParser() def formatter(prog):
parser.add_argument('-c', '--config', type=argparse.FileType('r'), default='./config.json', help='Path to config file in JSON format.') return argparse.ArgumentDefaultsHelpFormatter(prog, max_help_position=shutil.get_terminal_size().columns)
parser.add_argument('-l', '--listen', type=str, default='0.0.0.0', help='IP on which to listen.')
parser.add_argument('-p', '--port', type=int, default='8080', help='Port on which to serve the webserver.') parser = argparse.ArgumentParser(formatter_class=formatter)
parser.add_argument(
'-c', '--config', type=argparse.FileType('r'),
default=os.environ['TESTDATA_CONFIG'] if 'TESTDATA_CONFIG' in os.environ else './config.json',
help='Path to config file in JSON format.'
)
parser.add_argument(
'-l', '--listen', type=str,
default=os.environ['TESTDATA_HOST'] if 'TESTDATA_HOST' in os.environ else '0.0.0.0',
help='IP on which to listen.'
)
parser.add_argument(
'-p', '--port', type=int,
default=os.environ['TESTDATA_PORT'] if 'TESTDATA_PORT' in os.environ else 8080,
help='Port on which to serve the webserver.'
)
return parser.parse_args(args) return parser.parse_args(args)

View File

@ -1,10 +1,16 @@
import os import os
import json import json
import asyncio import asyncio
import inspect
import functools
import random
import importlib.metadata
from datetime import datetime
from typing_extensions import Annotated
import uvicorn 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 fastapi.responses import StreamingResponse
from pydantic import BaseModel, ConfigDict, Field, BeforeValidator, ValidationError from pydantic import BaseModel, ConfigDict, Field, BeforeValidator, ValidationError
@ -49,77 +55,158 @@ class Testdata:
_config: Config _config: Config
_api: FastAPI _api: FastAPI
_state: dict[str, int] _state: dict
_logger: logger.Logger _logger: logger.Logger
def __init__(self, config: Config): def __init__(self, config: Config):
self._config = config self._config = config
self._api = FastAPI(docs_url=None, redoc_url=None)
self._logger = logger.getLogger('testdata') self._logger = logger.getLogger('testdata')
self._api = self._setup_api()
# Store internal state # Store internal state
self._state = {'data-used': 0} self._state = {
'version': importlib.metadata.version('testdata'), # For future compatibility
'data-used': {f'{(today := datetime.today()).year}-{today.month:02}': 0} # math each months data usage
}
@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
current_date = f'{(today := datetime.today()).year}-{today.month:02}'
if current_date not in self._state['data-used']:
self._state['data-used'][current_date] = 0
if self._config.max_data < self._state['data-used'][current_date] + size:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail='Service not available.'
)
self._state['data-used'][current_date] += 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+' mode = 'r+' if os.path.exists(self._config.database) else 'w+'
with open(self._config.database, mode, encoding='utf-8') as file: with open(self._config.database, mode, encoding='utf-8') as file:
@ -130,7 +217,7 @@ class Testdata:
while True: while True:
file.seek(0) file.seek(0)
json.dump(self._state, file) json.dump(self._state, file, indent=2)
file.truncate() file.truncate()
await asyncio.sleep(self._config.database_update_interval) await asyncio.sleep(self._config.database_update_interval)
@ -142,7 +229,7 @@ class Testdata:
self._logger = logger.getLogger('testdata') self._logger = logger.getLogger('testdata')
self._logger.info('Server started.') 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: if self._config.database is not None:
coroutines.append(self._update_state()) coroutines.append(self._update_state())

View File

@ -53,6 +53,17 @@ def _server(request) -> Generator[str, None, None]:
break break
@pytest.mark.parametrize('_server', [({
'keys': ['one', 'two', 'three'],
'max-size': '100',
'max-data': 1234,
'buffer-size': '12MiB',
})], indirect=['_server'])
def test_invalid_api_key(_server):
response = requests.get(f'{PROTOCOL}://{HOST}:{PORT}/zeros?api_key=four&size=100', timeout=TIMEOUT)
assert response.status_code == 401
@pytest.mark.parametrize('_server', [({ @pytest.mark.parametrize('_server', [({
'keys': ['one', 'two', 'three'], 'keys': ['one', 'two', 'three'],
'max-size': '100', 'max-size': '100',
@ -85,13 +96,16 @@ def test_request_size_upper_bound(_server):
@pytest.mark.parametrize('_server', [({ @pytest.mark.parametrize('_server', [({
'keys': ['one', 'two', 'three'], 'keys': ['one', 'two', 'three'],
'max-size': '100', 'max-size': '100KB',
'max-data': 1234, 'max-data': '100KB',
'buffer-size': '12MiB', 'buffer-size': '12MiB',
})], indirect=['_server']) })], indirect=['_server'])
def test_invalid_api_key(_server): def test_request_max_data_used(_server):
response = requests.get(f'{PROTOCOL}://{HOST}:{PORT}/zeros?api_key=four&size=100', timeout=TIMEOUT) response = requests.get(f'{PROTOCOL}://{HOST}:{PORT}/zeros?api_key=one&size=100KB', timeout=TIMEOUT)
assert response.status_code == 401 assert response.status_code == 200
response = requests.get(f'{PROTOCOL}://{HOST}:{PORT}/zeros?api_key=one&size=1', timeout=TIMEOUT)
assert response.status_code == 500
@pytest.mark.parametrize('_server', [({ @pytest.mark.parametrize('_server', [({
@ -102,15 +116,29 @@ def test_invalid_api_key(_server):
'database-update-interval': 0.1 'database-update-interval': 0.1
})], indirect=['_server']) })], indirect=['_server'])
def test_check_database_update(_server): def test_check_database_update(_server):
import importlib.metadata
from datetime import datetime
database = _server database = _server
with open(database, 'r', encoding='utf-8') as file: with open(database, 'r', encoding='utf-8') as file:
file.seek(0) file.seek(0)
assert json.load(file) == {'data-used': 0} today = datetime.today()
assert json.load(file) == {
'version': importlib.metadata.version('testdata'),
'data-used': {
f'{today.year}-{today.month:02}': 0
}
}
response = requests.get(f'{PROTOCOL}://{HOST}:{PORT}/zeros?api_key=one&size=100', timeout=TIMEOUT) response = requests.get(f'{PROTOCOL}://{HOST}:{PORT}/zeros?api_key=one&size=100', timeout=TIMEOUT)
assert response.status_code == 200 assert response.status_code == 200
time.sleep(0.1) time.sleep(0.1)
file.seek(0) file.seek(0)
assert json.load(file) == {'data-used': 100} assert json.load(file) == {
'version': importlib.metadata.version('testdata'),
'data-used': {
f'{today.year}-{today.month:02}': 100
}
}