Skip to content

Bug fix: notebook/sdk_byot_custom_mcp adding list_tools() #116

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cemnyc wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Bug fix: notebook/sdk_byot_custom_mcp adding list_tools() #116

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 52 additions & 3 deletions notebook/sdk_byot_custom_mcp.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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",
Expand Down Expand Up @@ -167,4 +216,4 @@
},
"nbformat": 4,
"nbformat_minor": 4
}
}