Complete Guide to Enhancing Wazuh with Ollama: AI-Powered Cybersecurity Integration#

1
from ollama import chat
2

3
response = chat(
4
    model='llama3.2',
5
    messages=[{
6
        'role': 'user',
7
        'content': 'What is Wazuh?'
8
    }]
9
)
10
print(response['message']['content'])

1
from ollama import Client
2

3
client = Client(host='http://ollama:11434')
4

5
response = client.chat(
6
    model='llama3.2',
7
    messages=[{
8
        'role': 'user',
9
        'content': 'Analyze this IP: 203.0.113.5'
10
    }]
11
)
12
print(response['message']['content'])

1
import os
2

3
# Get Wazuh installation path
4
wazuh_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
5

6
# Socket address
7
socket_addr = f"{wazuh_path}/queue/sockets/queue"
8
# Typically: /var/ossec/queue/sockets/queue

1
from socket import socket, AF_UNIX, SOCK_DGRAM
2
import json
3

4
def send_event(msg, agent=None):
5
    """
6
    Send custom event to Wazuh via Unix socket
7

8
    Args:
9
        msg: Event data (dict)
10
        agent: Agent information (dict with id, name, ip)
11
    """
12
    wazuh_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
13
    socket_addr = f"{wazuh_path}/queue/sockets/queue"
14

15
    if not agent or agent["id"] == "000":
16
        # Manager-level event
17
        string = f"1:ollama:{json.dumps(msg)}"
18
    else:
19
        # Agent-level event
20
        string = f"1:[{agent['id']}] ({agent['name']}) {agent.get('ip', 'any')}->ollama:{json.dumps(msg)}"
21

22
    # Create and connect socket
23
    sock = socket(AF_UNIX, SOCK_DGRAM)
24
    sock.connect(socket_addr)
25

26
    # Send event
27
    sock.send(string.encode())
28

29
    # Close socket
30
    sock.close()

1
1:source_name:{"key": "value"}

1
1:[agent_id] (agent_name) agent_ip->source_name:{"key": "value"}

1
#!/var/ossec/framework/python/bin/python3
2
"""
3
Wazuh-Ollama Integration Script
4
Author: Anubhav Gain
5
Description: Enriches Wazuh security events with AI-powered analysis from Ollama
6
"""
7

8
import json
9
import sys
10
import time
11
import os
12
from socket import socket, AF_UNIX, SOCK_DGRAM
13
from ollama import Client
14

15
# ============================================================================
16
# GLOBAL CONFIGURATION
17
# ============================================================================
18

19
# Enable/disable debug logging
20
debug_enabled = True
21

22
# Get Wazuh installation path
23
pwd = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
24

25
# Empty alert dictionary
26
alert = {}
27

28
# Current timestamp
29
now = time.strftime("%a %b %d %H:%M:%S %Z %Y")
30

31
# Log file for integration logging
32
log_file = f"{pwd}/logs/integrations.log"
33

34
# Wazuh Unix socket address
35
socket_addr = f"{pwd}/queue/sockets/queue"
36

37

38
# ============================================================================
39
# HELPER FUNCTIONS
40
# ============================================================================
41

42
def debug(msg):
43
    """
44
    Write debug messages to integration log file
45

46
    Args:
47
        msg: Message to log
48
    """
49
    if debug_enabled:
50
        msg = f"{now}: {msg}\n"
51

52
    with open(log_file, "a") as integration_logs:
53
        integration_logs.write(str(msg))
54

55

56
def send_event(msg, agent=None):
57
    """
58
    Send custom event to Wazuh via Unix socket
59

60
    Args:
61
        msg: Event data dictionary
62
        agent: Agent information (optional)
63
    """
64
    if not agent or agent["id"] == "000":
65
        # Manager-level event
66
        string = f"1:ollama:{json.dumps(msg)}"
67
    else:
68
        # Agent-level event with full context
69
        string = "1:[{0}] ({1}) {2}->ollama:{3}".format(
70
            agent["id"],
71
            agent["name"],
72
            agent.get("ip", "any"),
73
            json.dumps(msg)
74
        )
75

