PR : Add reporting for API

aikido_firewall/__init__.py

@@ -4,6 +4,8 @@
 
 from dotenv import load_dotenv
 
+# Constants
+PKG_VERSION = "0.0.1"
 
 # Import logger
 from aikido_firewall.helpers.logging import logger

aikido_firewall/background_process/api/__init__.py

@@ -0,0 +1,28 @@
+""" 
+init.py file for api/ folder. Includes abstract class ReportingApi
+"""
+
+import json
+from aikido_firewall.helpers.logging import logger
+
+
+class ReportingApi:
+    """This is the super class for the reporting API's"""
+
+    def to_api_response(self, res):
+        """Converts results into an Api response obj"""
+        status = res.status_code
+        if status == 429:
+            return {"success": False, "error": "rate_limited"}
+        elif status == 401:
+            return {"success": False, "error": "invalid_token"}
+        elif status == 200:
+            try:
+                return json.loads(res.text)
+            except Exception as e:
+                logger.debug(e)
+                logger.debug(res.text)
+        return {"success": False, "error": "unknown_error"}
+
+    def report(self, token, event, timeout_in_sec):
+        """Report event to aikido server"""

aikido_firewall/background_process/api/http_api.py

@@ -0,0 +1,32 @@
+"""
+Exports the HTTP API class
+"""
+
+import requests
+from aikido_firewall.background_process.api import ReportingApi
+
+
+class ReportingApiHTTP(ReportingApi):
+    """HTTP Reporting API"""
+
+    def __init__(self, reporting_url):
+        self.reporting_url = reporting_url
+
+    def report(self, token, event, timeout_in_sec):
+        try:
+            res = requests.post(
+                self.reporting_url + "api/runtime/events",
+                json=event,
+                timeout=timeout_in_sec,
+                headers=get_headers(token),
+            )
+        except requests.exceptions.ConnectionError:
+            return {"success": False, "error": "timeout"}
+        except Exception as e:
+            raise e
+        return self.to_api_response(res)
+
+
+def get_headers(token):
+    """Returns headers"""
+    return {"Content-Type": "application/json", "Authorization": str(token)}

aikido_firewall/background_process/api/init_test.py

@@ -0,0 +1,44 @@
+import pytest
+from aikido_firewall.background_process.api import ReportingApi
+
+# Test ReportingApi Class :
+from requests.models import Response
+
+
+@pytest.fixture
+def reporting_api():
+    return ReportingApi()
+
+
+def test_to_api_response_rate_limited(reporting_api):
+    res = Response()
+    res.status_code = 429
+    assert reporting_api.to_api_response(res) == {
+        "success": False,
+        "error": "rate_limited",
+    }
+
+
+def test_to_api_response_invalid_token(reporting_api):
+    res = Response()
+    res.status_code = 401
+    assert reporting_api.to_api_response(res) == {
+        "success": False,
+        "error": "invalid_token",
+    }
+
+
+def test_to_api_response_unknown_error(reporting_api):
+    res = Response()
+    res.status_code = 500  # Simulating an unknown error status code
+    assert reporting_api.to_api_response(res) == {
+        "success": False,
+        "error": "unknown_error",
+    }
+
+
+def test_to_api_response_valid_json(reporting_api):
+    res = Response()
+    res.status_code = 200
+    res._content = b'{"key": "value"}'  # Simulating valid JSON response
+    assert reporting_api.to_api_response(res) == {"key": "value"}

aikido_firewall/background_process/heartbeats.py

@@ -0,0 +1,30 @@
+"""
+The code to send out a heartbeat is in here
+"""
+
+from aikido_firewall.helpers.logging import logger
+
+
+def send_heartbeats_every_x_secs(reporter, interval_in_secs, s):
+    """
+    Start sending out heartbeats every x seconds
+    """
+    if reporter.serverless:
+        logger.debug("Running in serverless environment, not starting heartbeats")
+        return
+    if not reporter.token:
+        logger.debug("No token provided, not starting heartbeats")
+        return
+
+    logger.debug("Starting heartbeats")
+
+    s.enter(0, 1, send_heartbeat_wrapper, (reporter, interval_in_secs, s))
+
+
+def send_heartbeat_wrapper(rep, interval_in_secs, s):
+    """
+    Wrapper function for send_heartbeat so we get an interval
+    """
+    s.enter(interval_in_secs, 1, send_heartbeat_wrapper, (rep, interval_in_secs, s))
+    logger.debug("Heartbeat...")
+    rep.send_heartbeat()

aikido_firewall/background_process/reporter.py

