Wazuh MCP Server: Complete Implementation Guide with Production Examples#

1
#!/bin/bash
2
# Create project structure
3
mkdir -p wazuh-mcp-deployment/{config,scripts,docker,k8s,tests}
4
cd wazuh-mcp-deployment
5

6
# Setup Python virtual environment
7
python3.11 -m venv venv
8
source venv/bin/activate
9

10
# Install dependencies
11
pip install --upgrade pip
12
pip install git+https://github.com/socfortress/wazuh-mcp-server.git
13
pip install pytest pytest-asyncio aiohttp python-dotenv

1
from pydantic import BaseSettings
2
from typing import Optional
3
import os
4

5
class WazuhMCPSettings(BaseSettings):
6
    """Production configuration for Wazuh MCP Server"""
7

8
    # Wazuh Configuration
9
    wazuh_url: str
10
    wazuh_username: str
11
    wazuh_password: str
12
    wazuh_ssl_verify: bool = True
13
    wazuh_timeout: int = 30
14

15
    # Server Configuration
16
    server_host: str = "0.0.0.0"
17
    server_port: int = 8000
18
    server_workers: int = 4
19

20
    # Security Configuration
21
    jwt_secret: str
22
    jwt_algorithm: str = "HS256"
23
    jwt_expiry: int = 3600
24

25
    # Feature Flags
26
    read_only_mode: bool = False
27
    enable_metrics: bool = True
28
    enable_tracing: bool = True
29

30
    # Rate Limiting
31
    rate_limit_enabled: bool = True
32
    rate_limit_requests: int = 100
33
    rate_limit_period: int = 60
34

35
    class Config:
36
        env_file = ".env"
37
        env_prefix = "MCP_"
38

39
settings = WazuhMCPSettings()

1
import asyncio
2
import logging
3
from aiohttp import web
4
from prometheus_client import make_asgi_app
5
import uvloop
6
from wazuh_mcp_server import WazuhMCPServer
7
from config.settings import settings
8

9
# Use uvloop for better performance
10
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
11

12
class ProductionMCPServer:
13
    def __init__(self):
14
        self.app = web.Application()
15
        self.mcp_server = WazuhMCPServer(settings)
16
        self.setup_routes()
17
        self.setup_middleware()
18

19
    def setup_routes(self):
20
        """Configure API routes"""
21
        self.app.router.add_route('*', '/sse/{path:.*}', self.mcp_server.handle_sse)
22
        self.app.router.add_get('/health', self.health_check)
23
        self.app.router.add_get('/ready', self.readiness_check)
24
        self.app.router.add_get('/metrics', self.metrics_handler)
25

26
    def setup_middleware(self):
27
        """Configure middleware stack"""
28
        self.app.middlewares.append(self.error_middleware)
29
        self.app.middlewares.append(self.auth_middleware)
30
        self.app.middlewares.append(self.rate_limit_middleware)
31
        self.app.middlewares.append(self.logging_middleware)
32

33
    @web.middleware
34
    async def error_middleware(self, request, handler):
35
        """Global error handling"""
36
        try:
37
            return await handler(request)
38
        except web.HTTPException:
39
            raise
40
        except Exception as ex:
41
            logging.error(f"Unhandled exception: {ex}")
42
            return web.json_response(
43
                {"error": "Internal server error"},
44
                status=500
45
            )
46

47
    @web.middleware
48
    async def auth_middleware(self, request, handler):
49
        """JWT authentication"""
50
        if request.path in ['/health', '/ready', '/metrics']:
51
            return await handler(request)
52

53
        token = request.headers.get('Authorization', '').replace('Bearer ', '')
54
        if not self.validate_jwt(token):
55
            return web.json_response(
56
                {"error": "Unauthorized"},
57
                status=401
58
            )
59
        return await handler(request)
60

61
    @web.middleware
62
    async def rate_limit_middleware(self, request, handler):
63
        """Rate limiting implementation"""
64
        if not settings.rate_limit_enabled:
65
            return await handler(request)
66

67
        client_id = request.remote
