Implement brute force protection on ALL endpoints

Handle zombie process error

docs/genindex.html

@@ -50,6 +50,7 @@ <h1 id="index">Index</h1>
  | <a href="#G"><strong>G</strong></a>
  | <a href="#H"><strong>H</strong></a>
  | <a href="#I"><strong>I</strong></a>
+ | <a href="#L"><strong>L</strong></a>
  | <a href="#M"><strong>M</strong></a>
  | <a href="#N"><strong>N</strong></a>
  | <a href="#P"><strong>P</strong></a>
@@ -66,16 +67,14 @@ <h2 id="A">A</h2>
       <li><a href="index.html#pyninja.models.Session.allowed_origins">allowed_origins (pyninja.models.Session attribute)</a>
 </li>
       <li><a href="index.html#pyninja.models.EnvConfig.api_secret">api_secret (pyninja.models.EnvConfig attribute)</a>
-</li>
-      <li><a href="index.html#pyninja.models.EnvConfig.apikey">apikey (pyninja.models.EnvConfig attribute)</a>
 </li>
   </ul></td>
   <td style="width: 33%; vertical-align: top;"><ul>
+      <li><a href="index.html#pyninja.models.EnvConfig.apikey">apikey (pyninja.models.EnvConfig attribute)</a>
+</li>
       <li><a href="index.html#pyninja.exceptions.APIResponse">APIResponse</a>
 </li>
       <li><a href="index.html#pyninja.models.Session.auth_counter">auth_counter (pyninja.models.Session attribute)</a>
-</li>
-      <li><a href="index.html#pyninja.auth.authenticator">authenticator() (in module pyninja.auth)</a>
 </li>
   </ul></td>
 </tr></table>
@@ -121,7 +120,7 @@ <h2 id="E">E</h2>
   <td style="width: 33%; vertical-align: top;"><ul>
       <li><a href="index.html#pyninja.models.EnvConfig.Config">EnvConfig.Config (class in pyninja.models)</a>
 </li>
-      <li><a href="index.html#pyninja.routers.epoch">epoch() (in module pyninja.routers)</a>
+      <li><a href="index.html#pyninja.auth.EPOCH">EPOCH() (in module pyninja.auth)</a>
 </li>
       <li><a href="index.html#pyninja.models.EnvConfig.Config.extra">extra (pyninja.models.EnvConfig.Config attribute)</a>
 </li>
@@ -135,6 +134,8 @@ <h2 id="F">F</h2>
 </li>
   </ul></td>
   <td style="width: 33%; vertical-align: top;"><ul>
+      <li><a href="index.html#pyninja.auth.forbidden">forbidden() (in module pyninja.auth)</a>
+</li>
       <li><a href="index.html#pyninja.models.EnvConfig.from_env_file">from_env_file() (pyninja.models.EnvConfig class method)</a>
 </li>
   </ul></td>
@@ -161,7 +162,7 @@ <h2 id="G">G</h2>
 <h2 id="H">H</h2>
 <table style="width: 100%" class="indextable genindextable"><tr>
   <td style="width: 33%; vertical-align: top;"><ul>
-      <li><a href="index.html#pyninja.routers.handle_auth_error">handle_auth_error() (in module pyninja.routers)</a>
+      <li><a href="index.html#pyninja.auth.handle_auth_error">handle_auth_error() (in module pyninja.auth)</a>
 </li>
   </ul></td>
   <td style="width: 33%; vertical-align: top;"><ul>
@@ -173,7 +174,7 @@ <h2 id="H">H</h2>
 <h2 id="I">I</h2>
 <table style="width: 100%" class="indextable genindextable"><tr>
   <td style="width: 33%; vertical-align: top;"><ul>
-      <li><a href="index.html#pyninja.routers.incrementer">incrementer() (in module pyninja.routers)</a>
+      <li><a href="index.html#pyninja.auth.incrementer">incrementer() (in module pyninja.auth)</a>
 </li>
   </ul></td>
   <td style="width: 33%; vertical-align: top;"><ul>
