Aik 3171 : Get context from WSGI to an Agent

aikido_firewall/__init__.py

@@ -4,20 +4,25 @@
 
 from dotenv import load_dotenv
 
-# Import sources
-import aikido_firewall.sources.django
-import aikido_firewall.sources.flask
-
-# Import sinks
-import aikido_firewall.sinks.pymysql
-
-# Import middleware
-import aikido_firewall.middleware.django
 
 # Import logger
 from aikido_firewall.helpers.logging import logger
 
+# Import agent
+from aikido_firewall.agent import start_agent
+
 # Load environment variables
 load_dotenv()
 
-logger.info("Aikido python firewall started")
+
+def protect():
+    """Start Aikido agent"""
+    # Import sources
+    import aikido_firewall.sources.django
+    import aikido_firewall.sources.flask
+
+    # Import sinks
+    import aikido_firewall.sinks.pymysql
+
+    logger.info("Aikido python firewall started")
+    start_agent()

aikido_firewall/agent/__init__.py

@@ -0,0 +1,74 @@
+"""
+Aikido agent, this will create a new thread and listen for stuff sent by our sources and sinks
+"""
+
+import time
+import queue
+from threading import Thread
+from aikido_firewall.helpers.logging import logger
+
+AGENT_SEC_INTERVAL = 60
+
+
+class AikidoThread:
+    """
+    Our agent thread
+    """
+
+    def __init__(self, q):
+        logger.debug("Agent thread started")
+        while True:
+            while not q.empty():
+                self.process_data(q.get())
+        time.sleep(AGENT_SEC_INTERVAL)
+        self.q = q
+        self.current_context = None
+
+    def process_data(self, item):
+        """Will process the data added to the queue"""
+        action, data = item
+        logger.debug("Action %s, Data %s", action, data)
+        if action == "REPORT":
+            logger.debug("Report")
+            self.current_context = data
+        else:
+            logger.error("Action `%s` is not defined. (Aikido Agent)", action)
+
+
+# pylint: disable=invalid-name # This variable does change
+agent = None
+
+
+def get_agent():
+    """Returns the globally stored agent"""
+    return agent
+
+
+def start_agent():
+    """
+    Starts a thread to handle incoming/outgoing data
+    """
+    # pylint: disable=global-statement # We need this to be global
+    global agent
+
+    # This creates a queue for Inter-Process Communication
+    logger.debug("Creating IPC Queue")
+    q = queue.Queue()
+
+    logger.debug("Starting a new agent thread")
+    agent_thread = Thread(target=AikidoThread, args=(q,))
+    agent_thread.start()
+    agent = Agent(q)
+
+
+class Agent:
+    """Agent class"""
+
+    def __init__(self, q):
+        self.q = q
+
+    def report(self, obj, action):
+        """
+        Report something to the agent
+        """
+        self.q.put((action, obj))

aikido_firewall/context/__init__.py

@@ -0,0 +1,52 @@
+"""
+Provides all the functionality for contexts
+"""
+
+# pylint: disable=invalid-name # This is not a const
+current_context = None
+
+
+def get_current_context():
+    """Returns the current context"""
+    return current_context
+
+
+class Context:
+    """
+    A context object, it stores everything that is important
+    for vulnerability detection
+    """
+
+    def __init__(self, req):
+        self.method = req.method
+        self.remote_address = req.remote_addr
+        self.url = req.url
+        self.body = req.form
+        self.headers = req.headers
+        self.query = req.args
+        self.cookies = req.cookies
+        self.source = "flask"
+
+    def __reduce__(self):
+        return (
+            self.__class__,
+            (
+                self.method,
+                self.remote_address,
+                self.url,
+                self.body,
+                self.headers,
+                self.query,
+                self.cookies,
+                self.source,
+            ),
+        )
+
+    def set_as_current_context(self):
+        """
+        Set the current context
+        """
+        # pylint: disable=global-statement # We need this so that
+        # Sinks can access context
+        global current_context
+        current_context = self

