#!/usr/bin/env python3
"""Connect to an MCP server, print everything it exposes, and call one tool.

Purpose: the review step from the Model Context Protocol lesson, as a program. It
    connects to a server, prints every tool name, full description and input schema
    exactly as the model would receive them, screens those descriptions for text that
    reads as an instruction rather than a description, and then calls one tool so you
    can see a real result. Run this before you connect any server you did not write,
    and again after it updates, because a description can change after you approved it.
Platform: all (pure Python over stdio or HTTP; use WSL2 on Windows for a stdio server)
Minimum memory: 8 GB on the machine serving the model; this script needs almost none
Assumes: Python 3.9 or later, no third-party packages, and mcpbridge.py beside this
    file. A server to inspect: this part's mcp-server.py, or another one.

Usage: python3 mcp-client-check.py --server "python3 mcp-server.py"
       python3 mcp-client-check.py --url http://127.0.0.1:8000/mcp
       python3 mcp-client-check.py --server "python3 mcp-server.py" \\
           --call search_documents --args '{"query": "backup window"}' --labbook labbook.md
"""

from __future__ import annotations

import argparse
import json
import re
import time
from pathlib import Path
from typing import Any, Dict, List

from mcpbridge import McpBridge, McpError

# Phrases that turn a description into an instruction. None of these is proof of an
# attack and none of their absence is proof of safety: the screen exists to put your
# eyes on the right paragraph, because the model reads all of it and a host UI does not.
INSTRUCTION_MARKERS = [
    (re.compile(r"ignore (?:all |any )?(?:previous|prior|earlier|above)", re.I), "overrides earlier instructions"),
    (re.compile(r"\bdo not (?:tell|mention|show|reveal|inform)\b", re.I), "asks to hide something from the user"),
    (re.compile(r"\b(?:before|prior to) (?:answering|responding|using this|calling)\b", re.I), "adds a step before the answer"),
    (re.compile(r"<\s*(?:important|system|secret|instructions?)\b", re.I), "uses a pseudo-tag to raise its authority"),
    (re.compile(r"\byou (?:must|should always|are required to)\b", re.I), "gives the model an obligation"),
    (re.compile(r"\bread (?:the )?(?:file|contents of|~/|/etc|\.env|id_rsa|ssh)", re.I), "names files to read"),
    (re.compile(r"\b(?:send|post|upload|exfiltrate|forward) (?:it|them|the (?:contents|result|output))\b", re.I), "asks for data to be sent somewhere"),
    (re.compile(r"\bapi[ _-]?key|\btoken\b|\bpassword\b|\bcredential", re.I), "mentions credentials"),
    (re.compile(r"\bsystem prompt\b", re.I), "refers to the system prompt"),
]

# A description far longer than the job needs is the shape tool poisoning takes, because
# the hidden part has to fit somewhere. This is a prompt to read, not a verdict.
LONG_DESCRIPTION_CHARS = 600


def screen(text: str) -> List[str]:
    """Every instruction-like marker found in one description."""
    findings = [why for pattern, why in INSTRUCTION_MARKERS if pattern.search(text or "")]
    if len(text or "") > LONG_DESCRIPTION_CHARS:
        findings.append("unusually long for a tool description (%d characters)" % len(text))
    return findings


def describe(bridge: McpBridge) -> Dict[str, Any]:
    """Print the server's whole surface and return what the screen found."""
    print("transport: %s" % ("Streamable HTTP" if bridge.url else "stdio"))
    print("protocol:  %s" % ("legacy initialize handshake" if bridge.legacy else "2026-07-28"))
    print("server:    %s" % json.dumps(bridge.server_info)[:400])
    print("prefix:    %s_  (added by the client so two servers cannot collide)" % bridge.label)

    flagged: Dict[str, List[str]] = {}
    for tool in bridge.tools:
        description = tool.get("description", "") or ""
        findings = screen(description)
        if findings:
            flagged[tool["name"]] = findings
        print("\n--- tool: %s" % tool["name"])
        if tool.get("title"):
            print("title: %s" % tool["title"])
        print("description, in full, exactly as the model receives it:")
        print("  " + (description.replace("\n", "\n  ") if description else "(none)"))
        print("input schema: %s" % json.dumps(tool.get("inputSchema", {}), indent=2))
        if tool.get("annotations"):
            print("annotations: %s   <- hints from the server; the specification says to treat "
                  "these as untrusted unless the server is trusted"
                  % json.dumps(tool["annotations"]))
        for finding in findings:
            print("  SCREEN: %s" % finding)

    for name in ("resources/list", "prompts/list"):
        try:
            result = bridge.request(name)
            items = result.get(name.split("/")[0], [])
            print("\n%s: %d item(s): %s" % (name, len(items),
                  ", ".join(str(i.get("uri") or i.get("name")) for i in items) or "-"))
        except McpError as exc:
            print("\n%s: not offered (%s)" % (name, exc))
    return flagged


def pick_call(bridge: McpBridge, wanted: str) -> str:
    names = [t["name"] for t in bridge.tools]
    if wanted:
        if wanted not in names:
            raise SystemExit("no tool named %r; this server offers: %s" % (wanted, ", ".join(names)))
        return wanted
    if not names:
        raise SystemExit("the server offers no tools to call")
    return names[0]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--server", default=None, help="command that starts a stdio server")
    parser.add_argument("--url", default=None, help="a Streamable HTTP MCP endpoint")
    parser.add_argument("--call", default=None, help="tool to call; defaults to the first one")
    parser.add_argument("--args", default='{"query": "backup window"}',
                        help="arguments for that call, as a JSON object")
    parser.add_argument("--no-call", action="store_true", help="inspect only, call nothing")
    parser.add_argument("--labbook", default=None)
    args = parser.parse_args()

    bridge = McpBridge(server=args.server, url=args.url)
    try:
        flagged = describe(bridge)
        called, result = None, None
        if not args.no_call:
            called = pick_call(bridge, args.call)
            print("\n=== calling %s with %s" % (called, args.args))
            result = bridge.call(bridge.label + "_" + called, json.loads(args.args))
            print(result[:1500])

        print("\n%d of %d description(s) contained instruction-like text."
              % (len(flagged), len(bridge.tools)))
        if flagged:
            for name, findings in flagged.items():
                print("  %s: %s" % (name, "; ".join(findings)))
            print("Read those descriptions again in full before connecting this server to "
                  "anything with your files or your network.")
    finally:
        bridge.close()

    if args.labbook:
        record = {"lab": "part-24/mcp-client-check",
                  "server": args.server or args.url,
                  "transport": "http" if args.url else "stdio",
                  "protocol": "legacy" if bridge.legacy else "2026-07-28",
                  "tools": [t["name"] for t in bridge.tools],
                  "flagged": flagged,
                  "called": called,
                  "result_chars": len(result or ""),
                  "recorded_at": time.strftime("%Y-%m-%dT%H:%M:%S")}
        with Path(args.labbook).open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(record) + "\n")
        print("recorded in %s" % args.labbook)


if __name__ == "__main__":
    try:
        main()
    except McpError as exc:
        raise SystemExit("mcp: %s" % exc)