Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions .github/workflows/unit-test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,10 @@ jobs:
- name: Add local.aikido.io to /etc/hosts
run: |
sudo echo "127.0.0.1 local.aikido.io" | sudo tee -a /etc/hosts
- name: Build and start aikido mock server
run: |
cd end2end/server && docker build -t mock_core .
docker run --name mock_core -d -p 5050:5000 mock_core
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
with:
Expand Down
12 changes: 7 additions & 5 deletions aikido_zen/background_process/api/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,15 @@
init.py file for api/ folder. Includes abstract class ReportingApi
"""

import json
from aikido_zen.background_process.requests.make_request import Response
from aikido_zen.helpers.logging import logger


class ReportingApi:
"""This is the super class for the reporting API's"""

def to_api_response(self, res):
@staticmethod
def to_api_response(res: Response):
"""Converts results into an Api response obj"""
status = res.status_code
if status == 429:
Expand All @@ -18,10 +19,11 @@ def to_api_response(self, res):
return {"success": False, "error": "invalid_token"}
elif status == 200:
try:
return json.loads(res.text)
return res.json()
except Exception as e:
logger.debug(e)
logger.debug(res.text)
logger.debug(
"Trying to load json, failed: %s (body=%s)", str(e), res.text
)
return {"success": False, "error": "unknown_error"}

def report(self, token, event, timeout_in_sec):
Expand Down
13 changes: 6 additions & 7 deletions aikido_zen/background_process/api/http_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,9 @@
Exports the HTTP API class
"""

import requests
import aikido_zen.background_process.requests as requests
from aikido_zen.background_process.api import ReportingApi
from aikido_zen.background_process.requests.errors import TimeoutExceeded
from aikido_zen.helpers.logging import logger


Expand All @@ -21,11 +22,10 @@ def report(self, token, event, timeout_in_sec):
timeout=timeout_in_sec,
headers=get_headers(token),
)
except requests.exceptions.ConnectionError as e:
logger.error(e)
except TimeoutExceeded as e:
return {"success": False, "error": "timeout"}
except Exception as e:
logger.error(e)
logger.error("Failed to report event : %s(%s)", str(e.__class__), str(e))
return {"success": False, "error": "unknown"}
return self.to_api_response(res)

Expand All @@ -47,11 +47,10 @@ def fetch_firewall_lists(self, token):
"Authorization": str(token),
},
)
except requests.exceptions.ConnectionError as e:
logger.error(e)
except TimeoutExceeded as e:
return {"success": False, "error": "timeout"}
except Exception as e:
logger.error(e)
logger.error("Failed to fetch firewall lists: %s", str(e))
return {"success": False, "error": "unknown"}
return self.to_api_response(res)

Expand Down
Original file line number Diff line number Diff line change
@@ -1,15 +1,14 @@
import pytest
from unittest.mock import patch
from aikido_zen.background_process.api.http_api import ReportingApiHTTP
from aikido_zen.helpers.get_current_unixtime_ms import get_unixtime_ms
from .http_api_ratelimited import ReportingApiHTTPRatelimited


@pytest.fixture
def reporting_api():
"""Fixture to create an instance of ReportingApiHTTPRatelimited."""
return ReportingApiHTTPRatelimited(
reporting_url="http://example.com",
reporting_url="https://example.com",
max_events_per_interval=3,
interval_in_ms=10000,
)
Expand Down Expand Up @@ -150,7 +149,7 @@ def test_report_event_expiry(reporting_api):
):
event3 = {"type": "detected_attack", "time": 11000}
response = reporting_api.report(token="token", event=event3, timeout_in_sec=5)
assert response == {"error": "timeout", "success": False}
assert response == {"error": "unknown", "success": False}
assert (
len(reporting_api.events) == 1
) # One event should have expired, and the new one is added
Expand All @@ -167,7 +166,7 @@ def test_report_event_at_boundary(reporting_api):
):
event2 = {"type": "detected_attack", "time": 10000} # Exactly at the boundary
response = reporting_api.report(token="token", event=event2, timeout_in_sec=5)
assert response == {"error": "timeout", "success": False}
assert response == {"error": "unknown", "success": False}
assert (
len(reporting_api.events) == 2
) # Should be added since it's at the boundary
Expand Down
96 changes: 42 additions & 54 deletions aikido_zen/background_process/api/http_api_test.py
Original file line number Diff line number Diff line change
@@ -1,72 +1,60 @@
import pytest
import requests
from unittest.mock import patch
from aikido_zen.background_process.api.http_api import (
ReportingApiHTTP,
) # Replace with the actual module name
from aikido_zen.background_process.api.http_api import ReportingApiHTTP

# 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/")
def test_report_data_401_code():
api = ReportingApiHTTP("https://guard.aikido.dev/")

# 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
response = api.report("wrong_token", sample_event, 5)

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
),
)
def test_report_local_valid():
api = ReportingApiHTTP("http://localhost:5050/")

# 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/")
assert response == {
"success": True,
"blockedUserIds": [],
"endpoints": [
{
"forceProtectionOff": False,
"graphql": False,
"method": "*",
"rateLimiting": {
"enabled": True,
"maxRequests": 2,
"windowSizeInMS": 5000,
},
"route": "/test_ratelimiting_1",
}
],
"receivedAnyStats": False,
}


def test_report_local_timeout():
api = ReportingApiHTTP("http://localhost:5050/timeout5/")

response = api.report("mocked_token", sample_event, 4)

# Mock the requests.post method to raise a generic exception
def mock_post(url, json, timeout, headers):
raise Exception("Some error occurred")
assert response == {"success": False, "error": "timeout"}

monkeypatch.setattr(requests, "post", mock_post)

# Call the report method
response = api.report("mocked_token", sample_event, 5)
def test_local_gzip():
api = ReportingApiHTTP("http://localhost:5050/")
response = api.fetch_firewall_lists("token")

# Assert that the response is None (or however you want to handle it)
assert response["error"] is "unknown"
assert not response["success"]
assert response == {
"success": True,
"allowedIPAddresses": [],
"blockedIPAddresses": [
{"description": "geo restrictions", "ips": ["1.2.3.4"], "source": "geoip"}
],
"blockedUserAgents": "",
}
2 changes: 1 addition & 1 deletion aikido_zen/background_process/realtime/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
"""