76
    debug(f"Sending event: {string}")
77

78
    # Create Unix socket
79
    sock = socket(AF_UNIX, SOCK_DGRAM)
80
    sock.connect(socket_addr)
81
    sock.send(string.encode())
82
    sock.close()
83

84

85
# ============================================================================
86
# OLLAMA API INTEGRATION
87
# ============================================================================
88

89
def query_ollama_api(src_ip, ollama_host='http://ollama:11434',
90
                     ollama_model='llama3.2'):
91
    """
92
    Query Ollama API for IP address intelligence
93

94
    Args:
95
        src_ip: Source IP address to analyze
96
        ollama_host: Ollama server URL
97
        ollama_model: Model to use for analysis
98

99
    Returns:
100
        str: AI-generated analysis of the IP address
101
    """
102
    try:
103
        # Create Ollama client
104
        client = Client(host=ollama_host)
105

106
        # Prepare prompt for IP analysis
107
        prompt = f"""Analyze this IP address from a cybersecurity perspective: {src_ip}
108

109
Provide:
110
1. Geolocation information if known
111
2. Known threat intelligence (if this is a known malicious IP)
112
3. Typical attack patterns from this region
113
4. Recommended security actions
114

115
Keep the response concise (2-3 sentences)."""
116

117
        # Send request to Ollama
118
        response = client.chat(
119
            model=ollama_model,
120
            messages=[{
121
                'role': 'user',
122
                'content': prompt
123
            }]
124
        )
125

126
        # Extract response content
127
        return response.model_dump().get('message').get('content')
128

129
    except Exception as e:
130
        debug(f"Error querying Ollama: {str(e)}")
131
        return f"Error analyzing IP address: {str(e)}"
132

133

134
# ============================================================================
135
# EVENT PROCESSING
136
# ============================================================================
137

138
def get_ollama_info(alert):
139
    """
140
    Process alert and create enriched event with Ollama analysis
141

142
    Args:
143
        alert: Original Wazuh alert dictionary
144

145
    Returns:
146
        dict: Enriched alert with Ollama analysis, or 0 if no source IP
147
    """
148
    alert_output = {}
149

150
    # Exit if alert doesn't contain source IP
151
    if "srcip" not in alert.get("data", {}):
152
        debug("No source IP found in alert, skipping Ollama analysis")
153
        return 0
154

155
    src_ip = alert["data"]["srcip"]
156
    debug(f"Analyzing source IP: {src_ip}")
157

158
    # Query Ollama for IP intelligence
159
    ollama_analysis = query_ollama_api(src_ip)
160

161
    # Build enriched alert structure
162
    alert_output["ollama"] = {
163
        "found": 1,
164
        "source": {
165
            "alert_id": alert.get("id", "unknown"),
166
            "rule": alert.get("rule", {}).get("id", "unknown"),
167
            "description": alert.get("rule", {}).get("description", "unknown"),
168
            "full_log": alert.get("full_log", ""),
169
            "srcip": src_ip
170
        },
171
        "info": ollama_analysis,
172
        "srcip": src_ip,
173
        "timestamp": alert.get("timestamp", now)
174
    }
175

176
    alert_output["integration"] = "custom-ollama"
177

178
    debug(f"Created enriched event: {json.dumps(alert_output, indent=2)}")
179

180
    return alert_output
181

182

183
# ============================================================================
184
# MAIN EXECUTION
185
# ============================================================================
186

187
if __name__ == "__main__":
188
    try:
189
        debug("=" * 80)
190
        debug("Starting Ollama integration")
191

192
        # Get alert file location from command line
193
        alert_file_location = sys.argv[1]
194
        debug(f"Alert file location: {alert_file_location}")
195

196
        # Get Ollama domain from command line (optional)
197
        ollama_domain = sys.argv[3] if len(sys.argv) > 3 else "http://ollama:11434"
198
        debug(f"Ollama domain: {ollama_domain}")
199

200
        # Load alert JSON file
201
        with open(alert_file_location) as alert_file:
202
            alert = json.load(alert_file)
203

204
        debug(f"Processing alert: {json.dumps(alert, indent=2)}")