68
        if self.is_rate_limited(client_id):
69
            return web.json_response(
70
                {"error": "Rate limit exceeded"},
71
                status=429
72
            )
73
        return await handler(request)
74

75
    async def health_check(self, request):
76
        """Kubernetes liveness probe"""
77
        return web.json_response({"status": "healthy"})
78

79
    async def readiness_check(self, request):
80
        """Kubernetes readiness probe"""
81
        try:
82
            await self.mcp_server.check_wazuh_connection()
83
            return web.json_response({"status": "ready"})
84
        except:
85
            return web.json_response(
86
                {"status": "not ready"},
87
                status=503
88
            )
89

90
    def run(self):
91
        """Start the server"""
92
        web.run_app(
93
            self.app,
94
            host=settings.server_host,
95
            port=settings.server_port,
96
            access_log_format='%t %a "%r" %s %b "%{User-Agent}i" %Tf'
97
        )
98

99
if __name__ == "__main__":
100
    server = ProductionMCPServer()
101
    server.run()

1
version: '3.8'
2

3
services:
4
  wazuh-mcp:
5
    build:
6
      context: ..
7
      dockerfile: docker/Dockerfile
8
    container_name: wazuh-mcp-server
9
    restart: unless-stopped
10
    ports:
11
      - "8000:8000"
12
    environment:
13
      MCP_WAZUH_URL: ${WAZUH_URL}
14
      MCP_WAZUH_USERNAME: ${WAZUH_USERNAME}
15
      MCP_WAZUH_PASSWORD: ${WAZUH_PASSWORD}
16
      MCP_JWT_SECRET: ${JWT_SECRET}
17
      MCP_ENABLE_METRICS: "true"
18
      MCP_LOG_LEVEL: "INFO"
19
    volumes:
20
      - ./config:/app/config:ro
21
      - ./logs:/var/log/mcp
22
      - ./certs:/app/certs:ro
23
    networks:
24
      - wazuh-network
25
    depends_on:
26
      - redis
27
      - prometheus
28
    logging:
29
      driver: "json-file"
30
      options:
31
        max-size: "10m"
32
        max-file: "3"
33

34
  redis:
35
    image: redis:7-alpine
36
    container_name: wazuh-mcp-redis
37
    restart: unless-stopped
38
    command: redis-server --appendonly yes
39
    volumes:
40
      - redis-data:/data
41
    networks:
42
      - wazuh-network
43

44
  prometheus:
45
    image: prom/prometheus:latest
46
    container_name: wazuh-mcp-prometheus
47
    restart: unless-stopped
48
    volumes:
49
      - ./prometheus.yml:/etc/prometheus/prometheus.yml:ro
50
      - prometheus-data:/prometheus
51
    command:
52
      - '--config.file=/etc/prometheus/prometheus.yml'
53
      - '--storage.tsdb.path=/prometheus'
54
    ports:
55
      - "9090:9090"
56
    networks:
57
      - wazuh-network
58

59
  grafana:
60
    image: grafana/grafana:latest
61
    container_name: wazuh-mcp-grafana
62
    restart: unless-stopped
63
    environment:
64
      GF_SECURITY_ADMIN_PASSWORD: ${GRAFANA_PASSWORD}
65
      GF_INSTALL_PLUGINS: grafana-piechart-panel
66
    volumes:
67
      - grafana-data:/var/lib/grafana
68
      - ./grafana/dashboards:/etc/grafana/provisioning/dashboards:ro
69
      - ./grafana/datasources:/etc/grafana/provisioning/datasources:ro
70
    ports:
71
      - "3000:3000"
72
    networks:
73
      - wazuh-network
74
    depends_on:
75
      - prometheus
76

77
volumes:
78
  redis-data:
79
  prometheus-data:
80
  grafana-data:
81

82
networks:
83
  wazuh-network:
84
    driver: bridge

1
apiVersion: v1
2
kind: Namespace
3
metadata:
4
  name: wazuh-mcp
5
  labels:
6
    name: wazuh-mcp