@@ -184,6 +185,18 @@ <h2 id="I">I</h2>
   </ul></td>
 </tr></table>
 
+<h2 id="L">L</h2>
+<table style="width: 100%" class="indextable genindextable"><tr>
+  <td style="width: 33%; vertical-align: top;"><ul>
+      <li><a href="index.html#pyninja.auth.level_1">level_1() (in module pyninja.auth)</a>
+</li>
+  </ul></td>
+  <td style="width: 33%; vertical-align: top;"><ul>
+      <li><a href="index.html#pyninja.auth.level_2">level_2() (in module pyninja.auth)</a>
+</li>
+  </ul></td>
+</tr></table>
+
 <h2 id="M">M</h2>
 <table style="width: 100%" class="indextable genindextable"><tr>
   <td style="width: 33%; vertical-align: top;"><ul>
@@ -235,6 +248,8 @@ <h2 id="P">P</h2>
       <li><a href="index.html#pyninja.models.EnvConfig.parse_api_secret">parse_api_secret() (pyninja.models.EnvConfig class method)</a>
 </li>
       <li><a href="index.html#pyninja.models.Payload">Payload (class in pyninja.models)</a>
+</li>
+      <li><a href="index.html#pyninja.squire.process_command">process_command() (in module pyninja.squire)</a>
 </li>
       <li><a href="index.html#pyninja.routers.process_status">process_status() (in module pyninja.routers)</a>
 </li>

pyninja/auth.py

@@ -4,8 +4,8 @@
 from datetime import datetime
 from http import HTTPStatus
 
-from fastapi import Depends, Request
-from fastapi.security import HTTPBasicCredentials, HTTPBearer
+from fastapi import Request
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 
 from pyninja import database, exceptions, models
 
@@ -14,26 +14,96 @@
 SECURITY = HTTPBearer()
 
 
-async def authenticator(token: HTTPBasicCredentials = Depends(SECURITY)) -> None:
-    """Validates the token if mentioned as a dependency.
+async def forbidden(request: Request) -> None:
+    """Validates if a request is part of the forbidden list.
 
     Args:
-        token: Takes the authorization header token as an argument.
+        request: Reference to the FastAPI request object.
+
+    Raises:
+        APIResponse:
+        - 403: If host address is forbidden.
+    """
+    # placeholder list, to avoid a DB search for every request
+    if request.client.host in models.session.forbid:
+        # Get timestamp until which the host has to be forbidden
+        timestamp = database.get_record(request.client.host)
+        if timestamp and timestamp > EPOCH():
+            LOGGER.warning(
+                "%s is forbidden until %s due to repeated login failures",
+                request.client.host,
+                datetime.fromtimestamp(timestamp).strftime("%c"),
+            )
+            raise exceptions.APIResponse(
+                status_code=HTTPStatus.FORBIDDEN.value,
+                detail=f"{request.client.host!r} is not allowed",
+            )
+
+
+async def level_1(request: Request, apikey: HTTPAuthorizationCredentials) -> None:
+    """Validates the auth request using HTTPBearer.
+
+    Args:
+        request: Takes the authorization header token as an argument.
+        apikey: Basic APIKey required for all the routes.
 
     Raises:
         APIResponse:
         - 401: If authorization is invalid.
