Finish up PoC

Ickerday · Ickerday · commit fcae942949be · 2025-10-02T14:30:56.000+02:00
diff --git a/splunklib/mcp/tools.py b/splunklib/mcp/tools.py
@@ -0,0 +1,110 @@
+import configparser
+import os
+from dataclasses import asdict, dataclass, field
+from typing import Any, Literal
+
+from fastmcp.client import Client
+from fastmcp.client.transports import PythonStdioTransport
+from mcp.types import Tool
+
+
+@dataclass
+class SplunkMeta:
+    permissions: list[str] = field()
+    tool_type: str = field(default="")
+    schema_version: str = field(default="")
+    execution_mode: str = field(default="")
+    execution_endpoint: str = field(default="")
+
+
+@dataclass
+class McpInputOutputSchema:
+    type: Literal["object"] = "object"
+    properties: dict[str, Any] = field(default_factory=lambda: {})  # pyright: ignore[reportExplicitAny]
+    required: list[str] = field(default_factory=lambda: [])
+
+
+tool_reg_prefix = "app:mcp_tool"
+
+
+def filter_sections(section_name: str) -> bool:
+    return section_name.startswith(tool_reg_prefix)
+
+
+def match_input_schema(input: Literal["query_string"] | Literal["other"]):
+    """Gets super messy :("""
+    match input:
+        case "query_string":
+            return {
+                "type": "object",
+                "properties": {
+                    "query_string": {
+                        "type": "string",
+                        "description": "SPL2 query string",
+                    }
+                },
+            }
+        case _:
+            raise NotImplementedError("We don't know what to put here lol")
+
+
+def parse_ai_conf(file_path: str) -> list[Tool]:
+    config = configparser.ConfigParser()
+    all_sections_len = config.read(file_path)
+    if len(all_sections_len) == 0:
+        return []
+
+    tool_reg_sections: list[str] = list(filter(filter_sections, config.sections()))
+    if len(tool_reg_sections) == 0:
+        return []
+
+    ini_tools: list[Tool] = []
+    for reg_section in tool_reg_sections:
+        reg_section_data = config[reg_section]
+
+        name: str = reg_section.split(":")[2]
+        description = reg_section_data["description"]
+        # https://modelcontextprotocol.io/specification/2025-06-18/schema#tool
+        inputSchema = McpInputOutputSchema(properties={}, required=[])
+        outputSchema = McpInputOutputSchema(properties={}, required=[])
+        meta = SplunkMeta(
+            permissions=[
+                perm.strip()
+                for perm in reg_section_data["permissions"].strip().split(",")
+            ],
+            tool_type="search",
+            schema_version=reg_section_data["schema_version"].strip(),
+        )
+
+        ini_tool = Tool(
+            name=name,
+            description=description,
+            inputSchema=asdict(inputSchema),
+            outputSchema=asdict(outputSchema),
+            _meta=asdict(meta),
+        )
+        ini_tools.append(ini_tool)
+
+    return ini_tools
+
+
+async def get_tools(server_path: str):
+    mcp_client = Client(server_path)
+
+    tools = []
+    async with mcp_client:
+        tools = await mcp_client.list_tools()
+
+        # TODO: Get registrations from ai.conf
+        # curr_path = os.path.join(os.getcwd(), "..", "default", "app.conf")
+        # yaml_tool_registrations: list[Tool] = parse_ai_conf(curr_path)
+
+    return tools
+
+
+async def register_tools_from(file_paths: list[str]) -> None:
+    """TODO
+    1. `POST /tools` with MCP payload
+    2.
+    """
+    print(file_paths)
diff --git a/tests/system/test_apps/mcp_enabled_app/bin/mcp_enabled_app.ipynb b/tests/system/test_apps/mcp_enabled_app/bin/mcp_enabled_app.ipynb
@@ -10,8 +10,23 @@
     "This example aspires to verify the points listed in [POC - AI with Splunk Apps](https://cisco-my.sharepoint.com/:w:/r/personal/hbalacha_cisco_com/Documents/POC%20-%20AI%20with%20Splunk%20Apps.docx?d=w2776e089011943abbd84c0fa30a53f34&csf=1&web=1&e=RxvShR)\n",
     "\n",
     "- Develop @tool Decorator\n",
