Removing mcp context to functions not needed and improved descriptions

Signed-off-by: S3B4SZ17 <[email protected]>

Author: S3B4SZ17

tests/conftest.py

@@ -1,17 +1,11 @@
+"""
+Module for general utilities and fixtures used in tests.
+"""
+
 import json
 import pytest
 import os
 from unittest.mock import patch
-from fastmcp import Context
-from fastmcp import FastMCP
-
-
-class MockMCP(FastMCP):
-    """
-    Mock class for FastMCP
-    """
-
-    pass
 
 
 def util_load_json(path):
@@ -43,21 +37,6 @@ def mock_success_response():
         patch.stopall()
 
 
-@pytest.fixture
-def mock_ctx():
-    """
-    Fixture to create a mock Context object with 'fastmcp' tags.
-    Returns:
-        Context: A mocked Context object with 'fastmcp' tags.
-    """
-
-    fastmcp: MockMCP = MockMCP(
-        tags=["sysdig", "mcp", "stdio"],
-    )
-    ctx = Context(fastmcp=fastmcp)
-    return ctx
-
-
 @pytest.fixture
 def mock_creds():
     """

tests/events_feed_test.py

@@ -8,8 +8,6 @@
 from unittest.mock import MagicMock, AsyncMock
 import os
 
-from fastmcp import Context
-
 # Get the absolute path of the current module file
 module_path = os.path.abspath(__file__)
 
@@ -19,11 +17,11 @@
 EVENT_INFO_RESPONSE = util_load_json(f"{module_directory}/test_data/events_feed/event_info_response.json")
 
 
-def test_get_event_info(mock_success_response: MagicMock | AsyncMock, mock_ctx: Context, mock_creds) -> None:
+def test_get_event_info(mock_success_response: MagicMock | AsyncMock, mock_creds) -> None:
     """Test the get_event_info tool method.
     Args:
         mock_success_response (MagicMock | AsyncMock): Mocked response object.
-        mock_ctx (Context): Mocked Context object.
+        mock_creds: Mocked credentials.
     """
 
     # Successful response
@@ -32,7 +30,7 @@ def test_get_event_info(mock_success_response: MagicMock | AsyncMock, mock_ctx:
 
     tools_client = EventsFeedTools()
     # Pass the mocked Context object
-    result: dict = tools_client.tool_get_event_info("12345", mock_ctx)
+    result: dict = tools_client.tool_get_event_info("12345")
     results: dict = result["results"]
 
     assert result.get("status_code") == HTTPStatus.OK

tools/events_feed/tool.py

@@ -11,7 +11,6 @@
 from datetime import datetime
 from typing import Optional, Annotated, Any, Dict
 from pydantic import Field
-from fastmcp import Context
 from sysdig_client import ApiException
 from fastmcp.prompts.prompt import PromptMessage, TextContent
 from fastmcp.exceptions import ToolError
@@ -38,21 +37,21 @@ class EventsFeedTools:
     This class provides methods to retrieve event information and list runtime events.
     """
 
-    def init_client(self, config_tags: set[str], old_api: bool = False) -> SecureEventsApi | OldSysdigApi:
+    def init_client(self, old_api: bool = False) -> SecureEventsApi | OldSysdigApi:
         """
         Initializes the SecureEventsApi client from the request state.
         If the request does not have the API client initialized, it will create a new instance
         using the Sysdig Secure token and host from the environment variables.
         Args:
-            config_tags (set[str]): The tags associated with the MCP server configuration, used to determine the transport mode.
+            old_api (bool): If True, initializes the OldSysdigApi client instead of SecureEventsApi.
         Returns:
             SecureEventsApi | OldSysdigApi: An instance of the SecureEventsApi or OldSysdigApi client.
         Raises:
             ValueError: If the SYSDIG_SECURE_TOKEN environment variable is not set.
         """
         secure_events_api: SecureEventsApi = None
         old_sysdig_api: OldSysdigApi = None
-        if "streamable-http" in config_tags:
+        if app_config.get("mcp", {}).get("transport", "") == "streamable-http":
             # Try to get the HTTP request
             log.debug("Attempting to get the HTTP request to initialize the Sysdig API client.")
             request: Request = get_http_request()
@@ -77,7 +76,7 @@ def init_client(self, config_tags: set[str], old_api: bool = False) -> SecureEve
             return old_sysdig_api
         return secure_events_api
 
-    def tool_get_event_info(self, event_id: str, ctx: Context) -> dict:
+    def tool_get_event_info(self, event_id: str) -> dict:
         """
         Retrieves detailed information for a specific security event.
 
@@ -88,7 +87,7 @@ def tool_get_event_info(self, event_id: str, ctx: Context) -> dict:
             Event: The Event object containing detailed information about the specified event.
         """
         # Init of the sysdig client
-        secure_events_api = self.init_client(config_tags=ctx.fastmcp.tags)
+        secure_events_api = self.init_client()
         try:
             # Get the HTTP request
             start_time = time.time()