7
    environment: production
8
---
9
# k8s/configmap.yaml
10
apiVersion: v1
11
kind: ConfigMap
12
metadata:
13
  name: wazuh-mcp-config
14
  namespace: wazuh-mcp
15
data:
16
  server.conf: |
17
    server:
18
      host: 0.0.0.0
19
      port: 8000
20
      workers: 4
21

22
    security:
23
      read_only_mode: false
24
      enable_audit: true
25

26
    features:
27
      enable_metrics: true
28
      enable_tracing: true
29
      enable_profiling: false
30
---
31
# k8s/secret.yaml
32
apiVersion: v1
33
kind: Secret
34
metadata:
35
  name: wazuh-mcp-secret
36
  namespace: wazuh-mcp
37
type: Opaque
38
stringData:
39
  wazuh-url: "https://wazuh.example.com:55000"
40
  wazuh-username: "admin"
41
  wazuh-password: "SecurePassword123!"
42
  jwt-secret: "your-jwt-secret-key"
43
---
44
# k8s/deployment.yaml
45
apiVersion: apps/v1
46
kind: Deployment
47
metadata:
48
  name: wazuh-mcp-server
49
  namespace: wazuh-mcp
50
  labels:
51
    app: wazuh-mcp
52
    version: v1
53
spec:
54
  replicas: 3
55
  strategy:
56
    type: RollingUpdate
57
    rollingUpdate:
58
      maxSurge: 1
59
      maxUnavailable: 0
60
  selector:
61
    matchLabels:
62
      app: wazuh-mcp
63
  template:
64
    metadata:
65
      labels:
66
        app: wazuh-mcp
67
        version: v1
68
      annotations:
69
        prometheus.io/scrape: "true"
70
        prometheus.io/port: "8000"
71
        prometheus.io/path: "/metrics"
72
    spec:
73
      securityContext:
74
        runAsNonRoot: true
75
        runAsUser: 1000
76
        fsGroup: 1000
77
      containers:
78
      - name: wazuh-mcp
79
        image: your-registry/wazuh-mcp-server:latest
80
        imagePullPolicy: Always
81
        ports:
82
        - containerPort: 8000
83
          name: http
84
          protocol: TCP
85
        env:
86
        - name: MCP_WAZUH_URL
87
          valueFrom:
88
            secretKeyRef:
89
              name: wazuh-mcp-secret
90
              key: wazuh-url
91
        - name: MCP_WAZUH_USERNAME
92
          valueFrom:
93
            secretKeyRef:
94
              name: wazuh-mcp-secret
95
              key: wazuh-username
96
        - name: MCP_WAZUH_PASSWORD
97
          valueFrom:
98
            secretKeyRef:
99
              name: wazuh-mcp-secret
100
              key: wazuh-password
101
        - name: MCP_JWT_SECRET
102
          valueFrom:
103
            secretKeyRef:
104
              name: wazuh-mcp-secret
105
              key: jwt-secret
106
        resources:
107
          requests:
108
            memory: "256Mi"
109
            cpu: "250m"
110
          limits:
111
            memory: "1Gi"
112
            cpu: "1000m"
113
        livenessProbe:
114
          httpGet:
115
            path: /health
116
            port: 8000
117
          initialDelaySeconds: 30
118
          periodSeconds: 10
119
          timeoutSeconds: 5
120
          failureThreshold: 3
121
        readinessProbe:
122
          httpGet:
123
            path: /ready
124
            port: 8000
125
          initialDelaySeconds: 10
126
          periodSeconds: 5
127
          timeoutSeconds: 3
128
          failureThreshold: 3
129
        volumeMounts:
130
        - name: config
131
          mountPath: /app/config
132
          readOnly: true
133
        - name: tmp
134
          mountPath: /tmp
135
      volumes:
136
      - name: config
137
        configMap:
138
          name: wazuh-mcp-config
139
      - name: tmp
140
        emptyDir: {}
141
      affinity:
142
        podAntiAffinity:
143
          preferredDuringSchedulingIgnoredDuringExecution:
