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/aikido_background_process.py

@@ -6,11 +6,17 @@
 import os
 import time
 import signal
+import sched
 from threading import Thread
 from queue import Queue
 from aikido_firewall.helpers.logging import logger
+from aikido_firewall.background_process.reporter import Reporter
+from aikido_firewall.helpers.should_block import should_block
+from aikido_firewall.helpers.token import get_token_from_env
+from aikido_firewall.background_process.api.http_api import ReportingApiHTTP
 
-REPORT_SEC_INTERVAL = 600  # 10 minutes
+
+EMPTY_QUEUE_INTERVAL = 5  # 5 seconds
 
 
 class AikidoBackgroundProcess:
@@ -31,6 +37,7 @@ def __init__(self, address, key):
             pid = os.getpid()
             os.kill(pid, signal.SIGTERM)  # Kill this subprocess
         self.queue = Queue()
+        self.reporter = None
         # Start reporting thread :
         Thread(target=self.reporting_thread).start()
 
@@ -52,21 +59,40 @@ def __init__(self, address, key):
                     conn.close()
                     pid = os.getpid()
                     os.kill(pid, signal.SIGTERM)  # Kill this subprocess
+                elif data[0] == "READ_PROPERTY":  # meant to get config props
+                    if hasattr(self.reporter, data[1]):
+                        conn.send(self.reporter.__dict__[data[1]])
 
     def reporting_thread(self):
         """Reporting thread"""
         logger.debug("Started reporting thread")
-        while True:
-            self.send_to_reporter()
-            time.sleep(REPORT_SEC_INTERVAL)
+        event_scheduler = sched.scheduler(
+            time.time, time.sleep
+        )  # Create an event scheduler
+        self.send_to_reporter(event_scheduler)
+
+        api = ReportingApiHTTP("http://app.local.aikido.io/")
+        # We need to pass along the scheduler so that the heartbeat also gets sent
+        self.reporter = Reporter(
+            block=should_block(),
+            api=api,
+            token=get_token_from_env(),
+            serverless=False,
+            event_scheduler=event_scheduler,
+        )
+
+        event_scheduler.run()
 
-    def send_to_reporter(self):
+    def send_to_reporter(self, event_scheduler):
         """
         Reports the found data to an Aikido server
         """
-        items_to_report = []
+        # Add back to event scheduler in EMPTY_QUEUE_INTERVAL secs :
+        event_scheduler.enter(
+            EMPTY_QUEUE_INTERVAL, 1, self.send_to_reporter, (event_scheduler,)
+        )
+        logger.debug("Checking queue")
         while not self.queue.empty():
-            items_to_report.append(self.queue.get())
-        logger.debug("Reporting to aikido server")
-        logger.critical("Items to report : %s", items_to_report)
-        # Currently not making API calls
+            attack = self.queue.get()
+            logger.debug("Reporting attack : %s", attack)
+            self.reporter.on_detected_attack(attack[0], attack[1])

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,34 @@
+"""
+Exports the HTTP API class
+"""
+
+import requests
+from aikido_firewall.background_process.api import ReportingApi
+from aikido_firewall.helpers.logging import logger
+
+
+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:
+            logger.error(e)
+            return {"success": False, "error": "unknown"}
+        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/http_api_test.py

@@ -0,0 +1,72 @@
+import pytest
+import requests
+from unittest.mock import patch
+from aikido_firewall.background_process.api.http_api import (
+    ReportingApiHTTP,
+)  # Replace with the actual module name
+
+# Sample event data for testing
+sample_event = {"event_type": "test_event", "data": {"key": "value"}}
+
+
+def test_report_data_401_code(monkeypatch):
+    # Create an instance of ReportingApiHTTP
+    api = ReportingApiHTTP("http://mocked-url.com/")
+
+    # Mock the requests.post method to return a successful response
+    class MockResponse:
+        def __init__(self, json_data, status_code):
+            self.json_data = json_data
+            self.status_code = status_code
+
+        def json(self):
+            return self.json_data
+
+    def mock_post(url, json, timeout, headers):
+        return MockResponse({"success": False}, 401)
+
+    monkeypatch.setattr(requests, "post", mock_post)
+
+    # Call the report method
+    response = api.report("mocked_token", sample_event, 5)
+
+    # Assert the response
+    assert response == {"success": False, "error": "invalid_token"}
+
+
+def test_report_connection_error(monkeypatch):
+    # Create an instance of ReportingApiHTTP
+    api = ReportingApiHTTP("http://mocked-url.com/")
+
+    # Mock the requests.post method to raise a ConnectionError
+    monkeypatch.setattr(
+        requests,
+        "post",
+        lambda *args, **kwargs: (_ for _ in ()).throw(
+            requests.exceptions.ConnectionError
+        ),
+    )
+
+    # Call the report method
+    response = api.report("mocked_token", sample_event, 5)
+
+    # Assert the response
+    assert response == {"success": False, "error": "timeout"}
+
+
+def test_report_other_exception(monkeypatch):
+    # Create an instance of ReportingApiHTTP
+    api = ReportingApiHTTP("http://mocked-url.com/")
+
+    # Mock the requests.post method to raise a generic exception
+    def mock_post(url, json, timeout, headers):
+        raise Exception("Some error occurred")
+
+    monkeypatch.setattr(requests, "post", mock_post)
+
+    # Call the report method
+    response = api.report("mocked_token", sample_event, 5)
+
+    # Assert that the response is None (or however you want to handle it)
+    assert response["error"] is "unknown"
+    assert not response["success"]

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/comms.py

