API Rate Limiting for Bad Actor; Small Code Refactor (#8)

* Cleanup Log.py

* Remove Unused Import

* Correct Typo

* Update Requirements to have rate-limiter package

* Update routes to use rate limiting

* rollback Rate Limit Changes

* Rollback Rate Limit; Implement Metrics

* Split Middleware to own files

* update requirements.txt

* remove unneeded requirement

* remove unused imports

* Adjust Middleware files to be better importable.

* remove unneeded API Singleton

* Update Authentication Import Path

* Properly implement new file layout.

* remove manual auth checking, 'prebuilt' middleware

* implement authentication middleware as middleware

* add authentication middleware to FastAPI App

* readd eof new line

---------

Co-authored-by: Chats <[email protected]>

README.md

@@ -1,4 +1,4 @@
-<img src="https://i.imgur.com/I1wosdV.png" width="55%">
+<img src="https://i.imgur.com/I1wosdV.png" width="55%" alt="HD2 Community API">
 
 [![Python](https://img.shields.io/badge/python-3670A0?style=for-the-badge&logo=python&logoColor=ffdd54)](https://www.python.org/)
 [![FastAPI](https://img.shields.io/badge/FastAPI-005571?style=for-the-badge&logo=fastapi)](https://fastapi.tiangolo.com/)[![Python Black](https://img.shields.io/badge/Python%20Black-000000?style=for-the-badge&logo=python&logoColor=FFFFFF&labelColor=000000&color=000000)](https://github.com/psf/black)
@@ -34,7 +34,7 @@ This can be renamed to .env and used as is, and it will use api.diveharder.com<b
 Or you may change the links to the AHGS API endpoints if you have them.
 
 SECURITY_TOKEN is what you use to access the /admin/* endpoints <br />
-SESSION_TOKEN is for accessing AHGS API's that require authentication
+SESSION_TOKEN is for accessing AHGS APIs that require authentication
 </details>
 
 <details>

requirements.txt

@@ -8,3 +8,4 @@ gitpython~=3.1.43
 gunicorn~=21.2.0
 uvicorn~=0.29.0
 brotli-asgi~=1.4.0
+prometheus-fastapi-instrumentator~=7.0.0

src/main.py

@@ -1,11 +1,19 @@
-from fastapi import FastAPI, APIRouter
-from fastapi.requests import Request
+from fastapi import FastAPI
+from starlette.middleware.base import BaseHTTPMiddleware
+
+# Middleware Improts
 from starlette.middleware.cors import CORSMiddleware
 from brotli_asgi import BrotliMiddleware
+from src.utils.middleware.case_sens_middleware import case_sens_middleware
+from src.utils.middleware.rate_limit import cutoff
+from src.utils.middleware.metrics import instrumentator
+from src.utils.middleware.authentication import authenticate
 
+# Routes Import
 from src.routes import base, admin, raw, v1
 
-app = FastAPI(
+# -------- INIT API -------- #
+app: FastAPI = FastAPI(
     title="HD2 Community API",
     openapi_url="/openapi.json",
     description="""
@@ -15,7 +23,7 @@
     Github: https://github.com/helldivers-2/diveharder_api.py/""",
 )
 
-# noinspection PyTypeChecker
+# -------- MIDDLEWARE -------- #
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -24,21 +32,14 @@
     allow_headers=["*"],
 )
 app.add_middleware(BrotliMiddleware)
+app.add_middleware(BaseHTTPMiddleware, dispatch=case_sens_middleware)
+app.add_middleware(BaseHTTPMiddleware, dispatch=cutoff)
+app.add_middleware(BaseHTTPMiddleware, dispatch=authenticate)
+instrumentator.instrument(app=app).expose(
+    app, include_in_schema=False, should_gzip=True
+)
 
-
-@app.middleware("http")
-async def case_sens_middleware(request: Request, call_next):
-    decode_format = "utf-8"
-    raw_query_str = request.scope["query_string"].decode(decode_format).lower()
-    request.scope["query_string"] = raw_query_str.encode(decode_format)
-
-    path = request.scope["path"].lower()
-    request.scope["path"] = path
-
-    response = await call_next(request)
-    return response
-
-
+# -------- ROUTES -------- #
 app.include_router(base.router)
 app.include_router(admin.router)
 app.include_router(raw.router)

src/routes/admin.py

@@ -3,59 +3,48 @@
 
 import src.utils.log as log
 from src.cfg.settings import ahgs_api as api_cfg
-from src.utils.authentication import verify_password
 
 router = APIRouter(
     prefix="/admin",
     tags=["admin"],
-    include_in_schema=True,
-    responses={401: {"description": "Unauthorized"}},
+    include_in_schema=False,
 )
 
 
-@router.get("/settings", include_in_schema=False)
+@router.get("/settings")
 async def settings(request: Request, source: str = None):
-    if verify_password(request.headers.get("Authorization")):
-        log.info(request, status.HTTP_200_OK, source)
-        return {
-            "log_level": log.logger.level,
-            "session_token": api_cfg["auth_headers"].get("Authorization", ""),
-        }
-    log.info(request, status.HTTP_401_UNAUTHORIZED, source)
-    raise HTTPException(status_code=401, detail="Unauthorized")
+    log.info(request, status.HTTP_200_OK, source)
+    return {
+        "log_level": log.logger.level,
+        "session_token": api_cfg["auth_headers"].get("Authorization", ""),
+    }
 
 
-@router.post("/session", include_in_schema=False)
+@router.post("/session")
 async def settings(request: Request, source: str = None):
-    if verify_password(request.headers.get("Authorization")):
-        request_json = await request.json()
-        token = request_json["token"]
-        if token:
-            api_cfg["auth_headers"]["Authorization"] = f"{token}"
-            log.info(request, status.HTTP_202_ACCEPTED, source)
-            raise HTTPException(status_code=status.HTTP_202_ACCEPTED)
-        else:
-            log.info(request, status.HTTP_400_BAD_REQUEST, source)
-            raise HTTPException(
-                status_code=status.HTTP_400_BAD_REQUEST, detail="Did not provide token"
-            )
-    log.info(request, status.HTTP_401_UNAUTHORIZED, source)
-    raise HTTPException(status_code=401, detail="Unauthorized")
-
-
-@router.post("/loglevel", include_in_schema=False)
+    request_json = await request.json()
+    token = request_json["token"]
+    if token:
+        api_cfg["auth_headers"]["Authorization"] = f"{token}"
+        log.info(request, status.HTTP_202_ACCEPTED, source)
+        raise HTTPException(status_code=status.HTTP_202_ACCEPTED)
+    else:
+        log.info(request, status.HTTP_400_BAD_REQUEST, source)
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST, detail="Did not provide token"
+        )
+
+
+@router.post("/loglevel")
 async def settings(request: Request, source: str = None):
-    if verify_password(request.headers.get("Authorization")):
-        request_json = await request.json()
-        loglevel_in = request_json["loglevel"]
-        if loglevel_in:
-            log.update_log_level(log_level="loglevel_in")
-            log.info(request, status.HTTP_202_ACCEPTED, source)
-            raise HTTPException(status_code=status.HTTP_202_ACCEPTED)
-        else:
-            log.info(request, status.HTTP_400_BAD_REQUEST, source)
-            raise HTTPException(
-                status_code=status.HTTP_400_BAD_REQUEST, detail="Did not provide token"
-            )
-    log.info(request, status.HTTP_401_UNAUTHORIZED, source)
-    raise HTTPException(status_code=401, detail="Unauthorized")
+    request_json = await request.json()
+    loglevel_in = request_json["loglevel"]
+    if loglevel_in:
+        log.update_log_level(log_level="loglevel_in")
+        log.info(request, status.HTTP_202_ACCEPTED, source)
+        raise HTTPException(status_code=status.HTTP_202_ACCEPTED)
+    else:
+        log.info(request, status.HTTP_400_BAD_REQUEST, source)
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST, detail="Did not provide token"
+        )

src/utils/log.py

@@ -25,7 +25,6 @@ def update_log_level(log_level):
             logger.setLevel(logging.NOTSET)
 
 
-# user="trainingmanual3"
 def info(request: Request, response_status: status, user: str = ""):
     if user == "" or user is None:
         source = request.headers.get("User-Agent")

src/utils/middleware/authentication.py

@@ -0,0 +1,45 @@
+import bcrypt
+
+from fastapi import status
+from fastapi.requests import Request
+from fastapi.responses import JSONResponse
+
+from src.cfg.settings import security
+
+TOKEN = security["token"]
+hashed = bcrypt.hashpw(bytes(TOKEN, "utf-8"), bcrypt.gensalt())
+
+
+async def authenticate(request: Request, call_next):
+    authenticated_prefixes = "admin"
+    response = None
+    headers = dict(request.scope["headers"])
+    if (
+        authenticated_prefixes in request.url.path
+        and "authorization" in request.headers
+    ):
+        if verify_password(request.headers["Authorization"]):
+            response = await call_next(request)
+        else:
+            response = JSONResponse(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                content={"error": "Unauthorized"},
+            )
+    else:
+        response = await call_next(request)
+    return response
+
+
+def verify_length(token):
+    length = len(token)
+    if length <= 7 or length >= 25:
+        return False
+    return True
+
+
+def verify_password(plaintext_token):
+    if verify_length(plaintext_token):
+        if bcrypt.checkpw(str.encode(plaintext_token), hashed):
+            return True
+        else:
+            return False

src/utils/middleware/case_sens_middleware.py

@@ -0,0 +1,13 @@
+from fastapi.requests import Request
+
+
+async def case_sens_middleware(request: Request, call_next):
+    decode_format = "utf-8"
+    raw_query_str = request.scope["query_string"].decode(decode_format).lower()
+    request.scope["query_string"] = raw_query_str.encode(decode_format)
+
+    path = request.scope["path"].lower()
+    request.scope["path"] = path
+
+    response = await call_next(request)
+    return response

src/utils/middleware/metrics.py

@@ -0,0 +1,40 @@
+from prometheus_fastapi_instrumentator import Instrumentator, metrics
+
+# from prometheus_fastapi_instrumentator.metrics import Info
+# from prometheus_client import Counter
+
+instrumentator = Instrumentator(
+    should_group_status_codes=True,
+    should_ignore_untemplated=False,
+    should_group_untemplated=False,
+    should_round_latency_decimals=True,
+    should_instrument_requests_inprogress=True,
+    excluded_handlers=[".*admin.*", "/metrics.py"],
+    inprogress_name="inprogress",
+    inprogress_labels=True,
+)
+
+instrumentator.add(
+    metrics.request_size(
+        should_include_method=True,
+        should_include_status=True,
+        should_include_handler=True,
+        metric_namespace="request",
+        metric_subsystem="size",
+    )
+).add(
+    metrics.response_size(
+        should_include_handler=True,
+        should_include_status=True,
+        should_include_method=False,
+        metric_namespace="response",
+        metric_subsystem="size",
+    )
+).add(
+    metrics.latency(
+        should_include_handler=True,
+        should_include_status=True,
+        should_include_method=False,
+        metric_namespace="latency",
+    )
+)

src/utils/middleware/rate_limit.py

@@ -0,0 +1,21 @@
+import src.utils.log as log
+from fastapi.requests import Request
+from fastapi.responses import JSONResponse
+
+
+async def cutoff(request: Request, call_next):
+    block_list = "Helldivers%20Companion/238 CFNetwork/1494.0.7 Darwin/23.4.0"
+    decode_format = "utf-8"
+    raw_query_str = request.headers.get("User-Agent", "")
+    if block_list == raw_query_str:
+        log.info(request, 429, block_list)
+        response = JSONResponse(
+            status_code=429,
+            content={
+                "limited": "Contact @chatterchats on Discord.",
+            },
+        )
+        return response
+    else:
+        response = await call_next(request)
+        return response