Complete Guide to Wazuh LLM and RAG: AI-Powered Security Analysis and Documentation#

1
import argparse
2
import os
3

4
from langchain_community.document_loaders import UnstructuredPDFLoader
5
from langchain_ollama import OllamaEmbeddings
6
from langchain_text_splitters import RecursiveCharacterTextSplitter
7
from langchain_community.vectorstores import Chroma
8

9

10
def upload(document_path, model_name, collection_name, ollama_base_url):
11
    """
12
    Load PDF documentation and create vector embeddings
13

14
    Args:
15
        document_path: Path to PDF file
16
        model_name: Ollama embedding model name
17
        collection_name: ChromaDB collection name
18
        ollama_base_url: Ollama server URL
19

20
    Returns:
21
        Chroma vector store object
22
    """
23
    # Get current script path
24
    current_path = os.path.dirname(os.path.realpath(__file__))
25
    chroma_persistent_directory = current_path + "/data"
26

27
    # Create data directory if it doesn't exist
28
    if not os.path.exists(chroma_persistent_directory):
29
        os.makedirs(chroma_persistent_directory, exist_ok=True)
30

31
    # Load PDF document
32
    loader = UnstructuredPDFLoader(file_path=document_path)
33
    data = loader.load()
34

35
    # Split text into chunks
36
    # chunk_size: 7500 characters per chunk
37
    # chunk_overlap: 100 characters overlap between chunks
38
    text_splitter = RecursiveCharacterTextSplitter(
39
        chunk_size=7500,
40
        chunk_overlap=100
41
    )
42
    chunks = text_splitter.split_documents(data)
43

44
    # Create vector embeddings and store in ChromaDB
45
    vector = Chroma.from_documents(
46
        documents=chunks,
47
        embedding=OllamaEmbeddings(
48
            base_url=ollama_base_url,
49
            model=model_name,
50
            show_progress=True
51
        ),
52
        collection_name=collection_name,
53
        persist_directory=chroma_persistent_directory
54
    )
55

56
    return vector
57

58

59
if __name__ == '__main__':
60
    # Creating argument parser
61
    parser = argparse.ArgumentParser(
62
        description='Upload a PDF to the vector store'
63
    )
64

65
    # Adding arguments
66
    parser.add_argument(
67
        '-p', '--path',
68
        type=str,
69
        help='Path to the PDF file',
70
        required=True
71
    )
72
    parser.add_argument(
73
        '-m', '--model',
74
        type=str,
75
        help='Name of the Ollama model for embedding',
76
        default='nomic-embed-text'
77
    )
78
    parser.add_argument(
79
        '-n', '--name',
80
        type=str,
81
        help='Collection name in ChromaDB',
82
        default='wazuh'
83
    )
84
    parser.add_argument(
85
        '-b', '--base-url',
86
        type=str,
87
        help='Base URL for the Ollama server',
88
        default='http://127.0.0.1:11434'
89
    )
90

91
    # Parsing arguments
92
    args = parser.parse_args()
93

94
    # Calling the upload function
95
    upload(
96
        document_path=args.path,
97
        model_name=args.model,
98
        collection_name=args.name,
99
        ollama_base_url=args.base_url
100
    )
101

102
    print(f"Successfully uploaded {args.path} to ChromaDB collection '{args.name}'")

1
import argparse
2
import os
3

4
from langchain.retrievers import MultiQueryRetriever
5
import chromadb
6
from langchain_core.output_parsers import StrOutputParser
7
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
8
from langchain_core.runnables import RunnablePassthrough
9
from langchain_ollama import ChatOllama
10
from langchain_ollama import OllamaEmbeddings
11
from langchain_chroma import Chroma
12

13

14
def ask_ollama(question, collection_name='wazuh',
15
               embedding_model='nomic-embed-text',
16
               local_model='llama3.2'):
17
    """
18
    Query the RAG system with a question
19

20
    Args:
21
        question: User question
22
        collection_name: ChromaDB collection name
23
        embedding_model: Embedding model name
24
        local_model: LLM model for generation
25

26
    Returns:
27
        Generated answer based on documentation
28
    """
29
    # Get current script path
30
    current_path = os.path.dirname(os.path.realpath(__file__))
31
    chroma_persistent_directory = current_path + "/data"
32

33
    # Initialize embedding model
34
    embedding = OllamaEmbeddings(model=embedding_model)
35

36
    # Connect to persistent ChromaDB
37
    persistent_client = chromadb.PersistentClient(
38
        path=chroma_persistent_directory
39
    )
40

41
    # Load vector database
42
    vector_db = Chroma(
43
        client=persistent_client,
44
        collection_name=collection_name,
45
        embedding_function=embedding,
46
    )
47

48
    # Initialize LLM
49
    load_ollama = ChatOllama(model=local_model)
50

51
    # Multi-query prompt template
52
    # This generates multiple versions of the question for better retrieval
53
    prompt_template = PromptTemplate(
54
        input_variables=["question"],
55
        template="""You are an AI language model assistant. Your task is to generate 2
56
different versions of the given user question to retrieve relevant documents from
57
a vector database. By generating multiple perspectives on the user question, your
58
goal is to help the user overcome some of the limitations of the distance-based
59
similarity search. Provide these alternative questions separated by newlines.
60
Original question: {question}""",
61
    )
