From 549a8e2ac48fa98fe9936db1d1f6a017363f640d Mon Sep 17 00:00:00 2001
From: Kristian Krsnik <git@krsnik.at>
Date: Wed, 10 Apr 2024 21:00:26 +0200
Subject: [PATCH] updated

---
 .gitignore     |  2 ++
 poetry.lock    | 13 ++++++++++++-
 pyproject.toml |  1 +
 src/main.py    | 36 +++++++++++++++++++++++++++++++-----
 4 files changed, 46 insertions(+), 6 deletions(-)

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