fix: sniffer

Author: tech4242

examples/generate_traffic/README.md

@@ -0,0 +1,63 @@
+# Traffic Generation Examples
+
+This directory contains example servers and clients for generating MCP traffic to test MCPHawk.
+
+## TCP-based MCP
+
+### Server
+```bash
+python3 tcp_server.py
+```
+Starts a TCP MCP server on port 12345.
+
+### Client
+```bash
+python3 tcp_client.py
+```
+Sends various MCP messages to the TCP server.
+
+## WebSocket-based MCP
+
+### Server
+```bash
+python3 ws_server.py
+```
+Starts a WebSocket MCP server on port 8765.
+
+### Client
+```bash
+python3 ws_client.py
+```
+Sends various MCP messages to the WebSocket server.
+
+## Generate All Traffic
+
+To generate both TCP and WebSocket traffic for testing:
+
+```bash
+python3 generate_all.py
+```
+
+This will:
+1. Start both TCP and WebSocket servers
+2. Send a variety of MCP messages to each
+3. Display the traffic being generated
+4. Clean up when done
+
+## Testing with MCPHawk
+
+In another terminal, run MCPHawk to capture the traffic:
+
+```bash
+# Capture TCP traffic
+sudo mcphawk sniff --port 12345
+
+# Capture WebSocket traffic
+sudo mcphawk sniff --port 8765
+
+# Capture both
+sudo mcphawk sniff --filter "tcp port 12345 or tcp port 8765"
+
+# Or use auto-detect
+sudo mcphawk sniff --auto-detect
+```

examples/generate_traffic/generate_all.py

@@ -0,0 +1,161 @@
+#!/usr/bin/env python3
+"""Generate both TCP and WebSocket MCP traffic for testing MCPHawk."""
+
+import asyncio
+import contextlib
+import json
+import logging
+import socket
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import websockets
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s [%(levelname)s] %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+def run_tcp_server():
+    """Run the TCP server in a subprocess."""
+    script_path = Path(__file__).parent / "tcp_server.py"
+    return subprocess.Popen([sys.executable, str(script_path)])
+
+
+def run_ws_server():
+    """Run the WebSocket server in a subprocess."""
+    script_path = Path(__file__).parent / "ws_server.py"
+    return subprocess.Popen([sys.executable, str(script_path)])
+
+
+def send_tcp_traffic():
+    """Send TCP MCP traffic."""
+    logger.info("Sending TCP traffic...")
+
+    try:
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.connect(('localhost', 12345))
+
+        messages = [
+            {"jsonrpc": "2.0", "method": "initialize", "params": {"protocolVersion": "2024-11-05"}, "id": 1},
+            {"jsonrpc": "2.0", "method": "tools/list", "id": 2},
+            {"jsonrpc": "2.0", "method": "notifications/tcp_test", "params": {"source": "tcp"}},
+            {"jsonrpc": "2.0", "method": "tools/call", "params": {"name": "test_tool", "arguments": {}}, "id": 3},
+        ]
+
+        for msg in messages:
+            data = json.dumps(msg).encode('utf-8')
+            sock.sendall(data)
+            logger.info(f"  TCP sent: {msg.get('method')}")
+
+            # Read response if it has an ID
+            if "id" in msg:
+                try:
+                    response = sock.recv(1024)
+                    if response:
+                        logger.info("  TCP received response")
+                except Exception:
+                    pass
+
+            time.sleep(0.2)
+
+        # Keep connection open a bit longer
+        time.sleep(1)
+        sock.close()
+
+    except Exception as e:
+        logger.error(f"TCP client error: {e}")
+
+
+async def send_ws_traffic():
+    """Send WebSocket MCP traffic."""
+    logger.info("Sending WebSocket traffic...")
+
+    try:
+        async with websockets.connect("ws://localhost:8765", compression=None) as ws:
+            messages = [
+                {"jsonrpc": "2.0", "method": "initialize", "params": {"protocolVersion": "2024-11-05"}, "id": 1},
+                {"jsonrpc": "2.0", "method": "tools/list", "id": 2},
+                {"jsonrpc": "2.0", "method": "notifications/ws_test", "params": {"source": "websocket"}},
+                {"jsonrpc": "2.0", "method": "tools/call", "params": {"name": "calculator", "arguments": {"a": 10, "b": 20}}, "id": 3},
+            ]
+
+            for msg in messages:
+                await ws.send(json.dumps(msg))
+                logger.info(f"  WS sent: {msg.get('method')}")
+
+                # Wait for response if it has an ID
+                if "id" in msg:
+                    await ws.recv()
+                    logger.info("  WS received response")
+                else:
+                    await asyncio.sleep(0.2)
+
+    except Exception as e:
+        logger.error(f"WebSocket client error: {e}")
+
+
+def check_port(port):
+    """Check if a port is available."""
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    result = sock.connect_ex(('localhost', port))
+    sock.close()
+    return result != 0
+
+
+async def main():
+    """Generate all traffic types."""
+    logger.info("MCPHawk Traffic Generator")
+    logger.info("========================\n")
+
+    # Check if ports are available
+    if not check_port(12345):
+        logger.error("Port 12345 is already in use. Please stop the existing TCP server.")
+        return
+
+    if not check_port(8765):
+        logger.error("Port 8765 is already in use. Please stop the existing WebSocket server.")
+        return
+
+    logger.info("Starting servers...")
+
+    # Start servers
+    tcp_server = run_tcp_server()
+    ws_server = run_ws_server()
+
+    # Give servers time to start
+    logger.info("Waiting for servers to start...")
+    await asyncio.sleep(2)
+
+    try:
+        # Send TCP traffic
+        send_tcp_traffic()
+
+        # Send WebSocket traffic
+        await send_ws_traffic()
+
+        logger.info("\n✅ All traffic sent successfully!")
+        logger.info("\nServers will continue running. Press Ctrl+C to stop.")
+
+        # Keep running
+        await asyncio.Future()
+
+    except KeyboardInterrupt:
+        logger.info("\nStopping servers...")
+    finally:
+        # Clean up
+        tcp_server.terminate()
+        ws_server.terminate()
+        tcp_server.wait()
+        ws_server.wait()
+        logger.info("Servers stopped.")
+
+
+if __name__ == "__main__":
+    with contextlib.suppress(KeyboardInterrupt):
+        asyncio.run(main())
+