-    "  - [ ] Capture e.g. tool_name, description, inputs, outputs\n",
-    "- MCP JSONSchema can (and most probably should) be used for tool registration in Splunk\n"
+    "  - [ ] Capture e.g. tool_name, description, inputs, outputs\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "18ab5550",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "363201e8",
+   "metadata": {},
+   "source": [
+    "- MCP JSONSchema can (and most probably should) be used for tool registration in Splunk"
    ]
   },
   {
@@ -98,117 +113,6 @@
     "  - [ ] Call MCP registry with credentials (from App Manager)\n",
     "  - [ ] Log success/fail to MCP audit\n"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "2e594272",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "/Users/bjedreck/Projects/spl-mcp-tool/bin/../default/app.conf\n",
-      "name='aws_logs_search' title=None description='Execute SPL queries against AWS logs' inputSchema={'type': 'object', 'properties': {}, 'required': []} outputSchema={'type': 'object', 'properties': {}, 'required': []} icons=None annotations=None meta={'permissions': ['role:search_admin', 'role:aws_analyst'], 'tool_type': 'search', 'schema_version': '1.0'}\n"
-     ]
-    }
-   ],
-   "source": [
-    "import configparser\n",
-    "import os\n",
-    "from dataclasses import asdict, dataclass, field\n",
-    "from typing import Any, Literal\n",
-    "\n",
-    "from mcp.types import Tool\n",
-    "\n",
-    "\n",
-    "@dataclass\n",
-    "class SplunkMeta:\n",
-    "    permissions: list[str]\n",
-    "    tool_type: str\n",
-    "    schema_version: str\n",
-    "\n",
-    "\n",
-    "@dataclass\n",
-    "class McpInputOutputSchema:\n",
-    "    type: Literal[\"object\"] = \"object\"\n",
-    "    properties: dict[str, Any] = field(default_factory=lambda: {})  # pyright: ignore[reportExplicitAny]\n",
-    "    required: list[str] = field(default_factory=lambda: [])\n",
-    "\n",
-    "\n",
-    "tool_reg_prefix = \"app:mcp_tool\"\n",
-    "\n",
-    "\n",
-    "def filter_sections(section_name: str):\n",
-    "    return section_name.startswith(tool_reg_prefix)\n",
-    "\n",
-    "\n",
-    "def match_input_schema(input: Literal[\"query_string\"] | Literal[\"other\"]):\n",
-    "    match input:\n",
-    "        case \"query_string\":\n",
-    "            return {\n",
-    "                \"type\": \"object\",\n",
-    "                \"properties\": {\n",
-    "                    \"query_string\": {\n",
-    "                        \"type\": \"string\",\n",
-    "                        \"description\": \"SPL2 query string\",\n",
-    "                    }\n",
-    "                },\n",
-    "            }\n",
-    "        case _:\n",
-    "            raise NotImplementedError(\"We don't know what to put here lol\")\n",
-    "\n",
-    "\n",
-    "def parse_app_conf_tool_registrations(file_path: str) -> list[Tool]:\n",
-    "    config = configparser.ConfigParser()\n",
-    "    all_sections_len = config.read(file_path)\n",
-    "    if len(all_sections_len) == 0:\n",
-    "        return []\n",
-    "\n",
-    "    tool_reg_sections: list[str] = list(filter(filter_sections, config.sections()))\n",
-    "    if len(tool_reg_sections) == 0:\n",
-    "        return []\n",
-    "\n",
-    "    ini_tools: list[Tool] = []\n",
-    "    for reg_section in tool_reg_sections:\n",
-    "        reg_section_data = config[reg_section]\n",
-    "\n",
-    "        name: str = reg_section.split(\":\")[2]\n",
-    "        description = reg_section_data[\"description\"]\n",
-    "        # https://modelcontextprotocol.io/specification/2025-06-18/schema#tool\n",
-    "        inputSchema = McpInputOutputSchema(properties={}, required=[])\n",
-    "        outputSchema = McpInputOutputSchema(properties={}, required=[])\n",
-    "        meta = SplunkMeta(\n",
-    "            permissions=[\n",
-    "                perm.strip()\n",
-    "                for perm in reg_section_data[\"permissions\"].strip().split(\",\")\n",
-    "            ],\n",
-    "            tool_type=\"search\",\n",
-    "            schema_version=reg_section_data[\"schema_version\"].strip(),\n",
-    "        )\n",
-    "\n",
-    "        ini_tool = Tool(\n",
-    "            name=name,\n",
-    "            description=description,\n",
-    "            inputSchema=asdict(inputSchema),\n",
-    "            outputSchema=asdict(outputSchema),\n",
-    "            _meta=asdict(meta),\n",
-    "        )\n",
-    "        ini_tools.append(ini_tool)\n",
-    "\n",
-    "    return ini_tools\n",
-    "\n",
-    "\n",
-    "async def post_install():\n",
-    "    curr_path = os.path.join(os.getcwd(), \"..\", \"default\", \"app.conf\")\n",
-    "    print(curr_path)\n",
-    "    yaml_tool_registrations: list[Tool] = parse_app_conf_tool_registrations(curr_path)\n",
-    "    [print(toolReg) for toolReg in yaml_tool_registrations]\n",
-    "\n",
-    "\n",
-    "await post_install()\n"
-   ]
   }
  ],
  "metadata": {
diff --git a/tests/system/test_apps/mcp_enabled_app/bin/post_install_example.py b/tests/system/test_apps/mcp_enabled_app/bin/post_install_example.py
@@ -0,0 +1,10 @@
+from splunklib.mcp import tools
+
+
+async def post_install(splunk_url: str, auth_token: str) -> None:
+    # TODO: Implement
+    try:
+        await tools.register_tools_from(["./tools.py", "../local/ai.conf"])
+    except Exception as e:
+        print(e)
+        raise e
diff --git a/tests/system/test_apps/mcp_enabled_app/bin/tools.py b/tests/system/test_apps/mcp_enabled_app/bin/tools.py
@@ -0,0 +1,58 @@
+import asyncio
+import json
+from dataclasses import asdict
+
+from fastmcp import FastMCP
+
+from splunklib import client
+from splunklib.mcp.tools import SplunkMeta
+from splunklib.results import JSONResultsReader
+
+MCP_SERVER_HOST: str = "0.0.0.0"
+MCP_SERVER_PORT: int = 2137
+
+MCP_SERVER_NAME: str = "GeneratingCSC PoC Server"
+app_mcp_server = FastMCP(MCP_SERVER_NAME)
+
+
+@app_mcp_server.tool(
+    description="""
+    The `generatingcsc` command generates a specific number of records.
+    
+    Example:
+    ``| generatingcsc count=4``
+    Returns a 4 records having text 'Test Event'.
+    """,
+    meta=asdict(
+        SplunkMeta(
+            permissions=["role:search_admin", "role:aws_analyst"],
+            tool_type="search",
+            schema_version="1.0",
+            execution_endpoint="",
+            execution_mode="",
+        )
+    ),
+    enabled=True,
+)
+def generating_csc(count: int = 10) -> list[str]:
+    service = client.connect(
+        scheme="https",
+        host="localhost",
+        port="8089",
+        username="admin",
+        password="changed!",
+        autologin=True,
+    )
+    stream = service.jobs.oneshot(f"| generatingcsc count={count}", output_mode="json")
+    results = JSONResultsReader(stream)
+
+    quuuuuux = [json.dumps(r) for r in list(results)]
+    print(quuuuuux)
+    return quuuuuux
+
+
+if __name__ == "__main__":
+    asyncio.run(
+        # app_mcp_server.run_streamable_http_async(MCP_SERVER_HOST, port=MCP_SERVER_PORT)
+        app_mcp_server.run_stdio_async(False)
+    )
