Updating the init and conventions config of the mcp server

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

Author: S3B4SZ17

.github/workflows/publish.yaml

@@ -5,6 +5,7 @@ on:
   push:
     branches:
       - main
+      - beta
     paths:
       - pyproject.toml
       - Dockerfile

main.py

@@ -3,7 +3,9 @@
 """
 
 import os
-import asyncio
+import signal
+import sys
+import logging
 from dotenv import load_dotenv
 
 # Application config loader
@@ -12,28 +14,49 @@
 # Register all tools so they attach to the MCP server
 from utils.mcp_server import run_stdio, run_http
 
+# Set up logging
+logging.basicConfig(
+    format="%(asctime)s-%(process)d-%(levelname)s- %(message)s",
+    level=os.environ.get("LOGLEVEL", "ERROR"),
+)
+log = logging.getLogger(__name__)
+
 # Load environment variables from .env
 load_dotenv()
 
 app_config = get_app_config()
 
 
+def handle_signals():
+    def signal_handler(sig, frame):
+        log.info(f"Received signal {sig}, shutting down...")
+        os._exit(0)
+
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
+    signal.signal(signal.SIGHUP, signal_handler)
+
+
 def main():
     # Choose transport: "stdio" or "sse" (HTTP/SSE)
+    handle_signals()
     transport = os.environ.get("MCP_TRANSPORT", app_config["mcp"]["transport"]).lower()
-    print("""
+    log.info("""
     ▄▖     ▌▘    ▖  ▖▄▖▄▖  ▄▖
     ▚ ▌▌▛▘▛▌▌▛▌  ▛▖▞▌▌ ▙▌  ▚ █▌▛▘▌▌█▌▛▘
     ▄▌▙▌▄▌▙▌▌▙▌  ▌▝ ▌▙▖▌   ▄▌▙▖▌ ▚▘▙▖▌
       ▄▌     ▄▌
     """)
     if transport == "stdio":
         # Run MCP server over STDIO (local)
-        asyncio.run(run_stdio())
+        run_stdio()
     else:
         # Run MCP server over streamable HTTP by default
         run_http()
 
 
 if __name__ == "__main__":
-    main()
+    try:
+        sys.exit(main())
+    except KeyboardInterrupt:
+        os._exit(0)

tests/conftest.py

@@ -6,6 +6,14 @@
 from fastmcp import FastMCP
 
 
+class MockMCP(FastMCP):
+    """
+    Mock class for FastMCP
+    """
+
+    pass
+
+
 def util_load_json(path):
     """
     Utility function to load a JSON file from the given path.
@@ -42,8 +50,8 @@ def mock_ctx():
     Returns:
         Context: A mocked Context object with 'fastmcp' tags.
     """