aikido_firewall/helpers/build_route_from_url.py

@@ -0,0 +1,82 @@
+"""
+Module with logic to build a route, i.e. find route params
+from a simple URL string
+"""
+
+import re
+import ipaddress
+from aikido_firewall.helpers.looks_like_a_secret import looks_like_a_secret
+from aikido_firewall.helpers.try_parse_url_path import try_parse_url_path
+
+UUID_REGEX = re.compile(
+    r"(?:[0-9a-f]{8}-[0-9a-f]{4}-[1-8][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}|00000000-0000-0000-0000-000000000000|ffffffff-ffff-ffff-ffff-ffffffffffff)$",
+    re.I,
+)
+NUMBER_REGEX = re.compile(r"^\d+$")
+DATE_REGEX = re.compile(r"^\d{4}-\d{2}-\d{2}|\d{2}-\d{2}-\d{4}$")
+EMAIL_REGEX = re.compile(
+    r"^[a-zA-Z0-9.!#$%&\'*+/=?^_`{|}~-]+@[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?(?:\.[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?)*$"
+)
+HASH_REGEX = re.compile(
+    r"^(?:[a-f0-9]{32}|[a-f0-9]{40}|[a-f0-9]{64}|[a-f0-9]{128})$", re.I
+)
+HASH_LENGTHS = [32, 40, 64, 128]
+
+
+def build_route_from_url(url):
+    """
+    Main helper function which will build the route
+    from a URL string as input
+    """
+    path = try_parse_url_path(url)
+
+    if not path:
+        return None
+
+    route = "/".join(
+        [replace_url_segment_with_param(segment) for segment in path.split("/")]
+    )
+
+    if route == "/":
+        return "/"
+
+    if route.endswith("/"):
+        return route[:-1]
+
+    return route
+
+
+def replace_url_segment_with_param(segment):
+    """
+    ??????????
+    """
+    if not segment:  # Check if segment is empty
+        return segment  # Return the segment as is if it's empty
+    char_code = ord(segment[0])
+    starts_with_number = 48 <= char_code <= 57  # ASCII codes for '0' to '9'
+
+    if starts_with_number and NUMBER_REGEX.match(segment):
+        return ":number"
+
+    if len(segment) == 36 and UUID_REGEX.match(segment):
+        return ":uuid"
+
+    if starts_with_number and DATE_REGEX.match(segment):
+        return ":date"
+
+    if "@" in segment and EMAIL_REGEX.match(segment):
+        return ":email"
+
+    try:
+        ipaddress.ip_address(segment)
+        return ":ip"
+    except ValueError:
+        pass
+
+    if len(segment) in HASH_LENGTHS and HASH_REGEX.match(segment):
+        return ":hash"
+
+    if looks_like_a_secret(segment):
+        return ":secret"
+
+    return segment

aikido_firewall/helpers/build_route_from_url_test.py