62

63
    # Answer generation template
64
    template = """Answer the question based ONLY on the following context:
65
{context}
66
Question: {question}
67
"""
68

69
    # Create multi-query retriever
70
    retriever = MultiQueryRetriever.from_llm(
71
        vector_db.as_retriever(),
72
        load_ollama,
73
        prompt=prompt_template
74
    )
75

76
    # Create chat prompt
77
    prompt = ChatPromptTemplate.from_template(template)
78

79
    # Build RAG chain
80
    chain = (
81
        {"context": retriever, "question": RunnablePassthrough()}
82
        | prompt
83
        | load_ollama
84
        | StrOutputParser()
85
    )
86

87
    return chain.invoke(question)
88

89

90
if __name__ == '__main__':
91
    parser = argparse.ArgumentParser(description='Ask Ollama a question')
92
    parser.add_argument(
93
        '-q', '--question',
94
        type=str,
95
        help='The question to ask Ollama',
96
        required=True
97
    )
98
    parser.add_argument(
99
        '-c', '--collection',
100
        type=str,
101
        help='ChromaDB collection name',
102
        default='wazuh'
103
    )
104
    parser.add_argument(
105
        '-e', '--embedding',
106
        type=str,
107
        help='Embedding model name',
108
        default='nomic-embed-text'
109
    )
110
    parser.add_argument(
111
        '-m', '--model',
112
        type=str,
113
        help='LLM model name',
114
        default='llama3.2'
115
    )
116

117
    args = parser.parse_args()
118

119
    print(ask_ollama(
120
        args.question,
121
        collection_name=args.collection,
122
        embedding_model=args.embedding,
123
        local_model=args.model
124
    ))

1
import requests
2
import json
3
from ask import ask_ollama
4

5
def comprehensive_security_analysis(security_event):
6
    """
7
    Perform comprehensive security analysis using both LLM and RAG
8

9
    Args:
10
        security_event: Wazuh security event (dict or string)
11

12
    Returns:
13
        Complete analysis with event assessment and documentation references
14
    """
15
    # Step 1: Analyze the event with fine-tuned LLM
16
    event_prompt = f"Analyze this Wazuh event: {json.dumps(security_event)}"
17

18
    llm_analysis = requests.post(
19
        'http://localhost:11434/api/generate',
20
        json={
21
            'model': 'mranv/siem-llama-3.1:v1',
22
            'prompt': event_prompt,
23
            'stream': False
24
        }
25
    ).json()
26

27
    # Step 2: Extract key terms for documentation lookup
28
    # (In production, use NLP to extract relevant terms)
29
    event_type = "brute force"  # Example extracted term
30

31
    # Step 3: Query RAG for relevant documentation
32
    doc_query = f"How to detect and respond to {event_type} attacks in Wazuh?"
33
    rag_response = ask_ollama(doc_query)
34

35
    # Step 4: Combine results
36
    return {
37
        "event_analysis": llm_analysis['response'],
38
        "documentation": rag_response,
39
        "timestamp": security_event.get('timestamp', 'N/A')
40
    }
41

42

43
# Example usage
44
event = {
45
    "timestamp": "2025-03-02T10:00:00Z",
46
    "rule": {
47
        "id": "5710",
48
        "level": 10,
49
        "description": "Brute force attack detected"
50
    },
51
    "data": {
52
        "srcip": "203.0.113.5",
53
        "attempts": 150
54
    }
55
}
56

57
result = comprehensive_security_analysis(event)
58
print("=== EVENT ANALYSIS ===")
59
print(result["event_analysis"])
60
print("\n=== RELEVANT DOCUMENTATION ===")
61
print(result["documentation"])

1
from flask import Flask, request, jsonify, render_template
2
import requests
3
from ask import ask_ollama
4

5
app = Flask(__name__)
6

7
@app.route('/')
8
def index():
9
    return render_template('dashboard.html')
10

11
@app.route('/analyze', methods=['POST'])
12
def analyze_event():
13
    """Analyze security event endpoint"""
14
    event_data = request.json
15

16
    # Analyze with LLM
17
    llm_response = requests.post(
18
        'http://localhost:11434/api/generate',
19
        json={
20
            'model': 'mranv/siem-llama-3.1:v1',
21
            'prompt': f"Analyze: {json.dumps(event_data)}",
22
            'stream': False
23
        }
24
    ).json()
25

26
    return jsonify({
27
        'analysis': llm_response['response'],
28
        'event': event_data
29
    })
30

31
@app.route('/docs', methods=['POST'])
32
def query_docs():
33
    """Query documentation endpoint"""
34
    question = request.json.get('question')
35
    answer = ask_ollama(question)
36

37
    return jsonify({
38
        'question': question,
39
        'answer': answer
40
    })
41

42
if __name__ == '__main__':
43
    app.run(debug=True, port=5000)

1
# Adjust chunk size for better performance
2
text_splitter = RecursiveCharacterTextSplitter(
3
    chunk_size=5000,  # Smaller chunks for faster retrieval
4
    chunk_overlap=200  # More overlap for better context
5
)

1
# Adjust chunk size
2
text_splitter = RecursiveCharacterTextSplitter(
3
    chunk_size=10000,  # Larger chunks
4
    chunk_overlap=500   # More overlap
5
)