@@ -0,0 +1,172 @@
+""" This file simply exports the Reporter class"""
+
+import time
+import socket
+import platform
+import json
+from copy import deepcopy
+from aikido_firewall.helpers.logging import logger
+from aikido_firewall.helpers.limit_length_metadata import limit_length_metadata
+from aikido_firewall.helpers.token import Token
+from aikido_firewall import PKG_VERSION
+from aikido_firewall.background_process.heartbeats import send_heartbeats_every_x_secs
+
+
+class Reporter:
+    """Reporter class"""
+
+    timeout_in_sec = 5
+    heartbeat_secs = 10
+
+    def __init__(self, block, api, token, serverless, event_scheduler):
+        self.block = block
+        self.api = api
+        self.token = token  # Should be instance of the Token class!
+
+        if isinstance(serverless, str) and len(serverless) == 0:
+            raise ValueError("Serverless cannot be an empty string")
+        self.serverless = serverless
+
+        self.on_start()
+        send_heartbeats_every_x_secs(self, self.heartbeat_secs, event_scheduler)
+
+    def on_detected_attack(self, attack, context):
+        """
+        This will send something to the API when an attack is detected
+        """
+        if not self.token:
+            return
+        # Modify attack so we can send it out :
+        try:
+            attack["user"] = None
+            attack["payload"] = json.dumps(attack["payload"])[:4096]
+            attack["metadata"] = limit_length_metadata(attack["metadata"], 4096)
+            attack["blocked"] = self.block
+
+            payload = {
+                "type": "detected_attack",
+                "time": get_unixtime_ms(),
+                "agent": self.get_reporter_info(),
+                "attack": attack,
+                "request": {
+                    "method": context.method,
+                    "url": context.url,
+                    "ipAddress": context.remote_address,
+                    "userAgent": get_ua_from_context(context),
+                    "body": context.body,
+                    "headers": context.headers,
+                    "source": context.source,
+                    "route": context.route,
+                },
+            }
+            logger.debug(json.dumps(payload))
+            result = self.api.report(
+                self.token,
+                payload,
+                self.timeout_in_sec,
+            )
+            logger.debug("Result : %s", result)
+        except Exception as e:
+            logger.debug(e)
+            logger.info("Failed to report attack")
+
+    def send_heartbeat(self):
+        """
+        This will send a heartbeat to the server
+        """
+        if not self.token:
+            return
+        logger.debug("Aikido Reporter : Sending out heartbeat")
+        res = self.api.report(
+            self.token,
+            {
+                "type": "heartbeat",
+                "time": get_unixtime_ms(),
+                "agent": self.get_reporter_info(),
+                "stats": {
+                    "sinks": {},
+                    "startedAt": 0,
+                    "endedAt": 0,
+                    "requests": {
+                        "total": 0,
+                        "aborted": 0,
+                        "attacksDetected": {
+                            "total": 0,
+                            "blocked": 0,
+                        },
+                    },
+                },
+                "hostnames": [],
+                "routes": [],
+                "users": [],
+            },
+            self.timeout_in_sec,
+        )
+        self.update_service_config(res)
+
+    def on_start(self):
+        """
+        This will send out an Event signalling the start to the server
+        """
+        if not self.token:
+            return
+        res = self.api.report(
+            self.token,
+            {
+                "type": "started",
+                "time": get_unixtime_ms(),
+                "agent": self.get_reporter_info(),
+            },
+            self.timeout_in_sec,
+        )
+        self.update_service_config(res)
+
+    def get_reporter_info(self):
+        """
+        This returns info about the reporter
+        """
+        return {
+            "dryMode": not self.block,
+            "hostname": socket.gethostname(),
+            "version": PKG_VERSION,
+            "library": "firewall_python",
+            "ipAddress": get_ip(),
+            "packages": [],
+            "serverless": bool(self.serverless),
+            "stack": [],
+            "os": {"name": platform.system(), "version": platform.release()},
+            "preventedPrototypePollution": False,  # Get this out of the API maybe?
+            "nodeEnv": "",
+        }
+
+    def update_service_config(self, res):
+        """
+        Update configuration based on the server's response
+        """
+        if res["success"] is False:
+            logger.debug(res)
+            return
+        if "block" in res.keys() and res["block"] != self.block:
+            logger.debug("Updating blocking, setting blocking to : %s", res["block"])
+            self.block = bool(res["block"])
+
+
+def get_unixtime_ms():
+    """Get the current unix time but in ms"""
+    return int(time.time() * 1000)
+
+
+def get_ip():
+    """Tries to fetch the IP and returns 0.0.0.0 on failure"""
+    try:
+        return socket.gethostbyname(socket.gethostname())
+    except Exception:
+        return "0.0.0.0"
+
+
+def get_ua_from_context(context):
+    """Tries to retrieve the user agent from context"""
+    for k, v in context.headers.items():
+        if k.lower() == "user-agent":
+            return v
+    return None

aikido_firewall/context/__init__.py

@@ -5,7 +5,8 @@
 import threading
 from urllib.parse import parse_qs
 from http.cookies import SimpleCookie
-
+from aikido_firewall.helpers.build_route_from_url import build_route_from_url
+from aikido_firewall.helpers.get_subdomains_from_url import get_subdomains_from_url
 
 SUPPORTED_SOURCES = ["django", "flask", "django-gunicorn"]
 UINPUT_SOURCES = ["body", "cookies", "query", "headers"]
@@ -51,7 +52,11 @@ class Context:
     for vulnerability detection
     """
 
-    def __init__(self, req, source):
+    def __init__(self, context_obj=None, req=None, source=None):
+        if context_obj:
+            self.__dict__.update(context_obj)
+            return
+
         if not source in SUPPORTED_SOURCES:
             raise ValueError(f"Source {source} not supported")
         self.source = source
@@ -63,6 +68,8 @@ def __init__(self, req, source):
             self.set_django_attrs(req)
         elif source == "django-gunicorn":
             self.set_django_gunicorn_attrs(req)
+        self.route = build_route_from_url(self.url)
+        self.subdomains = get_subdomains_from_url(self.url)
 
     def set_django_gunicorn_attrs(self, req):
         """Set properties that are specific to django-gunicorn"""
@@ -93,14 +100,20 @@ def __reduce__(self):
         return (
             self.__class__,
             (
-                self.method,
-                self.remote_address,
-                self.url,
-                self.body,
-                self.headers,
-                self.query,
-                self.cookies,
-                self.source,
+                {
+                    "method": self.method,
+                    "remote_address": self.remote_address,
+                    "url": self.url,
+                    "body": self.body,
+                    "headers": self.headers,
+                    "query": self.query,
+                    "cookies": self.cookies,
+                    "source": self.source,
+                    "route": self.route,
+                    "subdomains": self.subdomains,
+                },
+                None,
+                None,
             ),
         )