144
          - weight: 100
145
            podAffinityTerm:
146
              labelSelector:
147
                matchExpressions:
148
                - key: app
149
                  operator: In
150
                  values:
151
                  - wazuh-mcp
152
              topologyKey: kubernetes.io/hostname
153
---
154
# k8s/service.yaml
155
apiVersion: v1
156
kind: Service
157
metadata:
158
  name: wazuh-mcp-service
159
  namespace: wazuh-mcp
160
  labels:
161
    app: wazuh-mcp
162
spec:
163
  type: ClusterIP
164
  selector:
165
    app: wazuh-mcp
166
  ports:
167
  - port: 8000
168
    targetPort: 8000
169
    protocol: TCP
170
    name: http
171
---
172
# k8s/ingress.yaml
173
apiVersion: networking.k8s.io/v1
174
kind: Ingress
175
metadata:
176
  name: wazuh-mcp-ingress
177
  namespace: wazuh-mcp
178
  annotations:
179
    cert-manager.io/cluster-issuer: letsencrypt-prod
180
    nginx.ingress.kubernetes.io/ssl-redirect: "true"
181
    nginx.ingress.kubernetes.io/rate-limit: "100"
182
spec:
183
  ingressClassName: nginx
184
  tls:
185
  - hosts:
186
    - mcp.example.com
187
    secretName: wazuh-mcp-tls
188
  rules:
189
  - host: mcp.example.com
190
    http:
191
      paths:
192
      - path: /
193
        pathType: Prefix
194
        backend:
195
          service:
196
            name: wazuh-mcp-service
197
            port:
198
              number: 8000

1
apiVersion: autoscaling/v2
2
kind: HorizontalPodAutoscaler
3
metadata:
4
  name: wazuh-mcp-hpa
5
  namespace: wazuh-mcp
6
spec:
7
  scaleTargetRef:
8
    apiVersion: apps/v1
9
    kind: Deployment
10
    name: wazuh-mcp-server
11
  minReplicas: 2
12
  maxReplicas: 10
13
  metrics:
14
  - type: Resource
15
    resource:
16
      name: cpu
17
      target:
18
        type: Utilization
19
        averageUtilization: 70
20
  - type: Resource
21
    resource:
22
      name: memory
23
      target:
24
        type: Utilization
25
        averageUtilization: 80
26
  - type: Pods
27
    pods:
28
      metric:
29
        name: mcp_requests_per_second
30
      target:
31
        type: AverageValue
32
        averageValue: "1000"
33
  behavior:
34
    scaleDown:
35
      stabilizationWindowSeconds: 300
36
      policies:
37
      - type: Percent
38
        value: 50
39
        periodSeconds: 60
40
    scaleUp:
41
      stabilizationWindowSeconds: 0
42
      policies:
43
      - type: Percent
44
        value: 100
45
        periodSeconds: 15
46
      - type: Pods
47
        value: 2
48
        periodSeconds: 15

1
import asyncio
2
from datetime import datetime, timedelta
3
from typing import Dict, List
4
import json
5
from wazuh_mcp_client import WazuhMCPClient
6
from langchain_openai import ChatOpenAI
7
from langchain.agents import create_openai_tools_agent
8

9
class IncidentResponseAutomation:
10
    """Automated incident response using Wazuh MCP Server"""
11

12
    def __init__(self, mcp_url: str, openai_key: str):
13
        self.mcp_client = WazuhMCPClient(mcp_url)
14
        self.llm = ChatOpenAI(
15
            model="gpt-4",
16
            temperature=0,
17
            api_key=openai_key
18
        )
19
        self.setup_agent()
20

21
    def setup_agent(self):
22
        """Initialize AI agent with MCP tools"""
23
        tools = self.mcp_client.get_tools()
24
        self.agent = create_openai_tools_agent(self.llm, tools)
25

26
    async def investigate_security_event(self, event: Dict) -> Dict:
27
        """Comprehensive security event investigation"""
28