-    fastmcp: FastMCP = FastMCP(
-        name="Test",
+
+    fastmcp: MockMCP = MockMCP(
         tags=["sysdig", "mcp", "stdio"],
     )
     ctx = Context(fastmcp=fastmcp)

tools/events_feed/tool.py

@@ -14,6 +14,7 @@
 from fastmcp import Context
 from sysdig_client import ApiException
 from fastmcp.prompts.prompt import PromptMessage, TextContent
+from fastmcp.exceptions import ToolError
 from starlette.requests import Request
 from sysdig_client.api import SecureEventsApi
 from utils.sysdig.old_sysdig_api import OldSysdigApi
@@ -99,7 +100,7 @@ def tool_get_event_info(self, event_id: str, ctx: Context) -> dict:
             response = create_standard_response(results=raw, execution_time_ms=execution_time)
 
             return response
-        except ApiException as e:
+        except ToolError as e:
             logging.error("Exception when calling SecureEventsApi->get_event_v1: %s\n" % e)
             raise e
 
@@ -181,7 +182,7 @@ def tool_list_runtime_events(
                 execution_time_ms=duration_ms,
             )
             return response
-        except ApiException as e:
+        except ToolError as e:
             log.error(f"Exception when calling SecureEventsApi->get_events_v1: {e}\n")
             raise e
 
@@ -225,7 +226,7 @@ def tool_get_event_process_tree(self, ctx: Context, event_id: str) -> Dict[str,
             )
 
             return response
-        except ApiException as e:
+        except ToolError as e:
             log.error(f"Exception when calling Sysdig Sage API to get process tree: {e}")
             raise e
 

tools/inventory/tool.py

@@ -9,6 +9,7 @@
 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
 from sysdig_client.api import InventoryApi
@@ -69,133 +70,67 @@ def tool_list_resources(
             Field(
                 description=(
                     """
-                    Use the filter-query-language to filter the results.
-                    
+                    Sysdig Secure query filter expression to filter inventory resources.
+
+                    Use the resource://filter-query-language to get the expected filter expression format.
+               
                     List of supported fields:
-                    - field: accountName
-                      Description:  The account name that will be included in the results.
-                    - field: accountId
-                      Description:  The account id that will be included in the results.
-                    - field: cluster
-                      Description:  The kubernetes cluster that will be included in the results.
-                    - field: externalDNS
-                      Description:  The external DNS that will be included in the results.
-                    - field: distribution
-                      Description:  The kubernetes distribution that will be included in the results.
-                    - field: integrationName
-                      Description:  The name of the integration an IaC resource belongs to.
-                    - field: labels
-                      Description:  The resource labels that will be included in the results.
-                    - field: location
-                      Description:  The web address of an IaC Manifest.
-                    - field: name
-                      Description:  The names that will be included in the results.
-                    - field: namespace
-                      Description:  The namespace that will be included in the results.
-                    - field: nodeType
-                      Description:  The nodeType that will be included in the results.
-                    - field: osName
-                      Description:  The operating system that will be included in the results.
-                    - field: osImage
-                      Description:  The operating system image that will be included in the results.
-                    - field: organization
-                      Description:  The organization that will be included in the results.
-                    - field: platform
-                      Description:  The platform that will be included in the results.
-                    - field: control.accepted
-                      Description:  Include (or Exclude) only resources with accepted results.
-                                    Supported operators: exists and not exists.
-                    - field: policy
-                      Description:  Include resources that applied the selected policies.
-                                    Supported operators: in, not in, exists, not exists.
-                    - field: control.severity
-                      Description:  Include resources that have violated risks in the selected severities.
-                                    Supported operators: in, not in.
-                    - field: control.failed
-                      Description:  Include resources that have violated the selected risks.
-                                    Supported operators: in, not in, exists, not exists.
-                    - field: policy.failed
-                      Description:  Include resources that failed the selected policies.
-                                    Supported operators: in, not in, exists, not exists.
-                    - field: policy.passed
-                      Description:  Include resources that passed the selected policies.
-                                    Supported operators: in, not in, exists, not exists.
-                    - field: projectName
-                      Description:  The project name that will be included in the results.
-                    - field: projectId
-                      Description:  The project id that will be included in the results.
-                    - field: region
-                      Description:  The regions that will be included in the results.
-                    - field: repository
-                      Description:  The Repository an IaC resource belongs to.
-                    - field: resourceOrigin
-                      Description:  Origin of the resource. Supported values: Code, Deployed.
-                    - field: type
-                      Description:  The resource types that will be included in the results.
-                    - field: subscriptionName
-                      Description:  The Azure subscription name that will be included in the results.
-                    - field: subscriptionId
-                      Description:  The Azure subscription id that will be included in the results.
-                    - field: sourceType
-                      Description:  The source type of an IaC resource.
-                                    Supported values: YAML, Kustomize, Terraform, Helm.
-                    - field: version
-                      Description:  OCP Cluster versions that will be included in the results.
-                    - field: zone
-                      Description:  The zones that will be included in the results.
-                    - field: category
-                      Description:  The category that will be included in the results.
-                                    Supported operators: in, not in.
-                    - field: isExposed
-                      Description:  Specifies whether the resource to return is exposed to the internet.
-                                    Supported operators: exists and not exists.
-                    - field: validatedExposure
-                      Description:  Specifies whether the resource to return is exposed to the internet and could be reach
-                                    by our network exposure validator. Supported operators: exists and not exists.
-                    - field: arn
-                      Description:  The AWS ARN of the resource.
-                    - field: resourceId
-                      Description:  The Azure or GCP Resource Identifier of the resource.
-                    - field: container.name
-                      Description:  Filters the resource by a container.
-                    - field: architecture
-                      Description:  Image architecture.
-                    - field: baseOS
-                      Description:  Image Base OS.
-                    - field: digest
-                      Description:  Image Digest.
-                    - field: imageId
-                      Description:  Image Id.
-                    - field: os
-                      Description:  Image OS.
-                    - field: container.imageName
-                      Description:  Image Pullstring.
-                    - field: image.registry
-                      Description:  Image Registry.
-                    - field: image.tag
-                      Description:  Image tag.
-                    - field: package.inUse
-                      Description:  Package in use filter. Supported operators: exists and not exists.
-                    - field: package.info
-                      Description:  Filters by a package using the format [packge name] - field: [version].
-                    - field: package.path
-                      Description:  Filters by package path.
-                    - field: package.type
-                      Description:  Package type.
-                    - field: vuln.cvssScore
-                      Description:  Filter by vulnerability CVSS. Supported operators: = and >=.
-                    - field: vuln.hasExploit
-                      Description:  Filters resources by the existence of vulnerabilities with exploits.
-                      Supported operators: exists and not exists.
-                    - field: vuln.hasFix
-                      Description:  Filters resources by the existence of vulnerabilities with fixes.
-                                    Supported operators: exists and not exists.
-                    - field: vuln.name
-                      Description:  Filter by vulnerability name.
-                    - field: vuln.severity
-                      Description:  Filter by vulnerability severity. Supported operators: in, not in, exists and not exists.
-                    - field: machineImage
-                      Description:  Filter by host machine image.
+                    - accountName
+                    - accountId
+                    - cluster
+                    - externalDNS
+                    - distribution
+                    - integrationName
+                    - labels
+                    - location
+                    - name
+                    - namespace
+                    - nodeType
+                    - osName
+                    - osImage
+                    - organization
+                    - platform
+                    - control.accepted
+                    - policy
+                    - control.severity
+                    - control.failed
+                    - policy.failed
+                    - policy.passed
+                    - projectName
+                    - projectId
+                    - region
+                    - repository
+                    - resourceOrigin
+                    - type
+                    - subscriptionName
+                    - subscriptionId
+                    - sourceType
+                    - version
+                    - zone
+                    - category
+                    - isExposed
+                    - validatedExposure
+                    - arn
+                    - resourceId
+                    - container.name
+                    - architecture
+                    - baseOS
+                    - digest
+                    - imageId
+                    - os
+                    - container.imageName
+                    - image.registry
+                    - image.tag
+                    - package.inUse
+                    - package.info
+                    - package.path
+                    - package.type
+                    - vuln.cvssScore
+                    - vuln.hasExploit
+                    - vuln.hasFix
+                    - vuln.name
+                    - vuln.severity
+                    - machineImage
                 """
                 ),
                 examples=[
@@ -205,6 +140,7 @@ def tool_list_resources(
                     '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',
                 ],
             ),
         ] = 'platform in ("GCP", "AWS", "Azure", "Kubernetes")',
@@ -243,7 +179,7 @@ def tool_list_resources(
             response = create_standard_response(results=api_response, execution_time_ms=execution_time)
 
             return response
-        except ApiException as e:
+        except ToolError as e:
             logging.error("Exception when calling InventoryApi->get_resources: %s\n" % e)
             raise e
 
@@ -272,6 +208,6 @@ def tool_get_resource(
             response = create_standard_response(results=api_response, execution_time_ms=execution_time)
 
             return response
-        except ApiException as e:
+        except ToolError as e:
             log.error(f"Exception when calling InventoryApi->get_resource: {e}")
             raise e

tools/sysdig_sage/tool.py

@@ -9,6 +9,7 @@
 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
 from fastmcp.server.dependencies import get_http_request
@@ -74,7 +75,7 @@ async def tool_sysdig_sage(self, ctx: Context, question: str) -> Dict[str, Any]:
             Dict: JSON-decoded response of the executed SysQL query, or an error object.
 
         Raises:
-            Exception: If the SysQL query generation or execution fails.
+            ToolError: If the SysQL query generation or execution fails.
 
         Examples:
             # tool_sysdig_sage(question="Match Cloud Resource affected by Critical Vulnerability")
@@ -87,8 +88,8 @@ async def tool_sysdig_sage(self, ctx: Context, question: str) -> Dict[str, Any]:
             old_sysdig_api = self.init_client(config_tags=ctx.fastmcp.tags)
             sysql_response = await old_sysdig_api.generate_sysql_query(question)
             if sysql_response.status > 299:
-                raise Exception(f"Sysdig Sage returned an error: {sysql_response.status} - {sysql_response.data}")
-        except Exception as e:
+                raise ToolError(f"Sysdig Sage returned an error: {sysql_response.status} - {sysql_response.data}")
+        except ToolError as e:
             log.error(f"Failed to generate SysQL query: {e}")
             raise e
         json_resp = sysql_response.json() if sysql_response.data else {}
@@ -107,6 +108,6 @@ async def tool_sysdig_sage(self, ctx: Context, question: str) -> Dict[str, Any]:
             )
 
             return response
-        except Exception as e:
+        except ToolError as e:
             log.error(f"Failed to execute SysQL query: {e}")
             raise e