111 lines
3.0 KiB
Python
111 lines
3.0 KiB
Python
import sys
|
|
import asyncio
|
|
import argparse
|
|
import json
|
|
from os.path import exists
|
|
|
|
from fastapi import FastAPI, Request, HTTPException, Query
|
|
from fastapi.responses import StreamingResponse
|
|
from fastapi import status
|
|
from hypercorn.config import Config
|
|
from hypercorn.asyncio import serve
|
|
import ipaddress
|
|
|
|
from .utils import convert_to_bytes, generate_data, load_database, save_database
|
|
|
|
# Setup Parser
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('-c', '--config', type=argparse.FileType('r'),
|
|
default='./config.json', help='Path to config file in JSON format.')
|
|
|
|
args = parser.parse_args(sys.argv[1:])
|
|
|
|
# Load Config
|
|
CONFIG = json.load(args.config)
|
|
BUFFER_SIZE = convert_to_bytes(CONFIG['buffer-size'])
|
|
MAX_SIZE = convert_to_bytes(CONFIG['max-size'])
|
|
MAX_DATA = convert_to_bytes(CONFIG['max-data'])
|
|
AUTHORIZED_KEYS = CONFIG['keys']
|
|
DATABASE = CONFIG['database']
|
|
|
|
if not exists(DATABASE):
|
|
save_database(DATABASE, {'data-used': 0})
|
|
|
|
|
|
api = FastAPI(docs_url=None, redoc_url=None)
|
|
|
|
|
|
class MaxSizePerRequestError(Exception):
|
|
pass
|
|
|
|
|
|
class MinSizePerRequestError(Exception):
|
|
pass
|
|
|
|
|
|
@api.get('/')
|
|
async def test_data(api_key: str, size: str) -> StreamingResponse:
|
|
try:
|
|
if api_key not in AUTHORIZED_KEYS:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail='Invalid API Key.'
|
|
)
|
|
|
|
try:
|
|
size = convert_to_bytes(size)
|
|
except ValueError as err:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail='Invalid format format for size.'
|
|
) from err
|
|
|
|
if size < 0:
|
|
raise MinSizePerRequestError
|
|
elif MAX_SIZE < size:
|
|
raise MaxSizePerRequestError
|
|
|
|
database = load_database(DATABASE)
|
|
if MAX_DATA <= database['data-used'] + size:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail='Service not available.'
|
|
)
|
|
database['data-used'] += size
|
|
|
|
save_database(DATABASE, database)
|
|
|
|
return StreamingResponse(
|
|
status_code=status.HTTP_200_OK,
|
|
content=generate_data(size, BUFFER_SIZE),
|
|
media_type='application/octet-stream',
|
|
headers={
|
|
'Content-Length': size
|
|
}
|
|
)
|
|
|
|
except MinSizePerRequestError as err:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE,
|
|
detail='Size has to be not-negative.'
|
|
) from err
|
|
except MaxSizePerRequestError as err:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE,
|
|
detail=f'Exceeded max size per request of {MAX_SIZE} Bytes.'
|
|
) from err
|
|
|
|
|
|
def main():
|
|
asyncio.run(serve(
|
|
api,
|
|
Config().from_mapping(
|
|
bind=CONFIG['binds'],
|
|
accesslog='-'
|
|
)
|
|
))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|