29
        investigation_steps = {
30
            "initial_triage": await self.triage_event(event),
31
            "agent_analysis": await self.analyze_affected_agent(event),
32
            "lateral_movement": await self.check_lateral_movement(event),
33
            "persistence_check": await self.check_persistence_mechanisms(event),
34
            "network_analysis": await self.analyze_network_connections(event),
35
            "process_tree": await self.analyze_process_tree(event),
36
            "file_analysis": await self.analyze_file_modifications(event),
37
            "threat_enrichment": await self.enrich_with_threat_intel(event)
38
        }
39

40
        # Generate comprehensive report
41
        report = await self.generate_investigation_report(
42
            event,
43
            investigation_steps
44
        )
45

46
        # Determine response actions
47
        response_actions = await self.determine_response_actions(report)
48

49
        # Execute automated responses
50
        if report['severity'] >= 8:
51
            await self.execute_critical_response(event, response_actions)
52

53
        return {
54
            "event": event,
55
            "investigation": investigation_steps,
56
            "report": report,
57
            "actions_taken": response_actions,
58
            "timestamp": datetime.utcnow().isoformat()
59
        }
60

61
    async def triage_event(self, event: Dict) -> Dict:
62
        """Initial event triage and classification"""
63

64
        triage_prompt = f"""
65
        Analyze this security event and provide triage information:
66

67
        Event: {json.dumps(event, indent=2)}
68

69
        Please determine:
70
        1. Severity (1-10)
71
        2. Event category (malware, intrusion, data_exfiltration, etc.)
72
        3. Affected assets
73
        4. Immediate risk assessment
74
        5. Required investigation priority
75
        """
76

77
        result = await self.agent.ainvoke({"input": triage_prompt})
78
        return self.parse_triage_result(result)
79

80
    async def analyze_affected_agent(self, event: Dict) -> Dict:
81
        """Deep analysis of affected Wazuh agent"""
82

83
        agent_id = event.get('agent', {}).get('id')
84

85
        analysis_queries = [
86
            f"Get detailed information about agent {agent_id}",
87
            f"Show all running processes on agent {agent_id}",
88
            f"List network connections for agent {agent_id}",
89
            f"Check installed packages on agent {agent_id}",
90
            f"Review recent file modifications on agent {agent_id}"
91
        ]
92

93
        results = {}
94
        for query in analysis_queries:
95
            result = await self.agent.ainvoke({"input": query})
96
            results[query] = result
97

98
        return results
99

100
    async def check_lateral_movement(self, event: Dict) -> Dict:
101
        """Check for potential lateral movement indicators"""
102

103
        source_ip = event.get('data', {}).get('srcip')
104
        timeframe = (datetime.utcnow() - timedelta(hours=24)).isoformat()
105

106
        lateral_movement_query = f"""
107
        Find all agents that have connections from {source_ip} since {timeframe}.
108
        For each agent, check for:
109
        1. New user accounts created
110
        2. Unusual authentication patterns
111
        3. Service installations
112
        4. Registry modifications
113
        5. Scheduled task creations
114
        """
115

116
        result = await self.agent.ainvoke({"input": lateral_movement_query})
117
        return self.analyze_lateral_movement_indicators(result)
118

119
    async def check_persistence_mechanisms(self, event: Dict) -> Dict:
120
        """Identify potential persistence mechanisms"""
121

122
        agent_id = event.get('agent', {}).get('id')
123

124
        persistence_query = f"""
125
        On agent {agent_id}, check for persistence mechanisms:
126
        1. Startup folder items
127
        2. Registry run keys
128
        3. Scheduled tasks
129
        4. Services set to auto-start
130
        5. WMI event subscriptions
131
        6. Modified system binaries
132
        7. Suspicious kernel modules
133
        """
134

135
        result = await self.agent.ainvoke({"input": persistence_query})
136
        return self.parse_persistence_findings(result)
137

138
    async def execute_critical_response(self, event: Dict, actions: List[str]):
139
        """Execute automated response for critical events"""
140

141
        agent_id = event.get('agent', {}).get('id')
142

