updated

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