updated
This commit is contained in:
parent
27a6dcfe7c
commit
549a8e2ac4
2
.gitignore
vendored
2
.gitignore
vendored
@ -1,3 +1,5 @@
|
|||||||
/result
|
/result
|
||||||
|
|
||||||
/.direnv/
|
/.direnv/
|
||||||
|
|
||||||
|
*.json
|
||||||
|
13
poetry.lock
generated
13
poetry.lock
generated
@ -132,6 +132,17 @@ files = [
|
|||||||
{file = "idna-3.6.tar.gz", hash = "sha256:9ecdbbd083b06798ae1e86adcbfe8ab1479cf864e4ee30fe4e46a003d12491ca"},
|
{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]]
|
[[package]]
|
||||||
name = "priority"
|
name = "priority"
|
||||||
version = "2.0.0"
|
version = "2.0.0"
|
||||||
@ -309,4 +320,4 @@ h11 = ">=0.9.0,<1"
|
|||||||
[metadata]
|
[metadata]
|
||||||
lock-version = "2.0"
|
lock-version = "2.0"
|
||||||
python-versions = "^3.11"
|
python-versions = "^3.11"
|
||||||
content-hash = "c9da55bda660d1cadc77c14b7540377a8faa966892b8c4768abcf90cbda87e69"
|
content-hash = "a4a91c84503735a2120b16b6cd0a4c672588c4d58dc1fd410be0cdf8921982d2"
|
||||||
|
@ -11,6 +11,7 @@ python = "^3.11"
|
|||||||
fastapi = "^0.110.1"
|
fastapi = "^0.110.1"
|
||||||
hypercorn = "^0.16.0"
|
hypercorn = "^0.16.0"
|
||||||
pydantic = "^2.6.4"
|
pydantic = "^2.6.4"
|
||||||
|
ipaddress = "^1.0.23"
|
||||||
|
|
||||||
[tool.poetry.scripts]
|
[tool.poetry.scripts]
|
||||||
main = "src.main:main"
|
main = "src.main:main"
|
||||||
|
36
src/main.py
36
src/main.py
@ -1,10 +1,27 @@
|
|||||||
|
import sys
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
|
||||||
from fastapi import FastAPI
|
from fastapi import FastAPI, Request
|
||||||
from fastapi.responses import StreamingResponse
|
from fastapi.responses import StreamingResponse
|
||||||
from fastapi import status
|
from fastapi import status
|
||||||
from hypercorn.config import Config
|
from hypercorn.config import Config
|
||||||
from hypercorn.asyncio import serve
|
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()
|
api = FastAPI()
|
||||||
|
|
||||||
@ -17,11 +34,8 @@ async def generate_test_data(size: int | str, max_size=1024 * 1024) -> bytes:
|
|||||||
size_left = None
|
size_left = None
|
||||||
|
|
||||||
try:
|
try:
|
||||||
print('Here1')
|
|
||||||
size_left = int(size)
|
size_left = int(size)
|
||||||
except ValueError: # treat as string
|
except ValueError: # treat as string
|
||||||
print('Here2')
|
|
||||||
# TODO: maybe hardcode values
|
|
||||||
units = {
|
units = {
|
||||||
'GB': 10 ** 9, 'GiB': 2 ** 30,
|
'GB': 10 ** 9, 'GiB': 2 ** 30,
|
||||||
'MB': 10 ** 6, 'MiB': 2 ** 20,
|
'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
|
yield b'\0' * size_left
|
||||||
|
|
||||||
|
|
||||||
|
def check_policies(ip: str) -> None:
|
||||||
|
network = ipaddress.ip_network(ip)
|
||||||
|
print(network)
|
||||||
|
|
||||||
@api.get('/')
|
@api.get('/')
|
||||||
async def test_data(size: str) -> StreamingResponse:
|
async def test_data(size: str, request: Request) -> StreamingResponse:
|
||||||
try:
|
try:
|
||||||
|
check_policies(request.client.host)
|
||||||
|
|
||||||
return StreamingResponse(
|
return StreamingResponse(
|
||||||
status_code=status.HTTP_200_OK,
|
status_code=status.HTTP_200_OK,
|
||||||
content=generate_test_data(size)
|
content=generate_test_data(size)
|
||||||
)
|
)
|
||||||
|
# except TooManyRequestsError as err:
|
||||||
|
# ...
|
||||||
|
# except BlockedError as err:
|
||||||
|
# ...
|
||||||
|
# except OutOfQuotaError as err:
|
||||||
|
# ...
|
||||||
except err:
|
except err:
|
||||||
raise err
|
raise err
|
||||||
pass
|
pass
|
||||||
|
Loading…
Reference in New Issue
Block a user