diff --git a/.gitignore b/.gitignore index ebd38e7..d0176a5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ /result /.direnv/ + +*.json diff --git a/poetry.lock b/poetry.lock index f1ea35a..05b4642 100644 --- a/poetry.lock +++ b/poetry.lock @@ -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" diff --git a/pyproject.toml b/pyproject.toml index cb4baf5..58f23ea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/src/main.py b/src/main.py index f1b24c4..9a98e52 100644 --- a/src/main.py +++ b/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