Ready for pytest

.github/workflows/publish.yml

@@ -0,0 +1,50 @@
+name: Publish Python 🐍 distribution 📦 to TestPyPI
+
+on:
+  push:
+    tags:
+      - 'v*'
+
+jobs:
+  build:
+    name: Build distribution 📦
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Install poetry
+      run: pip install poetry
+    - name: Install dependencies
+      run: poetry install
+
+    - name: Build distribution packages
+      run: poetry build
+
+    - name: Store the distribution packages
+      uses: actions/upload-artifact@v3
+      with:
+        name: python-package-distributions
+        path: dist/
+  publish-to-testpypi:
+    name: Publish Python 🐍 distribution 📦 to TestPyPI
+    needs:
+    - build
+    runs-on: ubuntu-latest
+
+    environment:
+      name: testpypi
+      url: https://test.pypi.org/p/aikido_firewall
+
+    permissions:
+      id-token: write  # IMPORTANT: mandatory for trusted publishing
+
+    steps:
+    - name: Download all the dists
+      uses: actions/download-artifact@v3
+      with:
+        name: python-package-distributions
+        path: dist/
+    - name: Publish distribution 📦 to TestPyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
+      with:
+        repository-url: https://test.pypi.org/legacy/

aikido_firewall/__init__.py

@@ -5,7 +5,7 @@
 from dotenv import load_dotenv
 
 # Constants
-PKG_VERSION = "0.0.1"
+PKG_VERSION = "0.0.7"
 
 # Import logger
 from aikido_firewall.helpers.logging import logger

aikido_firewall/background_process/aikido_background_process.py

@@ -44,33 +44,45 @@
             conn = listener.accept()
             logger.debug("connection accepted from %s", listener.last_accepted)
             while True:
-                data = conn.recv()  #  because of this no sleep needed in thread
-                logger.debug("Incoming data : %s", data)
-                if data[0] == "ATTACK":
-                    self.queue.put(data[1])
-                elif data[0] == "CLOSE":  # this is a kind of EOL for python IPC
-                    conn.close()
-                    break
-                elif (
-                    data[0] == "KILL"
-                ):  # when main process quits , or during testing etc
-                    logger.debug("Killing subprocess")
-                    conn.close()
-                    sys.exit(0)
-                elif data[0] == "READ_PROPERTY":  # meant to get config props
-                    if hasattr(self.reporter, data[1]):
-                        conn.send(self.reporter.__dict__[data[1]])
-                elif data[0] == "ROUTE":
-                    # Called every time the user visits a route
-                    self.reporter.routes.add_route(method=data[1][0], path=data[1][1])
-                elif data[0] == "SHOULD_RATELIMIT":
-                    # Called to check if the context passed along as data should be
-                    # Rate limited
-                    conn.send(
-                        should_ratelimit_request(
-                            context=data[1], reporter=self.reporter
+                try:
+                    data = conn.recv()  #  because of this no sleep needed in thread
+                    logger.debug("Incoming data : %s", data)
+                    if data[0] == "ATTACK":
+                        self.queue.put(data[1])
+                    elif data[0] == "CLOSE":  # this is a kind of EOL for python IPC
+                        conn.close()
+                        break
+                    elif (
+                        data[0] == "KILL"
+                    ):  # when main process quits , or during testing etc
+                        logger.debug("Killing subprocess")
+                        conn.close()
+                        sys.exit(0)
+                    elif data[0] == "READ_PROPERTY":  # meant to get config props
+                        if hasattr(self.reporter, data[1]):
+                            conn.send(self.reporter.__dict__[data[1]])
+                        else:
+                            logger.debug(
+                                "Reporter has no attribute %s, current reporter: %s",
+                                data[1],
+                                self.reporter,
+                            )
+                            conn.send(None)
+                    elif data[0] == "ROUTE":
+                        # Called every time the user visits a route
+                        self.reporter.routes.add_route(
+                            method=data[1][0], path=data[1][1]
                         )
-                    )
+                    elif data[0] == "SHOULD_RATELIMIT":
+                        # Called to check if the context passed along as data should be
+                        # Rate limited
+                        conn.send(
+                            should_ratelimit_request(
+                                context=data[1], reporter=self.reporter
+                            )
+                        )
+                except Exception as e:
+                    logger.error("Exception occured in server thread : %s", e)
 
     def reporting_thread(self):
         """Reporting thread"""
@@ -80,7 +92,7 @@
         )  # Create an event scheduler
         self.send_to_reporter(event_scheduler)
 
