updated
This commit is contained in:
parent
27a6dcfe7c
commit
549a8e2ac4
2
.gitignore
vendored
2
.gitignore
vendored
@ -1,3 +1,5 @@
|
||||
/result
|
||||
|
||||
/.direnv/
|
||||
|
||||
*.json
|
||||
|
13
poetry.lock
generated
13
poetry.lock
generated
@ -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"
|
||||
|
@ -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"
|
||||
|
36
src/main.py
36
src/main.py
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user