143
        response_actions = {
144
            "isolate_host": f"Isolate agent {agent_id} from network",
145
            "kill_process": f"Terminate malicious process on agent {agent_id}",
146
            "block_ip": f"Block IP address in firewall",
147
            "disable_account": f"Disable compromised user account",
148
            "collect_forensics": f"Collect forensic data from agent {agent_id}"
149
        }
150

151
        executed_actions = []
152
        for action in actions:
153
            if action in response_actions:
154
                # Execute via MCP
155
                result = await self.agent.ainvoke({
156
                    "input": response_actions[action]
157
                })
158
                executed_actions.append({
159
                    "action": action,
160
                    "result": result,
161
                    "timestamp": datetime.utcnow().isoformat()
162
                })
163

164
        # Send notification
165
        await self.send_critical_alert_notification(event, executed_actions)
166

167
        return executed_actions

1
import asyncio
2
from typing import Dict, List
3
import pandas as pd
4
from datetime import datetime
5

6
class ComplianceMonitor:
7
    """Automated compliance monitoring using Wazuh MCP"""
8

9
    def __init__(self, mcp_client, frameworks: List[str]):
10
        self.mcp = mcp_client
11
        self.frameworks = frameworks  # ['PCI-DSS', 'HIPAA', 'CIS']
12
        self.compliance_rules = self.load_compliance_rules()
13

14
    async def run_compliance_scan(self) -> Dict:
15
        """Execute comprehensive compliance scan"""
16

17
        results = {
18
            "scan_id": self.generate_scan_id(),
19
            "timestamp": datetime.utcnow().isoformat(),
20
            "frameworks": {},
21
            "summary": {}
22
        }
23

24
        # Get all agents
25
        agents = await self.get_all_agents()
26

27
        for framework in self.frameworks:
28
            framework_results = await self.scan_framework_compliance(
29
                framework,
30
                agents
31
            )
32
            results["frameworks"][framework] = framework_results
33

34
        # Generate summary
35
        results["summary"] = self.generate_compliance_summary(results)
36

37
        # Create remediation plan
38
        results["remediation"] = await self.create_remediation_plan(results)
39

40
        return results
41

42
    async def scan_framework_compliance(
43
        self,
44
        framework: str,
45
        agents: List[Dict]
46
    ) -> Dict:
47
        """Scan compliance for specific framework"""
48

49
        framework_rules = self.compliance_rules[framework]
50
        results = {
51
            "compliant": [],
52
            "non_compliant": [],
53
            "warnings": [],
54
            "not_applicable": []
55
        }
56

57
        for rule in framework_rules:
58
            rule_result = await self.check_compliance_rule(rule, agents)
59

60
            if rule_result["status"] == "compliant":
61
                results["compliant"].append(rule_result)
62
            elif rule_result["status"] == "non_compliant":
63
                results["non_compliant"].append(rule_result)
64
            elif rule_result["status"] == "warning":
65
                results["warnings"].append(rule_result)
66
            else:
67
                results["not_applicable"].append(rule_result)
68

69
        return results
70

71
    async def check_compliance_rule(
72
        self,
73
        rule: Dict,
74
        agents: List[Dict]
75
    ) -> Dict:
76
        """Check specific compliance rule across agents"""
77

78
        rule_query = self.build_compliance_query(rule)
79

80
        result = await self.mcp.agent.ainvoke({
81
            "input": rule_query
82
        })
83

84
        return {
85
            "rule_id": rule["id"],
86
            "description": rule["description"],
87
            "status": self.evaluate_compliance_status(result),
88
            "affected_agents": self.extract_affected_agents(result),
89
            "evidence": result,
90
            "remediation": rule.get("remediation")
91
        }
92

93
    def build_compliance_query(self, rule: Dict) -> str:
94
        """Build MCP query for compliance rule"""
95