205

206
        # Get Ollama analysis
207
        enriched_event = get_ollama_info(alert)
208

209
        # Send enriched event to Wazuh if analysis was successful
210
        if enriched_event:
211
            send_event(enriched_event, alert.get("agent"))
212
            debug("Successfully sent enriched event to Wazuh")
213
        else:
214
            debug("No enriched event to send")
215

216
        debug("Ollama integration completed successfully")
217
        debug("=" * 80)
218

219
    except Exception as e:
220
        debug(f"ERROR in main execution: {str(e)}")
221
        import traceback
222
        debug(traceback.format_exc())
223
        sys.exit(1)

1
-rwxr-x--- 1 root wazuh 8192 Feb 24 10:00 custom-integration-ollama.py

1
<!-- Ollama Integration Configuration -->
2
<integration>
3
  <name>custom-integration-ollama.py</name>
4
  <hook_url>http://ollama:11434</hook_url>
5
  <level>10</level>
6
  <rule_id>100004,100005</rule_id>
7
  <alert_format>json</alert_format>
8
</integration>

1
<!-- SSH Authentication Rules for Public IPs -->
2
<group name="local,syslog,sshd,">
3

4
  <!-- Rule 100004: SSH authentication failed from public IP -->
5
  <rule id="100004" level="10">
6
    <if_sid>5760</if_sid>
7
    <match type="pcre2">\b(?!(10)|192\.168|172\.(2[0-9]|1[6-9]|3[0-1])|(25[6-9]|2[6-9][0-9]|[3-9][0-9][0-9]|99[1-9]))[0-9]{1,3}\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)</match>
8
    <description>sshd: Authentication failed from public IP address $(srcip)</description>
9
    <group>authentication_failed,pci_dss_10.2.4,pci_dss_10.2.5,</group>
10
  </rule>
11

12
  <!-- Rule 100005: SSH invalid user from public IP -->
13
  <rule id="100005" level="10">
14
    <if_sid>5710</if_sid>
15
    <match type="pcre2">\b(?!(10)|192\.168|172\.(2[0-9]|1[6-9]|3[0-1])|(25[6-9]|2[6-9][0-9]|[3-9][0-9][0-9]|99[1-9]))[0-9]{1,3}\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)</match>
16
    <description>sshd: Invalid user login attempt from public IP address $(srcip)</description>
17
    <group>authentication_failed,pci_dss_10.2.4,pci_dss_10.2.5,</group>
18
  </rule>
19

20
</group>
21

22
<!-- Ollama Enrichment Alert Rule -->
23
<group name="local,ollama,">
24

25
  <!-- Rule 100007: Display Ollama analysis results -->
26
  <rule id="100007" level="10">
27
    <field name="ollama.srcip">\.+</field>
28
    <description>[OLLAMA AI ANALYSIS] IP address $(ollama.srcip) attempting network access - AI threat intelligence added</description>
29
    <group>ollama_analysis,threat_intelligence,</group>
30
  </rule>
31

32
</group>

1
\b(?!(10)|192\.168|172\.(2[0-9]|1[6-9]|3[0-1])|(25[6-9]|2[6-9][0-9]|[3-9][0-9][0-9]|99[1-9]))[0-9]{1,3}\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)

1
def analyze_file_change(alert):
2
    """Analyze FIM alerts for potential threats"""
3

4
    if "syscheck" not in alert:
5
        return 0
6

7
    file_path = alert["syscheck"]["path"]
8
    file_change = alert["syscheck"]["event"]
9

10
    prompt = f"""Analyze this file integrity change:
11
File: {file_path}
12
Change: {file_change}
13

14
Is this suspicious? What could it indicate?
15
Provide risk assessment and recommended actions."""
16

17
    client = Client(host='http://ollama:11434')
18
    response = client.chat(
19
        model='llama3.2',
20
        messages=[{'role': 'user', 'content': prompt}]
21
    )
22

23
    return response.model_dump().get('message').get('content')

1
def analyze_vulnerability(alert):
2
    """Enhance vulnerability data with remediation advice"""
3

4
    if "vulnerability" not in alert["data"]:
