From b1ea5c66ac6d5772e5fec0ecdcaf8ac0860691e2 Mon Sep 17 00:00:00 2001 From: cemnyc Date: Tue, 2 Jun 2026 16:25:26 -0400 Subject: [PATCH] Bug fix: Adding list_tools() --- notebook/sdk_byot_custom_mcp.ipynb | 55 ++++++++++++++++++++++++++++-- 1 file changed, 52 insertions(+), 3 deletions(-) diff --git a/notebook/sdk_byot_custom_mcp.ipynb b/notebook/sdk_byot_custom_mcp.ipynb index 7d1cbfe..3f1ba30 100644 --- a/notebook/sdk_byot_custom_mcp.ipynb +++ b/notebook/sdk_byot_custom_mcp.ipynb @@ -36,7 +36,24 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": "import os\n\nAPI_KEY = os.environ.get(\"SEC_GEMINI_API_KEY\")\n\nif not API_KEY:\n try:\n from google.colab import userdata # type: ignore[import-not-found]\n\n API_KEY = userdata.get(\"SEC_GEMINI_API_KEY\")\n except (ImportError, Exception):\n pass\n\nif not API_KEY:\n API_KEY = \"YOUR_API_KEY_HERE\"\n\nassert API_KEY and API_KEY != \"YOUR_API_KEY_HERE\", \"Please set your API key\"" + "source": [ + "import os\n", + "\n", + "API_KEY = os.environ.get(\"SEC_GEMINI_API_KEY\")\n", + "\n", + "if not API_KEY:\n", + " try:\n", + " from google.colab import userdata # type: ignore[import-not-found]\n", + "\n", + " API_KEY = userdata.get(\"SEC_GEMINI_API_KEY\")\n", + " except (ImportError, Exception):\n", + " pass\n", + "\n", + "if not API_KEY:\n", + " API_KEY = \"YOUR_API_KEY_HERE\"\n", + "\n", + "assert API_KEY and API_KEY != \"YOUR_API_KEY_HERE\", \"Please set your API key\"" + ] }, { "cell_type": "markdown", @@ -53,7 +70,39 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": "from fastmcp import FastMCP\n\ncustom_mcp = FastMCP(\"custom-security-tools\")\n\n\n@custom_mcp.tool()\ndef lookup_hash(sha256: str) -> str:\n \"\"\"Look up a file hash in the local threat intelligence database.\"\"\"\n # Simulated lookup -- replace with real logic\n known_hashes = {\n \"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855\": \"Known clean: empty file\",\n \"a\" * 64: \"MALICIOUS: Emotet dropper variant (TLP:RED)\",\n }\n result = known_hashes.get(\n sha256.lower(), f\"Hash {sha256[:16]}... not found in local DB\"\n )\n return result\n\n\n@custom_mcp.tool()\ndef check_ip_reputation(ip_address: str) -> str:\n \"\"\"Check an IP address against the local reputation list.\"\"\"\n # Simulated check -- replace with real logic\n blocklist = {\"192.168.1.100\": \"Internal scanner\", \"10.0.0.1\": \"Gateway\"}\n if ip_address in blocklist:\n return f\"{ip_address}: {blocklist[ip_address]}\"\n return f\"{ip_address}: No reputation data in local DB\"\n\n\nprint(\n f\"Custom MCP defined with tools: {[t.name for t in custom_mcp.local_provider.tools.values()]}\" # type: ignore[attr-defined]\n)" + "source": [ + "from fastmcp import FastMCP\n", + "\n", + "custom_mcp = FastMCP(\"custom-security-tools\")\n", + "\n", + "\n", + "@custom_mcp.tool()\n", + "def lookup_hash(sha256: str) -> str:\n", + " \"\"\"Look up a file hash in the local threat intelligence database.\"\"\"\n", + " # Simulated lookup -- replace with real logic\n", + " known_hashes = {\n", + " \"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855\": \"Known clean: empty file\",\n", + " \"a\" * 64: \"MALICIOUS: Emotet dropper variant (TLP:RED)\",\n", + " }\n", + " result = known_hashes.get(\n", + " sha256.lower(), f\"Hash {sha256[:16]}... not found in local DB\"\n", + " )\n", + " return result\n", + "\n", + "\n", + "@custom_mcp.tool()\n", + "def check_ip_reputation(ip_address: str) -> str:\n", + " \"\"\"Check an IP address against the local reputation list.\"\"\"\n", + " # Simulated check -- replace with real logic\n", + " blocklist = {\"192.168.1.100\": \"Internal scanner\", \"10.0.0.1\": \"Gateway\"}\n", + " if ip_address in blocklist:\n", + " return f\"{ip_address}: {blocklist[ip_address]}\"\n", + " return f\"{ip_address}: No reputation data in local DB\"\n", + "\n", + "\n", + "tools = await custom_mcp.list_tools()\n", + "print(f\"Custom MCP defined with tools: {[t.name for t in tools]}\")" + ] }, { "cell_type": "markdown", @@ -167,4 +216,4 @@ }, "nbformat": 4, "nbformat_minor": 4 -} \ No newline at end of file +}