@@ -81,3 +81,15 @@ def target(address, key, data_array):
         t.start()
         # This joins the thread for 3 seconds, afterwards the thread is forced to close (daemon=True)
         t.join(timeout=3)
+
+    def poll_config(self, prop):
+        """
+        This will poll the config from the Background Process
+        """
+        conn = con.Client(self.address, authkey=self.key)
+        conn.send(("READ_PROPERTY", prop))
+        prop_value = conn.recv()
+        conn.send(("CLOSE", {}))
+        conn.close()
+        logger.debug("Received property %s as %s", prop, prop_value)
+        return prop_value

aikido_firewall/background_process/heartbeats.py

@@ -0,0 +1,37 @@
+"""
+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, event_scheduler):
+    """
+    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")
+
+    event_scheduler.enter(
+        0, 1, send_heartbeat_wrapper, (reporter, interval_in_secs, event_scheduler)
+    )
+
+
+def send_heartbeat_wrapper(rep, interval_in_secs, event_scheduler):
+    """
+    Wrapper function for send_heartbeat so we get an interval
+    """
+    event_scheduler.enter(
+        interval_in_secs,
+        1,
+        send_heartbeat_wrapper,
+        (rep, interval_in_secs, event_scheduler),
+    )
+    logger.debug("Heartbeat...")
+    rep.send_heartbeat()

aikido_firewall/background_process/heartbeats_test.py

@@ -0,0 +1,62 @@
+import pytest
+from unittest.mock import Mock, patch
+from aikido_firewall.background_process.heartbeats import (
+    send_heartbeats_every_x_secs,
+    send_heartbeat_wrapper,
+)
+
+
+def test_send_heartbeats_serverless():
+    reporter = Mock()
+    reporter.serverless = True
+    reporter.token = "mocked_token"
+    event_scheduler = Mock()
+
+    with patch("aikido_firewall.helpers.logging.logger.debug") as mock_debug:
+        send_heartbeats_every_x_secs(reporter, 5, event_scheduler)
+
+    mock_debug.assert_called_once_with(
+        "Running in serverless environment, not starting heartbeats"
+    )
+    event_scheduler.enter.assert_not_called()
+
+
+def test_send_heartbeats_no_token():
+    reporter = Mock()
+    reporter.serverless = False
+    reporter.token = None
+    event_scheduler = Mock()
+
+    with patch("aikido_firewall.helpers.logging.logger.debug") as mock_debug:
+        send_heartbeats_every_x_secs(reporter, 5, event_scheduler)
+
+    mock_debug.assert_called_once_with("No token provided, not starting heartbeats")
+    event_scheduler.enter.assert_not_called()
+
+
+def test_send_heartbeats_success():
+    reporter = Mock()
+    reporter.serverless = False
+    reporter.token = "mocked_token"
+    event_scheduler = Mock()
+
+    with patch("aikido_firewall.helpers.logging.logger.debug") as mock_debug:
+        send_heartbeats_every_x_secs(reporter, 5, event_scheduler)
+
+    mock_debug.assert_called_with("Starting heartbeats")
+    event_scheduler.enter.assert_called_once_with(
+        0, 1, send_heartbeat_wrapper, (reporter, 5, event_scheduler)
+    )
+
+
+def test_send_heartbeat_wrapper():
+    reporter = Mock()
+    reporter.send_heartbeat = Mock()
+    event_scheduler = Mock()
+
+    send_heartbeat_wrapper(reporter, 5, event_scheduler)
+
+    reporter.send_heartbeat.assert_called_once()
+    event_scheduler.enter.assert_called_once_with(
+        5, 1, send_heartbeat_wrapper, (reporter, 5, event_scheduler)
+    )