@@ -0,0 +1,116 @@
+from aikido_firewall.helpers.build_route_from_url import build_route_from_url
+import pytest
+import hashlib
+
+
+def generate_hash(algorithm):
+    data = "test"
+    if algorithm == "md5":
+        return hashlib.md5(data.encode()).hexdigest()
+    elif algorithm == "sha1":
+        return hashlib.sha1(data.encode()).hexdigest()
+    elif algorithm == "sha256":
+        return hashlib.sha256(data.encode()).hexdigest()
+    elif algorithm == "sha512":
+        return hashlib.sha512(data.encode()).hexdigest()
+    else:
+        return None
+
+
+def test_invalid_urls():
+    assert build_route_from_url("") == None
+    assert build_route_from_url("http") == None
+
+
+def test_root_urls():
+    assert build_route_from_url("/") == "/"
+    assert build_route_from_url("http://localhost/") == "/"
+
+
+def test_replace_numbers():
+    assert build_route_from_url("/posts/3") == "/posts/:number"
+    assert build_route_from_url("http://localhost/posts/3") == "/posts/:number"
+    assert build_route_from_url("http://localhost/posts/3/") == "/posts/:number"
+    assert (
+        build_route_from_url("http://localhost/posts/3/comments/10")
+        == "/posts/:number/comments/:number"
+    )
+    assert (
+        build_route_from_url("/blog/2023/05/great-article")
+        == "/blog/:number/:number/great-article"
+    )
+
+
+def test_replace_dates():
+    assert build_route_from_url("/posts/2023-05-01") == "/posts/:date"
+    assert build_route_from_url("/posts/2023-05-01/") == "/posts/:date"
+    assert (
+        build_route_from_url("/posts/2023-05-01/comments/2023-05-01")
+        == "/posts/:date/comments/:date"
+    )
+    assert build_route_from_url("/posts/01-05-2023") == "/posts/:date"
+
+
+def test_ignore_comma_numbers():
+    assert build_route_from_url("/posts/3,000") == "/posts/3,000"
+
+
+def test_ignore_api_version_numbers():
+    assert build_route_from_url("/v1/posts/3") == "/v1/posts/:number"
+
+
+def test_replace_uuids():
+    uuids = [
+        "d9428888-122b-11e1-b85c-61cd3cbb3210",
+        "000003e8-2363-21ef-b200-325096b39f47",
+        "a981a0c2-68b1-35dc-bcfc-296e52ab01ec",
+        "109156be-c4fb-41ea-b1b4-efe1671c5836",
+        "90123e1c-7512-523e-bb28-76fab9f2f73d",
+        "1ef21d2f-1207-6660-8c4f-419efbd44d48",
+        "017f22e2-79b0-7cc3-98c4-dc0c0c07398f",
+        "0d8f23a0-697f-83ae-802e-48f3756dd581",
+    ]
+    for uuid in uuids:
+        assert build_route_from_url(f"/posts/{uuid}") == "/posts/:uuid"
+
+
+def test_ignore_invalid_uuids():
+    assert (
+        build_route_from_url("/posts/00000000-0000-1000-6000-000000000000")
+        == "/posts/00000000-0000-1000-6000-000000000000"
+    )
+
+
+def test_ignore_strings():
+    assert build_route_from_url("/posts/abc") == "/posts/abc"
+
+
+def test_replace_email_addresses():
+    assert build_route_from_url("/login/john.doe@acme.com") == "/login/:email"
+    assert build_route_from_url("/login/john.doe+alias@acme.com") == "/login/:email"
+
+
+def test_replace_ip_addresses():
+    assert build_route_from_url("/block/1.2.3.4") == "/block/:ip"
+    assert (
+        build_route_from_url("/block/2001:2:ffff:ffff:ffff:ffff:ffff:ffff")
+        == "/block/:ip"
+    )
+    assert build_route_from_url("/block/64:ff9a::255.255.255.255") == "/block/:ip"
+    assert build_route_from_url("/block/100::") == "/block/:ip"
+    assert build_route_from_url("/block/fec0::") == "/block/:ip"
+    assert build_route_from_url("/block/227.202.96.196") == "/block/:ip"
+
+
+def test_replace_hashes():
+    assert build_route_from_url(f"/files/{generate_hash('md5')}") == "/files/:hash"
+    assert build_route_from_url(f"/files/{generate_hash('sha1')}") == "/files/:hash"
+    assert build_route_from_url(f"/files/{generate_hash('sha256')}") == "/files/:hash"
+    assert build_route_from_url(f"/files/{generate_hash('sha512')}") == "/files/:hash"
+
+
+def test_replace_secrets():
+    assert (
+        build_route_from_url("/confirm/CnJ4DunhYfv2db6T1FRfciRBHtlNKOYrjoz")
+        == "/confirm/:secret"
+    )