import os
import requests
import aikido_zen.background_process.requests as requests
from aikido_zen.helpers.logging import logger
from aikido_zen.helpers.urls.get_api_url import get_api_url

Expand Down
19 changes: 19 additions & 0 deletions aikido_zen/background_process/requests/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
from aikido_zen.background_process.requests.make_request import make_request
from json import dumps as json_dumps


def get(url, headers, timeout):
return make_request(method="GET", url=url, headers=headers, timeout=timeout)


def post(url, json, headers, timeout):
data = None
if headers is None:
headers = {}
if json is not None:
data = json_dumps(json).encode("utf-8")
headers["Content-Type"] = "application/json"
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Function modifies the input headers parameter directly. More info


return make_request(
method="POST", url=url, data=data, headers=headers, timeout=timeout
)
3 changes: 3 additions & 0 deletions aikido_zen/background_process/requests/errors/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
class TimeoutExceeded(Exception):
def __init__(self):
super().__init__("Aikido internal requests library: timeout was exceeded.")
47 changes: 47 additions & 0 deletions aikido_zen/background_process/requests/make_request.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
import gzip
import json
import socket
import urllib.request
import urllib.error

from aikido_zen.background_process.requests.errors import TimeoutExceeded


def make_request(method, url, timeout, data=None, headers=None):
req = urllib.request.Request(url, data=data, method=method)

# Add headers
if headers:
for key, value in headers.items():
req.add_header(key, value)
try:
with urllib.request.urlopen(req, timeout=timeout) as response:
return Response(response)
except urllib.error.HTTPError as e:
return FailedResponse(status_code=e.code)
except socket.timeout:
raise TimeoutExceeded()


class Response:
def __init__(self, response):
self.status_code = response.getcode()
if response.info().get("Content-Encoding") == "gzip":
# we need to do the gzip decoding ourselves, since urllib has no native support.
self.text = decode_gzip(response.read())
else:
self.text = response.read().decode("utf-8")

def json(self):
return json.loads(self.text)


class FailedResponse:
def __init__(self, status_code):
self.status_code = status_code
self.text = ""


def decode_gzip(raw_bytes, encoding="utf-8"):
f = gzip.decompress(raw_bytes)
return f.decode(encoding)
6 changes: 6 additions & 0 deletions end2end/server/mock_aikido_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,12 @@ def mock_set_config():
def mock_get_events():
return jsonify(events)

@app.route('/timeout5/api/runtime/events', methods=['POST'])
def mock_timeout_5_secs():
time.sleep(5)
return jsonify({"success": True})



@app.route('/mock/reset', methods=['GET'])
def mock_reset_events():
Expand Down
16 changes: 8 additions & 8 deletions poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading
Loading