Wazuh MCP Server

#!/usr/bin/env python3 """ Enhanced test to verify Phase 5.1 incident pipeline with real data processing. """ import sys import os import asyncio # Add src to path sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src')) def test_incident_pipeline_functionality(): """Test enhanced incident pipeline functionality.""" print("Testing enhanced incident pipeline functionality...") try: from wazuh_mcp_server.prompt_enhancement.pipelines import IncidentPipeline from wazuh_mcp_server.prompt_enhancement.cache import AsyncContextCache, ContextCache from wazuh_mcp_server.prompt_enhancement.context_aggregator import ContextRequest # Create pipeline instance cache = AsyncContextCache(ContextCache()) pipeline = IncidentPipeline(None, cache) # None server for testing # Test helper methods with sample data # Test alert trend calculation sample_alerts = [ {'rule': {'level': 5}, 'timestamp': '2024-01-01T10:00:00Z'}, {'rule': {'level': 8}, 'timestamp': '2024-01-01T11:00:00Z'}, {'rule': {'level': 12}, 'timestamp': '2024-01-01T12:00:00Z'} ] trend = pipeline._calculate_alert_trend(sample_alerts) assert trend in ['increasing', 'decreasing', 'stable'] # Test top triggered rules sample_alerts_with_rules = [ {'rule': {'id': '001', 'description': 'Test Rule 1'}}, {'rule': {'id': '001', 'description': 'Test Rule 1'}}, {'rule': {'id': '002', 'description': 'Test Rule 2'}} ] top_rules = pipeline._get_top_triggered_rules(sample_alerts_with_rules) assert len(top_rules) > 0 assert top_rules[0]['count'] == 2 # Rule 001 should have count 2 # Test agent health score calculation sample_agent = {'status': 'active', 'version': '4.8.0'} health_score = pipeline._calculate_agent_health_score(sample_agent, None) assert 0 <= health_score <= 100 # Test vulnerability analysis sample_vulns = [ {'severity': 'Critical', 'cve': 'CVE-2024-0001', 'description': 'Remote code execution'}, {'severity': 'High', 'cve': 'CVE-2024-0002', 'description': 'Buffer overflow'}, {'severity': 'Medium', 'cve': 'CVE-2024-0003', 'description': 'Information disclosure'} ] exploitable_count = len([v for v in sample_vulns if pipeline._is_exploitable_vulnerability(v)]) assert exploitable_count >= 2 # First two should be detected as exploitable cvss_dist = pipeline._analyze_cvss_scores([]) assert isinstance(cvss_dist, dict) # Test process analysis sample_processes = [ {'name': 'systemd', 'cpu': 5.0, 'rss': 50000, 'state': 'R'}, {'name': 'suspicious.exe', 'cpu': 95.0, 'rss': 2000000, 'state': 'R'}, {'name': 'nginx', 'cpu': 10.0, 'rss': 100000, 'state': 'S'} ] system_processes = pipeline._identify_system_processes(sample_processes) assert len(system_processes) >= 2 # systemd and nginx process_anomalies = pipeline._detect_process_anomalies(sample_processes) assert len(process_anomalies) > 0 # Should detect high CPU usage # Test port analysis sample_ports = [ {'local_port': 22, 'protocol': 'tcp', 'state': 'open'}, {'local_port': 4444, 'protocol': 'tcp', 'state': 'open'}, # Suspicious {'local_port': 80, 'protocol': 'tcp', 'state': 'listening'}, {'local_port': 55555, 'protocol': 'tcp', 'state': 'open'} # High port ] critical_ports = pipeline._identify_critical_service_ports(sample_ports) assert len(critical_ports) >= 2 # SSH and HTTP port_anomalies = pipeline._detect_port_anomalies(sample_ports) assert len(port_anomalies) >= 2 # Should detect 4444 and 55555 network_exposure = pipeline._assess_network_exposure(sample_ports) assert network_exposure in ['none', 'low', 'medium', 'high'] print("✓ Enhanced incident pipeline functionality works correctly") return True except Exception as e: print(f"✗ Enhanced incident pipeline test failed: {str(e)}") import traceback traceback.print_exc() return False def test_pattern_matching_enhancement(): """Test enhanced pattern matching capabilities.""" print("Testing enhanced pattern matching...") try: from wazuh_mcp_server.prompt_enhancement.context_aggregator import PromptPatternMatcher matcher = PromptPatternMatcher() # Test various incident detection patterns test_cases = [ ("investigate suspicious activity on agent 001", "incident", ["agent_id"]), ("hunt for lateral movement indicators", "hunting", []), ("check compliance with PCI-DSS requirements", "compliance", []), ("perform forensic analysis of the attack", "forensic", []), ("monitor system performance and health", "monitoring", []) ] for prompt, expected_type, expected_entities in test_cases: analysis = matcher.analyze_prompt(prompt, "test_tool", {"agent_id": "001"}) # Check that expected context type is detected assert expected_type in analysis['context_types'], f"Failed to detect {expected_type} in '{prompt}'" # Check confidence assert analysis['confidence'] > 0, f"No confidence for '{prompt}'" # Check entity extraction for entity_type in expected_entities: assert entity_type in analysis['entities'], f"Failed to extract {entity_type} from '{prompt}'" print("✓ Enhanced pattern matching works correctly") return True except Exception as e: print(f"✗ Enhanced pattern matching test failed: {str(e)}") return False async def test_context_aggregator_enhancement(): """Test context aggregator with enhanced pipeline.""" print("Testing context aggregator with enhanced pipelines...") try: from wazuh_mcp_server.prompt_enhancement.context_aggregator import PromptContextAggregator from wazuh_mcp_server.prompt_enhancement.context_aggregator import ContextRequest # Create mock server class MockServer: def __init__(self): self.config = type('Config', (), { 'enable_prompt_enhancement': True, 'enable_context_aggregation': True, 'context_cache_ttl': 300, 'max_context_size': 1000, 'enhancement_timeout': 5.0 })() # Mock API client self.api_client = type('APIClient', (), { 'get_alerts': lambda *args, **kwargs: {'data': {'affected_items': []}}, 'get_agents': lambda *args, **kwargs: {'data': {'affected_items': []}}, 'get_agent_vulnerabilities': lambda agent_id: {'data': {'affected_items': []}}, 'get_agent_stats': lambda agent_id: {'data': {}}, 'get_agent_processes': lambda agent_id: {'data': {'affected_items': []}}, 'get_agent_ports': lambda agent_id: {'data': {'affected_items': []}} })() # Create aggregator mock_server = MockServer() aggregator = PromptContextAggregator(mock_server) aggregator.setup_pipelines() # Test that pipelines are set up assert 'incident' in aggregator.pipelines assert 'hunting' in aggregator.pipelines assert 'compliance' in aggregator.pipelines assert 'forensic' in aggregator.pipelines # Test pattern analysis request = ContextRequest( prompt="investigate security incident on agent 001", tool_name="get_alerts", arguments={"agent_id": "001"} ) # This should not crash and should return analysis context = await aggregator._gather_context(request) # Context might be None due to mock data, but should not crash print("✓ Context aggregator with enhanced pipelines works correctly") return True except Exception as e: print(f"✗ Context aggregator enhancement test failed: {str(e)}") import traceback traceback.print_exc() return False def main(): """Run all enhanced tests.""" print("Phase 5.1 Enhanced Test Suite") print("=" * 35) tests = [ test_incident_pipeline_functionality, test_pattern_matching_enhancement, ] async_tests = [ test_context_aggregator_enhancement ] passed = 0 total = len(tests) + len(async_tests) # Run sync tests for test_func in tests: if test_func(): passed += 1 print() # Run async tests for test_func in async_tests: try: if asyncio.run(test_func()): passed += 1 except Exception as e: print(f"✗ Async test failed: {str(e)}") print() print(f"Test Results: {passed}/{total} tests passed") print() if passed == total: print("✓ All enhanced tests passed! Phase 5.1 incident pipeline is fully functional.") print("✓ Real data gathering and analysis capabilities are working.") print("✓ Enhanced context aggregation is ready for production use.") return True else: print("✗ Some tests failed. Please check the implementation.") return False if __name__ == "__main__": success = main() sys.exit(0 if success else 1)