5
        return 0
6

7
    cve_id = alert["data"]["vulnerability"]["cve"]
8
    severity = alert["data"]["vulnerability"]["severity"]
9

10
    prompt = f"""Provide remediation guidance for:
11
CVE: {cve_id}
12
Severity: {severity}
13

14
Include:
15
1. Exploit likelihood
16
2. Patch priority
17
3. Temporary mitigations
18
4. Related vulnerabilities"""
19

20
    # ... query Ollama and return

1
def analyze_log_anomaly(alert):
2
    """Use LLM to detect anomalous log patterns"""
3

4
    log_content = alert["full_log"]
5

6
    prompt = f"""Analyze this log entry for anomalies:
7
{log_content}
8

9
Identify:
10
1. Unusual patterns
11
2. Potential security implications
12
3. False positive likelihood
13
4. Investigation steps"""
14

15
    # ... query Ollama and return

1
import hashlib
2
import pickle
3
from pathlib import Path
4

5
CACHE_DIR = "/var/ossec/tmp/ollama_cache"
6

7
def get_cached_response(src_ip):
8
    """Check if we have a recent cached response"""
9

10
    cache_file = Path(CACHE_DIR) / f"{hashlib.md5(src_ip.encode()).hexdigest()}.pkl"
11

12
    if cache_file.exists():
13
        # Check if cache is less than 24 hours old
14
        if time.time() - cache_file.stat().st_mtime < 86400:
15
            with open(cache_file, 'rb') as f:
16
                return pickle.load(f)
17

18
    return None
19

20
def cache_response(src_ip, response):
21
    """Cache Ollama response for future use"""
22

23
    Path(CACHE_DIR).mkdir(exist_ok=True)
24
    cache_file = Path(CACHE_DIR) / f"{hashlib.md5(src_ip.encode()).hexdigest()}.pkl"
25

26
    with open(cache_file, 'wb') as f:
27
        pickle.dump(response, f)
28

29
def query_ollama_api(src_ip, ollama_host='http://ollama:11434',
30
                     ollama_model='llama3.2'):
31
    """Query with caching"""
32

33
    # Check cache first
34
    cached = get_cached_response(src_ip)
35
    if cached:
36
        debug(f"Using cached response for {src_ip}")
37
        return cached
38

39
    # Query Ollama
40
    client = Client(host=ollama_host)
41
    response = client.chat(
42
        model=ollama_model,
43
        messages=[{
44
            'role': 'user',
45
            'content': f'Analyze IP: {src_ip}'
46
        }]
47
    )
48

49
    result = response.model_dump().get('message').get('content')
50

51
    # Cache the response
52
    cache_response(src_ip, result)
53

54
    return result

1
import threading
2
from collections import deque
3
from time import time, sleep
4

5
class RateLimiter:
6
    """Simple rate limiter for Ollama queries"""
7

8
    def __init__(self, max_requests=10, time_window=60):
9
        self.max_requests = max_requests
10
        self.time_window = time_window
11
        self.requests = deque()
12
        self.lock = threading.Lock()
13

14
    def acquire(self):
15
        """Wait if necessary to maintain rate limit"""
16
        with self.lock:
17
            now = time()
18

19
            # Remove old requests outside time window
20
            while self.requests and self.requests[0] < now - self.time_window:
21
                self.requests.popleft()
22

23
            # If at limit, wait
24
            if len(self.requests) >= self.max_requests:
25
                sleep_time = self.time_window - (now - self.requests[0])
26
                if sleep_time > 0:
27
                    debug(f"Rate limit reached, sleeping {sleep_time}s")
28
                    sleep(sleep_time)
29

30
            self.requests.append(now)
31

32
# Global rate limiter
33
ollama_limiter = RateLimiter(max_requests=10, time_window=60)
34

35
def query_ollama_api(src_ip, ollama_host='http://ollama:11434',
36
                     ollama_model='llama3.2'):
37
    """Query with rate limiting"""
38

39
    # Wait for rate limiter
40
    ollama_limiter.acquire()
41

42
    # Proceed with query
43
    client = Client(host=ollama_host)
44
    # ... rest of function