96
        query_templates = {
97
            "password_policy": """
98
                Check password policy configuration on all Windows agents:
99
                - Minimum password length
100
                - Password complexity requirements
101
                - Password history
102
                - Maximum password age
103
            """,
104

105
            "encryption": """
106
                Verify encryption status on all agents:
107
                - Disk encryption status
108
                - TLS/SSL configuration
109
                - Certificate validation
110
                - Encrypted communication channels
111
            """,
112

113
            "logging": """
114
                Validate logging configuration:
115
                - Audit logging enabled
116
                - Log retention period
117
                - Log forwarding configuration
118
                - Critical event logging
119
            """,
120

121
            "access_control": """
122
                Review access control settings:
123
                - User permissions
124
                - Service accounts
125
                - Administrative access
126
                - Privilege escalation controls
127
            """
128
        }
129

130
        return query_templates.get(
131
            rule["category"],
132
            f"Check compliance for: {rule['description']}"
133
        )
134

135
    async def create_remediation_plan(self, scan_results: Dict) -> Dict:
136
        """Generate automated remediation plan"""
137

138
        remediation_plan = {
139
            "priority_high": [],
140
            "priority_medium": [],
141
            "priority_low": [],
142
            "automated_fixes": []
143
        }
144

145
        for framework, results in scan_results["frameworks"].items():
146
            for finding in results["non_compliant"]:
147
                remediation = {
148
                    "framework": framework,
149
                    "rule": finding["rule_id"],
150
                    "description": finding["description"],
151
                    "affected_agents": finding["affected_agents"],
152
                    "remediation_steps": finding.get("remediation", []),
153
                    "can_automate": self.can_automate_fix(finding)
154
                }
155

156
                # Prioritize based on severity
157
                if finding.get("severity", "medium") == "high":
158
                    remediation_plan["priority_high"].append(remediation)
159
                elif finding.get("severity") == "low":
160
                    remediation_plan["priority_low"].append(remediation)
161
                else:
162
                    remediation_plan["priority_medium"].append(remediation)
163

164
                # Add to automated fixes if possible
165
                if remediation["can_automate"]:
166
                    remediation_plan["automated_fixes"].append(remediation)
167

168
        return remediation_plan

1
import asyncio
2
from typing import Dict, List, Optional
3
from datetime import datetime, timedelta
4
import re
5

6
class ThreatHunter:
7
    """Automated threat hunting using Wazuh MCP and AI"""
8

9
    def __init__(self, mcp_client, threat_feeds: List[str]):
10
        self.mcp = mcp_client
11
        self.threat_feeds = threat_feeds
12
        self.hunting_patterns = self.load_hunting_patterns()
13

14
    async def hunt_threats(self, timeframe_hours: int = 24) -> Dict:
15
        """Execute automated threat hunting"""
16

17
        hunt_results = {
18
            "hunt_id": self.generate_hunt_id(),
19
            "start_time": datetime.utcnow().isoformat(),
20
            "timeframe": f"{timeframe_hours} hours",
21
            "hunts_executed": [],
22
            "threats_found": [],
23
            "suspicious_activity": [],
24
            "recommendations": []
25
        }
26

27
        # Execute different hunting techniques
28
        techniques = [
29
            self.hunt_living_off_the_land(),
30
            self.hunt_data_exfiltration(),
31
            self.hunt_privilege_escalation(),
32
            self.hunt_command_and_control(),
33
            self.hunt_defense_evasion(),
34
            self.hunt_anomalous_behavior()
35
        ]
36

37
        for technique in techniques:
38
            result = await technique
39
            hunt_results["hunts_executed"].append(result)
40

41
            if result["threats_found"]:
42
                hunt_results["threats_found"].extend(result["threats_found"])
43

44
            if result["suspicious_activity"]:
45
                hunt_results["suspicious_activity"].extend(
46
                    result["suspicious_activity"]
47
                )
48

49
        # Correlate findings
50
        hunt_results["correlations"] = await self.correlate_findings(
51
            hunt_results
52
        )
53

54
        # Generate recommendations
55
        hunt_results["recommendations"] = await self.generate_recommendations(
56
            hunt_results
57
        )
58

59
        return hunt_results
60

61
    async def hunt_living_off_the_land(self) -> Dict:
62
        """Hunt for Living off the Land techniques"""
63

64
        lolbins = [
65
            "certutil.exe", "bitsadmin.exe", "powershell.exe",
66
            "wmic.exe", "mshta.exe", "rundll32.exe", "regsvr32.exe",
67
            "cscript.exe", "wscript.exe", "msiexec.exe"
68
        ]
69

70
        hunt_query = f"""
71
        Find all processes using these binaries: {', '.join(lolbins)}
72
        For each process, analyze:
73
        1. Command line arguments for suspicious patterns
74
        2. Parent process relationships
75
        3. Network connections made
76
        4. Files accessed or created
77
        5. Registry modifications
78

79
        Flag as suspicious if:
80
        - Downloading from external URLs
81
        - Base64 encoded commands
82
        - Unusual parent-child relationships
83
        - Execution from temp directories
84
        """
85

86
        result = await self.mcp.agent.ainvoke({"input": hunt_query})
87

88
        return self.analyze_lolbin_activity(result)
89

90
    async def hunt_data_exfiltration(self) -> Dict:
91
        """Hunt for data exfiltration indicators"""
92

93
        exfil_query = """
94
        Analyze network traffic patterns for data exfiltration:
95
        1. Large outbound data transfers to unusual destinations
96
        2. Connections to cloud storage services
97
        3. DNS tunneling patterns (high DNS query volume)
98
        4. Encrypted channels to suspicious IPs
99
        5. Data staging in archive files
100

101
        Check for:
102
        - Compressed files created in last 24 hours
103
        - Unusual protocols on standard ports
104
        - Beaconing behavior
105
        - Data transfers during off-hours
106
        """
107

108
        result = await self.mcp.agent.ainvoke({"input": exfil_query})
109

110
        return self.analyze_exfiltration_patterns(result)
111

112
    async def hunt_privilege_escalation(self) -> Dict:
113
        """Hunt for privilege escalation attempts"""
114

115
        privesc_query = """
116
        Search for privilege escalation indicators:
117
        1. Process token manipulation
118
        2. Service creation or modification
119
        3. Scheduled task creation with SYSTEM privileges
120
        4. UAC bypass attempts
121
        5. Kernel exploit indicators
122
        6. SUID/SGID binary execution (Linux)
123
        7. Sudo configuration changes
124

125
        Look for:
126
        - Processes running with unexpected privileges
127
        - Recent privilege changes
128
        - Failed authentication followed by success
129
        - Suspicious service installations
130
        """
131

132
        result = await self.mcp.agent.ainvoke({"input": privesc_query})
133

134
        return self.analyze_privilege_escalation(result)
135

136
    async def hunt_anomalous_behavior(self) -> Dict:
137
        """Hunt using behavioral analytics"""
138

139
        anomaly_query = """
140
        Identify anomalous behavior patterns:
141
        1. Users accessing systems outside normal hours
142
        2. Processes with unusual resource consumption
143
        3. Rare process executions
144
        4. Abnormal network traffic patterns
145
        5. Unusual file access patterns
146
        6. Services starting at unusual times
147

148
        Compare against baseline behavior and flag deviations
149
        """
150

151
        result = await self.mcp.agent.ainvoke({"input": anomaly_query})
152

153
        return self.analyze_anomalies(result)
154

155
    async def correlate_findings(self, hunt_results: Dict) -> List[Dict]:
156
        """Correlate hunting findings to identify attack chains"""
157

158
        correlation_prompt = f"""
159
        Analyze these threat hunting findings and identify potential attack chains:
160

161
        Findings: {hunt_results}
162

163
        Look for:
164
        1. Related events across different hunt techniques
165
        2. Temporal relationships between events
166
        3. Common actors (users, processes, IPs)
167
        4. MITRE ATT&CK technique chains
168
        5. Kill chain progression
169

170
        Provide correlation confidence scores and evidence
171
        """
172

173
        correlations = await self.mcp.agent.ainvoke({
174
            "input": correlation_prompt
175
        })
176

177
        return self.parse_correlations(correlations)