+        - 403: If host address is forbidden.
     """
-    auth = token.model_dump().get("credentials", "")
-    if auth.startswith("\\"):
-        auth = bytes(auth, "utf-8").decode(encoding="unicode_escape")
+    await forbidden(request)
+    if apikey.credentials.startswith("\\"):
+        auth = bytes(apikey.credentials, "utf-8").decode(encoding="unicode_escape")
+    else:
+        auth = apikey.credentials
     if secrets.compare_digest(auth, models.env.apikey):
+        LOGGER.info(
+            "Connection received from client-host: %s, host-header: %s, x-fwd-host: %s",
+            request.client.host,
+            request.headers.get("host"),
+            request.headers.get("x-forwarded-host"),
+        )
+        if user_agent := request.headers.get("user-agent"):
+            LOGGER.info("User agent: %s", user_agent)
         return
+    # Adds host address to the forbidden set
+    await handle_auth_error(request)
     raise exceptions.APIResponse(
         status_code=HTTPStatus.UNAUTHORIZED.real, detail=HTTPStatus.UNAUTHORIZED.phrase
     )
 
 
+async def level_2(
+    request: Request, apikey: HTTPAuthorizationCredentials, token: str
+) -> None:
+    """Validates the auth request using HTTPBearer and additionally a secure token.
+
+    Args:
+        request: Takes the authorization header token as an argument.
+        apikey: Basic APIKey required for all the routes.
+        token: Additional token for critical requests.
+
+    Raises:
+        APIResponse:
+        - 401: If authorization is invalid.
+        - 403: If host address is forbidden.
+    """
+    await level_1(request, apikey)
+    if not all((models.env.remote_execution, models.env.api_secret)):
+        raise exceptions.APIResponse(
+            status_code=HTTPStatus.NOT_IMPLEMENTED.real,
+            detail="Remote execution has been disabled on the server.",
+        )
+    if token and secrets.compare_digest(token, models.env.api_secret):
+        return
+    await handle_auth_error(request)
+    raise exceptions.APIResponse(
+        status_code=HTTPStatus.UNAUTHORIZED.real,
+        detail=HTTPStatus.UNAUTHORIZED.phrase,
+    )
+
+
 async def incrementer(attempt: int) -> int:
     """Increments block time for a host address based on the number of failed attempts.
 

pyninja/main.py

@@ -7,7 +7,7 @@
 import pyninja
 from pyninja import models, routers, squire
 
-LOGGER = logging.getLogger(__name__)
+LOGGER = logging.getLogger("uvicorn.error")
 
 
 def start(**kwargs) -> None:

pyninja/process.py

@@ -8,22 +8,32 @@
 
 
 def get_process_status(process_name: str) -> Generator[Dict[str, int]]:
-    """Get process ID for a particular service.
+    """Get process information by name.
 
     Args:
-        service_name (str): Name of the service.
+        process_name: Name of the process.
 
     Yields:
         Generator[Dict[str, int]]:
         Yields the process metrics as a dictionary of key-value pairs.
     """
+    # todo: implement concurrency
     for proc in psutil.process_iter(["pid", "name"]):
-        if proc.info["name"].lower() == process_name.lower():
-            process = psutil.Process(proc.info["pid"])
-            yield get_performance(process_name, process)
-
-
-def get_performance(process_name: str, process: psutil.Process) -> Dict[str, int]:
+        if proc.name().lower() == process_name.lower():
+            process = psutil.Process(proc.pid)
+            process._name = process_name
+            try:
+                perf_report = get_performance(process)
+                LOGGER.info({f"{process_name} [{process.pid}]": perf_report})
+                perf_report["pname"] = process_name
+                perf_report["zombie"] = False
+                yield perf_report
+            except psutil.ZombieProcess as warn:
+                LOGGER.warning(warn)
+                yield {"zombie": True, "process_name": process_name}
+
+
+def get_performance(process: psutil.Process) -> Dict[str, int | float]:
     """Checks performance by monitoring CPU utilization, number of threads and open files.
 
     Args:
@@ -36,8 +46,9 @@ def get_performance(process_name: str, process: psutil.Process) -> Dict[str, int
     cpu = process.cpu_percent(interval=0.5)
     threads = process.num_threads()
     open_files = len(process.open_files())
-    info_dict = {"cpu": cpu, "threads": threads, "open_files": open_files}
-    LOGGER.info({f"{process_name} [{process.pid}]": info_dict})
-    info_dict["pid"] = process.pid.real
-    info_dict["pname"] = process_name
-    return info_dict
+    return {
+        "cpu": cpu,
+        "threads": threads,
+        "open_files": open_files,
+        "pid": process.pid.real,
+    }