Patched code sample

import re
import signal

class RegexTimeoutError(Exception):
    """Custom exception to indicate a regex operation timed out."""
    pass

def _handle_timeout(signum, frame):
    """Raises a RegexTimeoutError when the signal is received."""
    raise RegexTimeoutError("Regular expression search timed out due to excessive complexity.")

class MCPToolIndex:
    """
    A mock class representing the PraisonAI tool index. The search_tools method
    is patched to prevent Regular Expression Denial of Service (ReDoS) attacks.
    """
    def __init__(self):
        # A sample list of tools to search through.
        self.tools = [
            "data_analyzer_tool_v1",
            "code_generator_v3",
            "file_system_manager_alpha",
            "api_integration_helper",
            "auth_service_with_many_aaaaaaaaaaaaaaaaaaaaaaaaab_features",
        ]

    def search_tools(self, query: str, timeout: int = 2):
        """
        Searches for tools using a caller-supplied regular expression with a timeout.

        This patched version mitigates CVE-2026-34939 by implementing a timeout
        for the regular expression matching process. This prevents catastrophic
        backtracking from blocking the Python thread and causing a service outage.

        Args:
            query: The regular expression string from the user.
            timeout: The maximum number of seconds the search is allowed to run.

        Returns:
            A list of matching tool names, or an empty list if an error or timeout occurs.
        """
        # Set the signal handler for the alarm signal (SIGALRM).
        # This is the core of the fix: establishing a timeout mechanism.
        try:
            signal.signal(signal.SIGALRM, _handle_timeout)
            signal.alarm(timeout)  # Set the alarm to fire after `timeout` seconds.
        except ValueError:
            # signal.signal may fail if not in the main thread.
            # In a real-world web server, a different async-compatible timeout
            # or a third-party regex library with built-in timeouts (like `regex`)
            # would be a more robust solution. This implementation demonstrates
            # the principle for a single-threaded blocking context.
            print("Warning: Could not set a timeout. Running without protection.")
            pass


        found_tools = []
        try:
            # It's good practice to compile the regex first. This operation can
            # also be complex, so it's included within the timeout block.
            pattern = re.compile(query)
            for tool in self.tools:
                # The search operation is where catastrophic backtracking occurs.
                # With the alarm set, this will be interrupted if it takes too long.
                if pattern.search(tool):
                    found_tools.append(tool)

        except (re.error, RegexTimeoutError) as e:
            # If the regex is invalid (re.error) or times out (RegexTimeoutError),
            # we fail safely by logging the error and returning no results.
            # This prevents the service from crashing or hanging.
            print(f"Aborted search for query '{query}'. Reason: {e}")
            return []
        finally:
            # Crucially, disable the alarm in the `finally` block.
            # This ensures the alarm is cancelled whether the search succeeds,
            # fails, or times out, preventing it from firing later.
            try:
                signal.alarm(0)
            except ValueError:
                pass # Fails if signal was not set; can be ignored.

        return found_tools