@@ -106,7 +105,6 @@ def tool_get_event_info(self, event_id: str, ctx: Context) -> dict:
 
     def tool_list_runtime_events(
         self,
-        ctx: Context,
         cursor: Optional[str] = None,
         scope_hours: int = 1,
         limit: int = 50,
@@ -148,15 +146,15 @@ def tool_list_runtime_events(
             cursor (Optional[str]): Cursor for pagination.
             scope_hours (int): Number of hours back from now to include events. Defaults to 1.
             severity_level (Optional[str]): One of "info", "low", "medium", "high". If provided, filters by that severity.
-            If None, includes all severities.
+                If None, includes all severities.
             cluster_name (Optional[str]): Name of the Kubernetes cluster to filter events. If None, includes all clusters.
             limit (int): Maximum number of events to return. Defaults to 50.
             filter_expr (Optional[str]): An optional filter expression to further narrow down events.
 
         Returns:
-            List[Event]: A list of Event objects matching the criteria.
+            dict: A dictionary containing the results of the runtime events query, including pagination information.
         """
-        secure_events_api = self.init_client(config_tags=ctx.fastmcp.tags)
+        secure_events_api = self.init_client()
         start_time = time.time()
         # Compute time window
         now_ns = time.time_ns()
@@ -188,13 +186,12 @@ def tool_list_runtime_events(
 
     # A tool to retrieve all the process-tree information for a specific event.Add commentMore actions
 
-    def tool_get_event_process_tree(self, ctx: Context, event_id: str) -> Dict[str, Any]:
+    def tool_get_event_process_tree(self, event_id: str) -> dict:
         """
         Retrieves the process tree for a specific security event.
         Not every event has a process tree, so this may return an empty tree.
 
         Args:
-            ctx (Context): The context object containing request-specific information.
             event_id (str): The unique identifier of the security event.
 
         Returns:
@@ -203,7 +200,7 @@ def tool_get_event_process_tree(self, ctx: Context, event_id: str) -> Dict[str,
         try:
             start_time = time.time()
             # Get process tree branches
-            old_api_client = self.init_client(config_tags=ctx.fastmcp.tags, old_api=True)
+            old_api_client = self.init_client(old_api=True)
             branches = old_api_client.request_process_tree_branches(event_id)
             # Get process tree
             tree = old_api_client.request_process_tree_trees(event_id)

tools/inventory/tool.py

@@ -8,7 +8,6 @@
 from typing import Annotated
 from pydantic import Field
 from fastmcp.server.dependencies import get_http_request
-from fastmcp import Context
 from fastmcp.exceptions import ToolError
 from starlette.requests import Request
 from sysdig_client import ApiException
@@ -32,20 +31,18 @@ class InventoryTools:
     This class provides methods to list resources and retrieve a single resource by its hash.
     """
 
-    def init_client(self, config_tags: set[str]) -> InventoryApi:
+    def init_client(self) -> InventoryApi:
         """
         Initializes the InventoryApi client from the request state.
         If the request does not have the API client initialized, it will create a new instance
         using the Sysdig Secure token and host from the environment variables.
-        Args:
-            config_tags (set[str]): The tags associated with the MCP server configuration, used to determine the transport mode.
         Returns:
             InventoryApi: An instance of the InventoryApi client.
         Raises:
             ValueError: If the SYSDIG_SECURE_TOKEN environment variable is not set.
         """
         secure_events_api: InventoryApi = None
-        if "streamable-http" in config_tags:
+        if app_config.get("mcp", {}).get("transport", "") == "streamable-http":
             # Try to get the HTTP request
             log.debug("Attempting to get the HTTP request to initialize the Sysdig API client.")
             request: Request = get_http_request()
@@ -64,7 +61,6 @@ def init_client(self, config_tags: set[str]) -> InventoryApi:
 
     def tool_list_resources(
         self,
-        ctx: Context,
         filter_exp: Annotated[
             str,
             Field(
@@ -154,20 +150,28 @@ def tool_list_resources(
         List inventory items based on a filter expression, with optional pagination.
 
         Args:
-            filter_exp (str): Sysdig Secure query filter expression.
+            filter_exp (str): Sysdig query filter expression to filter inventory resources.
+                Use the resource://filter-query-language to get the expected filter expression format.
+                Supports operators: =, !=, in, exists, contains, startsWith.
+                Combine with and/or/not.
                 Examples:
-                - not isExposed exists
-                - category in ("IAM") and isExposed exists
-                - category in ("IAM","Audit & Monitoring")
+                - zone in ("zone1") and machineImage = "ami-0b22b359fdfabe1b5"
+                - (projectId = "1235495521" or projectId = "987654321") and vuln.severity in ("Critical")
+                - vuln.name in ("CVE-2023-0049")
+                - vuln.cvssScore >= "3"
+                - container.name in ("sysdig-container") and not labels exists
+                - imageId in ("sha256:3768ff6176e29a35ce1354622977a1e5c013045cbc4f30754ef3459218be8ac")
+                - platform in ("GCP", "AWS", "Azure", "Kubernetes") and isExposed exists
             page_number (int): Page number for pagination (1-based).
             page_size (int): Number of items per page.
             with_enrich_containers (bool): Include enriched container information.
 
         Returns:
-            InventoryResourceResponse: The API response containing inventory items.
+            dict: A dictionary containing the results of the inventory query, including pagination information.
+            Or a dict containing an error message if the call fails.
         """
         try:
-            inventory_api = self.init_client(config_tags=ctx.fastmcp.tags)
+            inventory_api = self.init_client()
             start_time = time.time()
 
             api_response = inventory_api.get_resources_without_preload_content(
@@ -185,7 +189,6 @@ def tool_list_resources(
 
     def tool_get_resource(
         self,
-        ctx: Context,
         resource_hash: Annotated[str, Field(description="The unique hash of the inventory resource to retrieve.")],
     ) -> dict:
         """
@@ -195,11 +198,10 @@ def tool_get_resource(
             resource_hash (str): The hash identifier of the resource.
 
         Returns:
-            InventoryResourceExtended: The detailed resource object.
-            Or a dict containing an error message if the call fails.
+            dict: A dictionary containing the details of the requested inventory resource.
         """
         try:
-            inventory_api = self.init_client(config_tags=ctx.fastmcp.tags)
+            inventory_api = self.init_client()
             start_time = time.time()
 
             api_response = inventory_api.get_resource_without_preload_content(hash=resource_hash)

tools/sysdig_sage/tool.py

@@ -8,7 +8,6 @@
 import os
 import time
 from typing import Any, Dict
-from fastmcp import Context
 from fastmcp.exceptions import ToolError
 from utils.sysdig.old_sysdig_api import OldSysdigApi
 from starlette.requests import Request
@@ -31,21 +30,19 @@ class SageTools:
     language questions and execute them against the Sysdig API.
     """
 
-    def init_client(self, config_tags: set[str]) -> OldSysdigApi:
+    def init_client(self) -> OldSysdigApi:
         """
         Initializes the OldSysdigApi client from the request state.
         If the request does not have the API client initialized, it will create a new instance
         using the Sysdig Secure token and host from the environment variables.
-        Args:
-            config_tags (set[str]): The tags associated with the MCP server configuration, used to determine the transport mode.
         Returns:
             OldSysdigApi: An instance of the OldSysdigApi client.
 
         Raises:
             ValueError: If the SYSDIG_SECURE_TOKEN environment variable is not set.
         """
         old_sysdig_api: OldSysdigApi = None
-        if "streamable-http" in config_tags:
+        if app_config.get("mcp", {}).get("transport", "") == "streamable-http":
             # Try to get the HTTP request
             log.debug("Attempting to get the HTTP request to initialize the Sysdig API client.")
             request: Request = get_http_request()
@@ -63,7 +60,7 @@ def init_client(self, config_tags: set[str]) -> OldSysdigApi:
             old_sysdig_api = OldSysdigApi(api_client)
         return old_sysdig_api
 
-    async def tool_sysdig_sage(self, ctx: Context, question: str) -> Dict[str, Any]:
+    async def tool_sage_to_sysql(self, question: str) -> dict:
         """
         Queries Sysdig Sage with a natural language question, retrieves a SysQL query,
         executes it against the Sysdig API, and returns the results.
@@ -72,20 +69,20 @@ async def tool_sysdig_sage(self, ctx: Context, question: str) -> Dict[str, Any]:
             question (str): A natural language question to send to Sage.
 
         Returns:
-            Dict: JSON-decoded response of the executed SysQL query, or an error object.
+            dict: A dictionary containing the results of the SysQL query execution and the query text.
 
         Raises:
             ToolError: If the SysQL query generation or execution fails.
 
         Examples:
-            # tool_sysdig_sage(question="Match Cloud Resource affected by Critical Vulnerability")
-            # tool_sysdig_sage(question="Match Kubernetes Workload affected by Critical Vulnerability")
-            # tool_sysdig_sage(question="Match AWS EC2 Instance that violates control 'EC2 - Instances should use IMDSv2'")
+            # tool_sage_to_sysql(question="Match Cloud Resource affected by Critical Vulnerability")
+            # tool_sage_to_sysql(question="Match Kubernetes Workload affected by Critical Vulnerability")
+            # tool_sage_to_sysql(question="Match AWS EC2 Instance that violates control 'EC2 - Instances should use IMDSv2'")
         """
         # 1) Generate SysQL query
         try:
             start_time = time.time()
-            old_sysdig_api = self.init_client(config_tags=ctx.fastmcp.tags)
+            old_sysdig_api = self.init_client()
             sysql_response = await old_sysdig_api.generate_sysql_query(question)
             if sysql_response.status > 299:
                 raise ToolError(f"Sysdig Sage returned an error: {sysql_response.status} - {sysql_response.data}")