examples/generate_traffic/tcp_client.py

@@ -0,0 +1,100 @@
+#!/usr/bin/env python3
+"""TCP MCP client for testing."""
+
+import json
+import socket
+import time
+
+
+def send_mcp_message(sock, message):
+    """Send a JSON-RPC message."""
+    data = json.dumps(message).encode('utf-8')
+    sock.sendall(data)
+    print(f"Sent: {message.get('method', message.get('result', 'response'))}")
+
+    # Read response if message has an ID
+    if "id" in message:
+        try:
+            response = sock.recv(1024)
+            if response:
+                print("  Received response")
+        except Exception:
+            pass
+
+
+def main():
+    """Connect to TCP MCP server and send test messages."""
+    print("Connecting to TCP MCP server on localhost:12345...")
+
+    try:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.connect(('localhost', 12345))
+            print("Connected!")
+
+            # Send initialize
+            send_mcp_message(sock, {
+                "jsonrpc": "2.0",
+                "method": "initialize",
+                "params": {
+                    "protocolVersion": "2024-11-05",
+                    "capabilities": {}
+                },
+                "id": 1
+            })
+            time.sleep(0.5)
+
+            # Send tools/list
+            send_mcp_message(sock, {
+                "jsonrpc": "2.0",
+                "method": "tools/list",
+                "id": 2
+            })
+            time.sleep(0.5)
+
+            # Send a notification (no id)
+            send_mcp_message(sock, {
+                "jsonrpc": "2.0",
+                "method": "notifications/progress",
+                "params": {
+                    "progress": 50,
+                    "operation": "processing"
+                }
+            })
+            time.sleep(0.5)
+
+            # Send tools/call
+            send_mcp_message(sock, {
+                "jsonrpc": "2.0",
+                "method": "tools/call",
+                "params": {
+                    "name": "calculator",
+                    "arguments": {"a": 5, "b": 3}
+                },
+                "id": 3
+            })
+            time.sleep(0.5)
+
+            # Send a batch of messages quickly
+            print("\nSending batch of messages...")
+            for i in range(4, 8):
+                send_mcp_message(sock, {
+                    "jsonrpc": "2.0",
+                    "method": f"test/message_{i}",
+                    "params": {"value": i * 10},
+                    "id": i
+                })
+                time.sleep(0.1)
+
+            print("\nAll messages sent!")
+            # Keep connection open briefly to avoid reset
+            time.sleep(2)
+
+    except ConnectionRefusedError:
+        print("Error: Could not connect to server. Make sure tcp_server.py is running.")
+    except Exception as e:
+        print(f"Error: {e}")
+
+
+if __name__ == "__main__":
+    main()
+

dummy_mcp_server.py renamed to examples/generate_traffic/tcp_server.py

@@ -57,3 +57,4 @@ def start_server():
 
 if __name__ == "__main__":
     start_server()
+