This commit is contained in:
Kristian Krsnik 2024-04-10 21:00:26 +02:00
parent 27a6dcfe7c
commit 549a8e2ac4
Signed by: Kristian
GPG Key ID: FD1330AC9F909E85
4 changed files with 46 additions and 6 deletions

2
.gitignore vendored
View File

@ -1,3 +1,5 @@
/result
/.direnv/
*.json

13
poetry.lock generated
View File

@ -132,6 +132,17 @@ files = [
{file = "idna-3.6.tar.gz", hash = "sha256:9ecdbbd083b06798ae1e86adcbfe8ab1479cf864e4ee30fe4e46a003d12491ca"},
]
[[package]]
name = "ipaddress"
version = "1.0.23"
description = "IPv4/IPv6 manipulation library"
optional = false
python-versions = "*"
files = [
{file = "ipaddress-1.0.23-py2.py3-none-any.whl", hash = "sha256:6e0f4a39e66cb5bb9a137b00276a2eff74f93b71dcbdad6f10ff7df9d3557fcc"},
{file = "ipaddress-1.0.23.tar.gz", hash = "sha256:b7f8e0369580bb4a24d5ba1d7cc29660a4a6987763faf1d8a8046830e020e7e2"},
]
[[package]]
name = "priority"
version = "2.0.0"
@ -309,4 +320,4 @@ h11 = ">=0.9.0,<1"
[metadata]
lock-version = "2.0"
python-versions = "^3.11"
content-hash = "c9da55bda660d1cadc77c14b7540377a8faa966892b8c4768abcf90cbda87e69"
content-hash = "a4a91c84503735a2120b16b6cd0a4c672588c4d58dc1fd410be0cdf8921982d2"

View File

@ -11,6 +11,7 @@ python = "^3.11"
fastapi = "^0.110.1"
hypercorn = "^0.16.0"
pydantic = "^2.6.4"
ipaddress = "^1.0.23"
[tool.poetry.scripts]
main = "src.main:main"

View File

@ -1,10 +1,27 @@
import sys
import asyncio
import argparse
import json
from fastapi import FastAPI
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from fastapi import status
from hypercorn.config import Config
from hypercorn.asyncio import serve
import ipaddress
# Setup Parser
parser = argparse.ArgumentParser()
parser.add_argument('-c', '--config', type=argparse.FileType('r'),
default='./config', help='Path to config file in JSON format.')
# parser.add_argument('-db', '--database', type=argparse.FileType('r'), # TODO: read+write
# default='./db.json', help='Path to database file in JSON format.')
args = parser.parse_args(sys.argv[1:])
# Load Config
config = json.load(args.config)
api = FastAPI()
@ -17,11 +34,8 @@ async def generate_test_data(size: int | str, max_size=1024 * 1024) -> bytes:
size_left = None
try:
print('Here1')
size_left = int(size)
except ValueError: # treat as string
print('Here2')
# TODO: maybe hardcode values
units = {
'GB': 10 ** 9, 'GiB': 2 ** 30,
'MB': 10 ** 6, 'MiB': 2 ** 20,
@ -44,13 +58,25 @@ async def generate_test_data(size: int | str, max_size=1024 * 1024) -> bytes:
yield b'\0' * size_left
def check_policies(ip: str) -> None:
network = ipaddress.ip_network(ip)
print(network)
@api.get('/')
async def test_data(size: str) -> StreamingResponse:
async def test_data(size: str, request: Request) -> StreamingResponse:
try:
check_policies(request.client.host)
return StreamingResponse(
status_code=status.HTTP_200_OK,
content=generate_test_data(size)
)
# except TooManyRequestsError as err:
# ...
# except BlockedError as err:
# ...
# except OutOfQuotaError as err:
# ...
except err:
raise err
pass