-        api = ReportingApiHTTP("http://app.local.aikido.io/")
+        api = ReportingApiHTTP("https://app.aikido.dev/")
         # We need to pass along the scheduler so that the heartbeat also gets sent
         self.reporter = Reporter(
             block=should_block(),

aikido_firewall/background_process/reporter.py

@@ -22,7 +22,7 @@ class Reporter:
     """Reporter class"""
 
     timeout_in_sec = 5  # Timeout of API calls to Aikido Server
-    heartbeat_secs = 600  # Heartbeat every 10 minutes
+    heartbeat_secs = 10  # Heartbeat every 10 minutes
 
     def __init__(self, block, api, token, serverless, event_scheduler):
         self.block = block

aikido_firewall/helpers/blocking_enabled.py

@@ -0,0 +1,13 @@
+"""Helper function file, see function docstring"""
+
+from aikido_firewall.background_process import get_comms
+
+
+def is_blocking_enabled():
+    """
+    Checks with the background process if blocking is enabled
+    """
+    should_block_res = get_comms().send_data_to_bg_process(
+        action="READ_PROPERTY", obj="block", receive=True
+    )
+    return should_block_res["success"] and should_block_res["data"]

aikido_firewall/sinks/os.py

@@ -50,7 +50,6 @@ def generate_aikido_function(op, former_func):
     """
 
     def aikido_new_func(*args, op=op, former_func=former_func, **kwargs):
-        logger.debug("`os` wrapper, filepath : `%s`; OP : `%s`", args[0], op)
         context = get_current_context()
         if not context:
             return former_func(*args, **kwargs)

aikido_firewall/vulnerabilities/path_traversal/check_context_for_path_traversal.py

@@ -22,7 +22,6 @@ def check_context_for_path_traversal(
         return {}
 
     for source in SOURCES:
-        logger.debug("Checking source %s for path traversal", source)
         if hasattr(context, source):
             user_inputs = extract_strings_from_user_input(getattr(context, source))
             for user_input, path in user_inputs.items():

aikido_firewall/vulnerabilities/ssrf/inspect_getaddrinfo_result.py

@@ -8,10 +8,12 @@
 from aikido_firewall.helpers.logging import logger
 from aikido_firewall.background_process import get_comms
 from aikido_firewall.errors import AikidoSSRF
+from aikido_firewall.helpers.blocking_enabled import is_blocking_enabled
 from .imds import is_trusted_hostname, is_imds_ip_address
 from .is_private_ip import is_private_ip
 from .find_hostname_in_context import find_hostname_in_context
 
+
 #  gets called when the result of the DNS resolution has come in
 def inspect_getaddrinfo_result(dns_results, hostname, port):
     """Inspect the results of a getaddrinfo() call"""
@@ -22,17 +24,12 @@
 
     context = get_current_context()
 
-    should_block_res = get_comms().send_data_to_bg_process(
-        action="READ_PROPERTY", obj="block", receive=True
-    )
-    should_block = should_block_res["success"] and should_block_res["data"]
-
     ip_addresses = extract_ip_array_from_results(dns_results)
     if resolves_to_imds_ip(ip_addresses, hostname):
         #  Block stored SSRF attack that target IMDS IP addresses
         #  An attacker could have stored a hostname in a database that points to an IMDS IP address
         #  We don't check if the user input contains the hostname because there's no context
-        if should_block:
+        if is_blocking_enabled():
             raise AikidoSSRF()
 
     if not context:
@@ -46,6 +43,7 @@
     if not found:
         return
 
+    should_block = is_blocking_enabled()
     stack = " ".join(traceback.format_stack())
     attack = {
         "module": "socket",