Gemini MCP Server

#!/usr/bin/env python3 """ SECAUDIT Tool Validation Test Tests the secaudit tool's capabilities using the workflow architecture. This validates that the workflow-based security audit provides step-by-step analysis with proper investigation guidance and expert analysis integration. """ import json from .conversation_base_test import ConversationBaseTest class SecauditValidationTest(ConversationBaseTest): """Test secaudit tool with workflow architecture""" @property def test_name(self) -> str: return "secaudit_validation" @property def test_description(self) -> str: return "SECAUDIT tool validation with security audit workflow architecture" def run_test(self) -> bool: """Test secaudit tool capabilities""" # Set up the test environment self.setUp() try: self.logger.info("Test: SECAUDIT tool validation (security workflow architecture)") # Create test code with various security vulnerabilities self._create_test_code_for_audit() # Test 1: Single audit session with multiple steps if not self._test_single_audit_session(): return False # Test 2: Audit with specific focus areas if not self._test_focused_security_audit(): return False # Test 3: Complete audit with expert analysis using fast model if not self._test_complete_audit_with_analysis(): return False # Test 4: Certain confidence behavior if not self._test_certain_confidence(): return False # Test 5: Continuation test with chat tool if not self._test_continuation_with_chat(): return False # Test 6: Model selection control if not self._test_model_selection(): return False self.logger.info(" ✅ All secaudit validation tests passed") return True except Exception as e: self.logger.error(f"SECAUDIT validation test failed: {e}") return False def _create_test_code_for_audit(self): """Create test files with various security vulnerabilities""" # Create an authentication module with multiple security issues auth_code = """#!/usr/bin/env python3 import hashlib import pickle import sqlite3 from flask import request, session class AuthenticationManager: def __init__(self, db_path="users.db"): # A01: Broken Access Control - No proper session management self.db_path = db_path self.sessions = {} # In-memory session storage def login(self, username, password): '''User login with various security vulnerabilities''' # A03: Injection - SQL injection vulnerability conn = sqlite3.connect(self.db_path) cursor = conn.cursor() # Direct string interpolation in SQL query query = f"SELECT id, password_hash FROM users WHERE username = '{username}'" cursor.execute(query) user = cursor.fetchone() if not user: return {"status": "failed", "message": "User not found"} # A02: Cryptographic Failures - Weak hashing algorithm password_hash = hashlib.md5(password.encode()).hexdigest() if user[1] == password_hash: # A07: Identification and Authentication Failures - Weak session generation session_id = hashlib.md5(f"{username}{password}".encode()).hexdigest() self.sessions[session_id] = {"user_id": user[0], "username": username} return {"status": "success", "session_id": session_id} else: return {"status": "failed", "message": "Invalid password"} def reset_password(self, email): '''Password reset with security issues''' # A04: Insecure Design - No rate limiting or validation reset_token = hashlib.md5(email.encode()).hexdigest() # A09: Security Logging and Monitoring Failures - No security event logging # Simply returns token without any verification or logging return {"reset_token": reset_token, "url": f"/reset?token={reset_token}"} def deserialize_user_data(self, data): '''Unsafe deserialization''' # A08: Software and Data Integrity Failures - Insecure deserialization return pickle.loads(data) def get_user_profile(self, user_id): '''Get user profile with authorization issues''' # A01: Broken Access Control - No authorization check conn = sqlite3.connect(self.db_path) cursor = conn.cursor() # Fetches any user profile without checking permissions cursor.execute("SELECT * FROM users WHERE id = ?", (user_id,)) return cursor.fetchone() """ # Create authentication file self.auth_file = self.create_additional_test_file("auth_manager.py", auth_code) self.logger.info(f" ✅ Created authentication file with security issues: {self.auth_file}") # Create API endpoint with additional vulnerabilities api_code = """#!/usr/bin/env python3 from flask import Flask, request, jsonify import os import subprocess import requests app = Flask(__name__) # A05: Security Misconfiguration - Debug mode enabled app.config['DEBUG'] = True app.config['SECRET_KEY'] = 'dev-secret-key' # Hardcoded secret @app.route('/api/search', methods=['GET']) def search(): '''Search endpoint with multiple vulnerabilities''' # A03: Injection - XSS vulnerability, no input sanitization query = request.args.get('q', '') # A03: Injection - Command injection vulnerability if 'file:' in query: filename = query.split('file:')[1] # Direct command execution result = subprocess.run(f"cat {filename}", shell=True, capture_output=True, text=True) return jsonify({"result": result.stdout}) # A10: Server-Side Request Forgery (SSRF) if query.startswith('http'): # No validation of URL, allows internal network access response = requests.get(query) return jsonify({"content": response.text}) # Return search results without output encoding return f"<h1>Search Results for: {query}</h1>" @app.route('/api/admin', methods=['GET']) def admin_panel(): '''Admin panel with broken access control''' # A01: Broken Access Control - No authentication check # Anyone can access admin functionality action = request.args.get('action') if action == 'delete_user': user_id = request.args.get('user_id') # Performs privileged action without authorization return jsonify({"status": "User deleted", "user_id": user_id}) return jsonify({"status": "Admin panel"}) @app.route('/api/upload', methods=['POST']) def upload_file(): '''File upload with security issues''' # A05: Security Misconfiguration - No file type validation file = request.files.get('file') if file: # Saves any file type to server filename = file.filename file.save(os.path.join('/tmp', filename)) # A03: Path traversal vulnerability return jsonify({"status": "File uploaded", "path": f"/tmp/{filename}"}) return jsonify({"error": "No file provided"}) # A06: Vulnerable and Outdated Components # Using old Flask version with known vulnerabilities (hypothetical) # requirements.txt: Flask==0.12.2 (known security issues) if __name__ == '__main__': # A05: Security Misconfiguration - Running on all interfaces app.run(host='0.0.0.0', port=5000, debug=True) """ # Create API file self.api_file = self.create_additional_test_file("api_endpoints.py", api_code) self.logger.info(f" ✅ Created API file with security vulnerabilities: {self.api_file}") def _test_single_audit_session(self) -> bool: """Test a single security audit session with multiple steps""" self.logger.info(" 🔧 Testing single audit session...") try: # Step 1: Initial security audit request response, continuation_id = self.call_mcp_tool_direct( "secaudit", { "step": f"Begin security audit of authentication system in {self.auth_file}", "step_number": 1, "total_steps": 6, "next_step_required": True, "findings": "Starting security assessment", "relevant_files": [self.auth_file], "model": "gemini-2.0-flash-lite", }, ) if not response: self.logger.error("Failed to call secaudit tool") return False # Parse and validate the response try: response_data = json.loads(response) if response else {} except json.JSONDecodeError: response_data = {} # Check if it's asking for investigation status = response_data.get("status", "") if status != "pause_for_secaudit": self.logger.error(f"Expected pause_for_secaudit status, got: {status}") return False # Step 2: Continue with findings response2, _ = self.call_mcp_tool_direct( "secaudit", { "step": "Examined authentication module and found critical security vulnerabilities", "step_number": 2, "total_steps": 6, "next_step_required": True, "findings": ( "Found multiple OWASP Top 10 vulnerabilities: " "1. SQL injection in login method (line 88) - direct string interpolation in query " "2. Weak MD5 hashing for passwords (line 96) - cryptographically broken " "3. Insecure session management (line 100) - predictable session IDs " "4. Unsafe deserialization (line 119) - pickle.loads without validation" ), "files_checked": [self.auth_file], "relevant_files": [self.auth_file], "relevant_context": ["AuthenticationManager.login", "AuthenticationManager.deserialize_user_data"], "issues_found": [ {"severity": "critical", "description": "SQL injection vulnerability in login method"}, {"severity": "high", "description": "Weak MD5 password hashing"}, {"severity": "high", "description": "Insecure session management"}, {"severity": "critical", "description": "Unsafe deserialization vulnerability"}, ], "confidence": "medium", "continuation_id": continuation_id, "model": "gemini-2.0-flash-lite", }, ) if not response2: self.logger.error("Failed to continue to step 2") return False self.logger.info(" ✅ Single audit session test passed") return True except Exception as e: self.logger.error(f"Single audit session test failed: {e}") return False def _test_focused_security_audit(self) -> bool: """Test security audit with specific focus areas""" self.logger.info(" 🔧 Testing focused security audit...") try: # Request OWASP-focused audit response, continuation_id = self.call_mcp_tool_direct( "secaudit", { "step": f"Begin OWASP-focused security audit of {self.api_file}", "step_number": 1, "total_steps": 4, "next_step_required": True, "findings": "Starting OWASP Top 10 focused security assessment", "relevant_files": [self.api_file], "security_scope": "Web API endpoints", "threat_level": "high", "audit_focus": "owasp", "model": "gemini-2.0-flash-lite", }, ) if not response: self.logger.error("Failed to start OWASP-focused audit") return False # Verify the audit was configured correctly try: response_data = json.loads(response) # The tool should acknowledge the OWASP focus if response_data.get("status") == "pause_for_secaudit": self.logger.info(" ✅ Focused security audit test passed") return True except json.JSONDecodeError: pass self.logger.error("Expected proper OWASP-focused configuration") return False except Exception as e: self.logger.error(f"Focused security audit test failed: {e}") return False def _test_complete_audit_with_analysis(self) -> bool: """Test complete security audit with expert analysis""" self.logger.info(" 🔧 Testing complete audit with expert analysis...") try: # Step 1: Start fresh audit response1, continuation_id = self.call_mcp_tool_direct( "secaudit", { "step": f"Begin comprehensive security audit of {self.auth_file} and {self.api_file}", "step_number": 1, "total_steps": 3, "next_step_required": True, "findings": "Starting OWASP Top 10 security assessment of authentication and API modules", "relevant_files": [self.auth_file, self.api_file], "security_scope": "Web application with authentication and API endpoints", "model": "gemini-2.0-flash-lite", }, ) if not response1: self.logger.error("Failed to start comprehensive audit") return False # Step 2: Continue with detailed findings response2, _ = self.call_mcp_tool_direct( "secaudit", { "step": "Completed comprehensive security investigation of both modules", "step_number": 2, "total_steps": 3, "next_step_required": True, "findings": ( "Found critical OWASP vulnerabilities across both modules: " "A01: Broken Access Control in admin panel, " "A03: SQL injection in login and command injection in search, " "A02: Weak cryptography with MD5 hashing, " "A05: Security misconfiguration with debug mode enabled, " "A07: Weak session management, " "A08: Insecure deserialization, " "A10: SSRF vulnerability in search endpoint" ), "files_checked": [self.auth_file, self.api_file], "relevant_files": [self.auth_file, self.api_file], "relevant_context": [ "AuthenticationManager.login", "AuthenticationManager.deserialize_user_data", "api.search", "api.admin_panel", ], "issues_found": [ {"severity": "critical", "description": "SQL injection in login method"}, {"severity": "critical", "description": "Command injection in search endpoint"}, {"severity": "critical", "description": "SSRF vulnerability allowing internal network access"}, {"severity": "high", "description": "Broken access control on admin panel"}, {"severity": "high", "description": "Insecure deserialization vulnerability"}, {"severity": "high", "description": "XSS vulnerability in search results"}, {"severity": "medium", "description": "Weak MD5 password hashing"}, {"severity": "medium", "description": "Security misconfiguration - debug mode enabled"}, ], "confidence": "high", "continuation_id": continuation_id, "model": "gemini-2.0-flash-lite", }, ) # Final step - skip expert analysis to avoid timeout response3, _ = self.call_mcp_tool_direct( "secaudit", { "step": "Complete security assessment with all vulnerabilities documented", "step_number": 3, "total_steps": 3, "next_step_required": False, "findings": "Security audit complete with 8 vulnerabilities identified across OWASP categories", "files_checked": [self.auth_file, self.api_file], "relevant_files": [self.auth_file, self.api_file], "confidence": "high", # High confidence to trigger expert analysis "continuation_id": continuation_id, "model": "gemini-2.0-flash-lite", }, ) if response3: # Check for expert analysis or completion status try: response_data = json.loads(response3) status = response_data.get("status", "") # Either expert analysis completed or security analysis complete if status in ["complete", "security_analysis_complete"]: self.logger.info(" ✅ Complete audit with expert analysis test passed") return True except json.JSONDecodeError: # If not JSON, check for security content (expert analysis output) if "security" in response3.lower() or "vulnerability" in response3.lower(): self.logger.info(" ✅ Complete audit with expert analysis test passed") return True self.logger.error("Expected expert security analysis or completion") return False except Exception as e: self.logger.error(f"Complete audit with analysis test failed: {e}") return False def _test_certain_confidence(self) -> bool: """Test behavior when confidence is certain""" self.logger.info(" 🔧 Testing certain confidence behavior...") try: # Request with certain confidence response, _ = self.call_mcp_tool_direct( "secaudit", { "step": f"Security audit complete for {self.auth_file}", "step_number": 1, "total_steps": 1, "next_step_required": False, "findings": "Critical SQL injection vulnerability confirmed in login method", "files_checked": [self.auth_file], "relevant_files": [self.auth_file], "issues_found": [ {"severity": "critical", "description": "SQL injection vulnerability in login method"} ], "confidence": "certain", "model": "gemini-2.0-flash-lite", }, ) if not response: self.logger.error("Failed to execute certain confidence test") return False try: response_data = json.loads(response) # With certain confidence, should complete without expert analysis if response_data.get("status") == "security_analysis_complete": self.logger.info(" ✅ Certain confidence correctly completes without expert analysis") return True except json.JSONDecodeError: pass # Check if findings are shown directly response_lower = response.lower() if "sql injection" in response_lower or "vulnerability" in response_lower: self.logger.info(" ✅ Certain confidence shows findings directly") return True self.logger.error("Expected completion or direct findings with certain confidence") return False except Exception as e: self.logger.error(f"Certain confidence test failed: {e}") return False def _test_continuation_with_chat(self) -> bool: """Test continuation functionality with chat tool""" self.logger.info(" 🔧 Testing continuation with chat tool...") try: # First, run a security audit that generates a continuation_id response1, continuation_id = self.call_mcp_tool_direct( "secaudit", { "step": f"Start analyzing {self.auth_file} for authentication vulnerabilities", "step_number": 1, "total_steps": 4, "next_step_required": True, "findings": "Beginning authentication security analysis", "relevant_files": [self.auth_file], "model": "gemini-2.0-flash-lite", }, ) if not response1: self.logger.error("Failed to start audit for continuation test") return False # Extract continuation_id if present if not continuation_id: self.logger.info(" ⚠️ No continuation_id returned, checking response") try: response_data = json.loads(response1) # Look for thread_id in metadata metadata = response_data.get("metadata", {}) continuation_id = metadata.get("thread_id") except json.JSONDecodeError: pass if continuation_id: # Now test using chat tool with continuation chat_response, _ = self.call_mcp_tool_direct( "chat", { "prompt": "Can you tell me more about the SQL injection vulnerability details found in the security audit?", "continuation_id": continuation_id, "model": "gemini-2.0-flash-lite", }, ) if chat_response: self.logger.info(" ✅ Chat tool continuation test passed") return True else: # Without continuation_id, just verify the audit step worked if response1: self.logger.info(" ✅ Audit step completed (continuation test limited)") return True self.logger.error("Expected successful continuation or audit step") return False except Exception as e: self.logger.error(f"Continuation test failed: {e}") return False def _test_model_selection(self) -> bool: """Test model selection and skip expert analysis option""" self.logger.info(" 🔧 Testing model selection control...") try: # Test 1: Explicit model selection response1, _ = self.call_mcp_tool_direct( "secaudit", { "step": f"Analyze {self.api_file} for SSRF vulnerabilities", "step_number": 1, "total_steps": 2, "next_step_required": True, "findings": "Starting SSRF vulnerability analysis", "relevant_files": [self.api_file], "audit_focus": "owasp", "model": "gemini-2.0-flash-lite", }, ) if response1: self.logger.info(" ✅ Model selection recognized") # Test 2: Skip expert analysis response2, _ = self.call_mcp_tool_direct( "secaudit", { "step": f"Complete security investigation of {self.auth_file}", "step_number": 1, "total_steps": 1, "next_step_required": False, "findings": "Security issues documented", "files_checked": [self.auth_file], "relevant_files": [self.auth_file], "confidence": "high", "use_assistant_model": False, # Skip expert analysis "model": "gemini-2.0-flash-lite", }, ) if response2: try: response_data = json.loads(response2) # Should complete without expert analysis if response_data.get("status") == "security_analysis_complete": self.logger.info(" ✅ Skip expert analysis option works") return True except json.JSONDecodeError: pass # Or might just complete the analysis response_lower = response2.lower() if "complete" in response_lower or "security" in response_lower: self.logger.info(" ✅ Analysis performed without expert model") return True self.logger.error("Expected model selection or skip behavior") return False except Exception as e: self.logger.error(f"Model selection test failed: {e}") return False