forked from sevalla-templates/python-demo-mcp-server
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathsecurity_test.py
More file actions
301 lines (249 loc) · 11.8 KB
/
security_test.py
File metadata and controls
301 lines (249 loc) · 11.8 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
#!/usr/bin/env python3
"""
Advanced security testing for MCP server
Tests for common security vulnerabilities and proper input validation
"""
import json
import requests
import time
import sys
from typing import Dict, Any, List
class MCPSecurityTester:
def __init__(self, base_url: str = "http://localhost:8080"):
self.base_url = base_url
self.messages_url = f"{base_url}/messages/"
self.results = []
def log_result(self, test_name: str, status: str, message: str):
"""Log test results"""
result = {
"test": test_name,
"status": status,
"message": message,
"timestamp": time.time()
}
self.results.append(result)
status_emoji = {"PASS": "✅", "FAIL": "❌", "WARN": "⚠️"}.get(status, "ℹ️")
print(f" {status_emoji} {test_name}: {message}")
def create_mcp_message(self, method: str, params: Dict[str, Any] = None) -> Dict[str, Any]:
"""Create a properly formatted MCP message"""
message = {
"jsonrpc": "2.0",
"id": int(time.time()),
"method": method
}
if params:
message["params"] = params
return message
def test_input_validation(self):
"""Test input validation for security vulnerabilities"""
print("\n🔒 Testing Input Validation:")
# Test SQL injection attempts
sql_payloads = [
"'; DROP TABLE users; --",
"' OR '1'='1",
"\"; DELETE FROM data; --"
]
for payload in sql_payloads:
message = self.create_mcp_message("tools/call", {
"name": "get_current_weather",
"arguments": {"city": payload}
})
try:
response = requests.post(
self.messages_url,
json=message,
headers={"Content-Type": "application/json"},
timeout=10
)
if response.status_code == 200:
result = response.json()
# Check if the payload was properly handled
if "error" in str(result).lower() and "sql" not in str(result).lower():
self.log_result("SQL Injection Test", "PASS", f"Payload rejected safely: {payload[:20]}...")
else:
self.log_result("SQL Injection Test", "WARN", f"Payload processed: {payload[:20]}...")
else:
self.log_result("SQL Injection Test", "PASS", f"Request rejected (status {response.status_code})")
except Exception as e:
self.log_result("SQL Injection Test", "PASS", f"Request failed safely: {str(e)[:50]}...")
def test_xss_protection(self):
"""Test XSS protection"""
print("\n🔒 Testing XSS Protection:")
xss_payloads = [
"<script>alert('xss')</script>",
"javascript:alert('xss')",
"<img src=x onerror=alert('xss')>",
"{{7*7}}" # Template injection
]
for payload in xss_payloads:
message = self.create_mcp_message("tools/call", {
"name": "get_current_weather",
"arguments": {"city": payload}
})
try:
response = requests.post(
self.messages_url,
json=message,
headers={"Content-Type": "application/json"},
timeout=10
)
if response.status_code == 200:
result = response.json()
response_text = json.dumps(result)
# Check if dangerous content is escaped or filtered
if "<script>" not in response_text and "javascript:" not in response_text:
self.log_result("XSS Protection", "PASS", f"Payload safely handled: {payload[:20]}...")
else:
self.log_result("XSS Protection", "FAIL", f"Potential XSS vulnerability: {payload[:20]}...")
else:
self.log_result("XSS Protection", "PASS", f"Request rejected (status {response.status_code})")
except Exception as e:
self.log_result("XSS Protection", "PASS", f"Request failed safely: {str(e)[:50]}...")
def test_rate_limiting(self):
"""Test for rate limiting"""
print("\n⏱️ Testing Rate Limiting:")
# Send multiple requests quickly
start_time = time.time()
success_count = 0
for i in range(10):
message = self.create_mcp_message("tools/call", {
"name": "add",
"arguments": {"a": 1, "b": 1}
})
try:
response = requests.post(
self.messages_url,
json=message,
headers={"Content-Type": "application/json"},
timeout=5
)
if response.status_code == 200:
success_count += 1
elif response.status_code == 429: # Too Many Requests
self.log_result("Rate Limiting", "PASS", "Rate limiting detected")
return
except Exception:
pass
elapsed = time.time() - start_time
if success_count == 10 and elapsed < 2:
self.log_result("Rate Limiting", "WARN", f"No rate limiting detected ({success_count} requests in {elapsed:.2f}s)")
else:
self.log_result("Rate Limiting", "INFO", f"Normal processing ({success_count} requests in {elapsed:.2f}s)")
def test_malformed_requests(self):
"""Test handling of malformed requests"""
print("\n🔧 Testing Malformed Request Handling:")
malformed_tests = [
("Invalid JSON", "not valid json"),
("Missing method", {"jsonrpc": "2.0", "id": 1}),
("Invalid JSON-RPC version", {"jsonrpc": "1.0", "method": "test", "id": 1}),
("Huge payload", {"jsonrpc": "2.0", "method": "tools/call", "id": 1, "params": {"data": "x" * 10000}}),
("Null method", {"jsonrpc": "2.0", "method": None, "id": 1}),
]
for test_name, payload in malformed_tests:
try:
if isinstance(payload, str):
response = requests.post(
self.messages_url,
data=payload,
headers={"Content-Type": "application/json"},
timeout=5
)
else:
response = requests.post(
self.messages_url,
json=payload,
headers={"Content-Type": "application/json"},
timeout=5
)
if response.status_code >= 400:
self.log_result(f"Malformed Request ({test_name})", "PASS", f"Properly rejected (status {response.status_code})")
else:
self.log_result(f"Malformed Request ({test_name})", "WARN", f"Accepted malformed request (status {response.status_code})")
except Exception as e:
self.log_result(f"Malformed Request ({test_name})", "PASS", f"Connection failed safely: {str(e)[:50]}...")
def test_server_headers(self):
"""Test security-related HTTP headers"""
print("\n🛡️ Testing Security Headers:")
try:
response = requests.get(f"{self.base_url}/sse", timeout=5)
headers = response.headers
# Check for security headers
security_headers = {
"X-Content-Type-Options": "nosniff",
"X-Frame-Options": "DENY",
"X-XSS-Protection": "1; mode=block",
"Strict-Transport-Security": "max-age=31536000",
"Content-Security-Policy": None, # Any CSP is good
}
for header, expected_value in security_headers.items():
if header in headers:
if expected_value is None or headers[header] == expected_value:
self.log_result(f"Security Header ({header})", "PASS", f"Present: {headers[header]}")
else:
self.log_result(f"Security Header ({header})", "WARN", f"Present but unexpected value: {headers[header]}")
else:
self.log_result(f"Security Header ({header})", "INFO", "Not present (consider adding)")
# Check for information disclosure in Server header
if "Server" in headers:
server_header = headers["Server"]
if "python" in server_header.lower() or "uvicorn" in server_header.lower():
self.log_result("Server Header", "WARN", f"Server information disclosed: {server_header}")
else:
self.log_result("Server Header", "PASS", f"Server header present: {server_header}")
else:
self.log_result("Server Header", "PASS", "Server header not disclosed")
except Exception as e:
self.log_result("Security Headers", "FAIL", f"Could not test headers: {str(e)}")
def run_all_tests(self):
"""Run all security tests"""
print("🔐 MCP Server Advanced Security Testing")
print("=" * 60)
# Check if server is running
try:
response = requests.get(f"{self.base_url}/sse", timeout=5)
if response.status_code != 200:
print(f"❌ Server not accessible at {self.base_url}")
return False
except Exception as e:
print(f"❌ Cannot connect to server: {e}")
return False
print(f"✅ Server accessible at {self.base_url}")
# Run all tests
self.test_server_headers()
self.test_malformed_requests()
self.test_input_validation()
self.test_xss_protection()
self.test_rate_limiting()
# Print summary
self.print_summary()
return True
def print_summary(self):
"""Print test summary"""
print("\n📊 Security Test Summary:")
print("-" * 40)
pass_count = sum(1 for r in self.results if r["status"] == "PASS")
warn_count = sum(1 for r in self.results if r["status"] == "WARN")
fail_count = sum(1 for r in self.results if r["status"] == "FAIL")
info_count = sum(1 for r in self.results if r["status"] == "INFO")
print(f"✅ Passed: {pass_count}")
print(f"⚠️ Warnings: {warn_count}")
print(f"❌ Failed: {fail_count}")
print(f"ℹ️ Info: {info_count}")
if fail_count > 0:
print(f"\n❌ Critical security issues found: {fail_count}")
print("Please review and fix the failed tests.")
elif warn_count > 0:
print(f"\n⚠️ Security warnings: {warn_count}")
print("Consider addressing the warnings for better security.")
else:
print(f"\n✅ All security tests passed!")
def main():
"""Main function"""
base_url = "http://localhost:8080"
if len(sys.argv) > 1:
base_url = sys.argv[1]
tester = MCPSecurityTester(base_url)
success = tester.run_all_tests()
return 0 if success else 1
if __name__ == "__main__":
sys.exit(main())