Complete LangChain Integration with Wazuh MCP Server: Building AI-Powered Security Agents#

1
import asyncio
2
import aiohttp
3
from typing import Dict, List, Optional, Any
4
from langchain.tools import BaseTool
5
from langchain.pydantic_v1 import BaseModel, Field
6
import json
7

8
class MCPClient:
9
    """Enhanced MCP Client for LangChain integration"""
10

11
    def __init__(self, base_url: str = "http://localhost:8000"):
12
        self.base_url = base_url
13
        self.session = None
14
        self.jwt_token = None
15

16
    async def __aenter__(self):
17
        self.session = aiohttp.ClientSession()
18
        await self.authenticate()
19
        return self
20

21
    async def __aexit__(self, exc_type, exc_val, exc_tb):
22
        if self.session:
23
            await self.session.close()
24

25
    async def authenticate(self):
26
        """Authenticate with Wazuh MCP Server"""
27
        async with self.session.post(
28
            f"{self.base_url}/auth",
29
            json={"username": "admin", "password": "admin"}
30
        ) as resp:
31
            data = await resp.json()
32
            self.jwt_token = data['token']
33

34
    async def call_tool(self, tool_name: str, params: Dict) -> Dict:
35
        """Call an MCP tool"""
36
        headers = {"Authorization": f"Bearer {self.jwt_token}"}
37

38
        async with self.session.post(
39
            f"{self.base_url}/tools/{tool_name}",
40
            json=params,
41
            headers=headers
42
        ) as resp:
43
            return await resp.json()
44

45
class MCPToolWrapper(BaseTool):
46
    """Wrapper to convert MCP tools to LangChain tools"""
47

48
    name: str = "mcp_tool"
49
    description: str = "MCP tool for Wazuh operations"
50
    mcp_client: Optional[MCPClient] = None
51
    tool_name: str = ""
52

53
    class Config:
54
        arbitrary_types_allowed = True
55

56
    async def _arun(self, query: str) -> str:
57
        """Async execution of MCP tool"""
58
        params = json.loads(query) if isinstance(query, str) else query
59
        result = await self.mcp_client.call_tool(self.tool_name, params)
60
        return json.dumps(result, indent=2)
61

62
    def _run(self, query: str) -> str:
63
        """Sync wrapper for async execution"""
64
        return asyncio.run(self._arun(query))

1
from typing import List
2
from langchain.tools import Tool
3

4
class WazuhToolRegistry:
5
    """Registry of Wazuh MCP tools for LangChain"""
6

7
    @staticmethod
8
    def get_tools(mcp_client: MCPClient) -> List[Tool]:
9
        """Get all available Wazuh tools"""
10

11
        tools = []
12

13
        # Agent Management Tools
14
        tools.append(Tool(
15
            name="get_agents",
16
            description="Get list of all Wazuh agents with their status",
17
            func=lambda x: asyncio.run(
18
                mcp_client.call_tool("GetAgentsTool", json.loads(x))
19
            )
20
        ))
21

22
        # Process Analysis Tools
23
        tools.append(Tool(
24
            name="get_agent_processes",
25
            description="Get running processes on a specific agent",
26
            func=lambda x: asyncio.run(
27
                mcp_client.call_tool("GetAgentProcessesTool", json.loads(x))
28
            )
29
        ))
30

31
        # Network Analysis Tools
32
        tools.append(Tool(
33
            name="get_agent_ports",
34
            description="Get open network ports on a specific agent",
35
            func=lambda x: asyncio.run(
36
                mcp_client.call_tool("GetAgentPortsTool", json.loads(x))
37
            )
38
        ))
39

40
        # Package Management Tools
41
        tools.append(Tool(
42
            name="get_agent_packages",
43
            description="Get installed packages on a specific agent",
44
            func=lambda x: asyncio.run(
45
                mcp_client.call_tool("GetAgentPackagesTool", json.loads(x))
46
            )
47
        ))
48

49
        # Security Assessment Tools
50
        tools.append(Tool(
51
            name="get_agent_sca",
52
            description="Get security configuration assessment for an agent",
53
            func=lambda x: asyncio.run(
54
                mcp_client.call_tool("GetAgentSCATool", json.loads(x))
55
            )
56
        ))
57

58
        # Rule Management Tools
59
        tools.append(Tool(
60
            name="list_rules",
61
            description="List Wazuh detection rules",
62
            func=lambda x: asyncio.run(
63
                mcp_client.call_tool("ListRulesTool", json.loads(x))
64
            )
65
        ))
66

67
        return tools

1
from langchain_openai import ChatOpenAI
2
from langchain.agents import create_openai_tools_agent, AgentExecutor
3
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
4
from langchain.memory import ConversationBufferWindowMemory
5
import os
6
from dotenv import load_dotenv
7

8
load_dotenv()
9

10
class OpenAISecurityAgent:
11
    """OpenAI-powered security analysis agent"""
12

13
    def __init__(self, mcp_client: MCPClient):
14
        self.mcp_client = mcp_client
15
        self.llm = ChatOpenAI(
16
            model="gpt-4-turbo-preview",
17
            temperature=0,
18
            api_key=os.getenv("OPENAI_API_KEY")
19
        )
20
        self.tools = WazuhToolRegistry.get_tools(mcp_client)
21
        self.memory = ConversationBufferWindowMemory(
22
            k=10,
23
            return_messages=True,
24
            memory_key="chat_history"
25
        )
26
        self.agent = self._create_agent()
27

28
    def _create_agent(self):
29
        """Create the OpenAI agent with Wazuh tools"""
30

31
        system_prompt = """You are a senior security analyst with expertise in:
32
        - Incident response and threat hunting
33
        - SIEM analysis using Wazuh
34
        - Network and endpoint security
35
        - Malware analysis and forensics
36

37
        You have access to Wazuh MCP tools to query security data.
38
        Always provide detailed analysis with evidence.
39
        Prioritize findings by severity and business impact.
40
        Suggest remediation steps for identified issues.
41
        """
42

43
        prompt = ChatPromptTemplate.from_messages([
44
            ("system", system_prompt),
45
            MessagesPlaceholder(variable_name="chat_history"),
46
            ("human", "{input}"),
47
            MessagesPlaceholder(variable_name="agent_scratchpad")
48
        ])
49

50
        agent = create_openai_tools_agent(
51
            llm=self.llm,
52
            tools=self.tools,
53
            prompt=prompt
54
        )
55

56
        return AgentExecutor(
57
            agent=agent,
58
            tools=self.tools,
59
            memory=self.memory,
60
            verbose=True,
61
            max_iterations=10,
62
            handle_parsing_errors=True
63
        )
64

65
    async def analyze(self, query: str) -> str:
66
        """Analyze security query using OpenAI"""
67
        result = await self.agent.ainvoke({"input": query})
68
        return result["output"]
69

70
    async def investigate_alert(self, alert: Dict) -> Dict:
71
        """Investigate security alert with context"""
72

73
        investigation_prompt = f"""
74
        Investigate this security alert:
75

76
        Alert Details:
77
        - Rule: {alert.get('rule', {}).get('description')}
78
        - Agent: {alert.get('agent', {}).get('name')} ({alert.get('agent', {}).get('id')})
79
        - Severity: {alert.get('rule', {}).get('level')}
80
        - Time: {alert.get('timestamp')}
81

82
        Perform the following:
83
        1. Check the agent's current status
84
        2. List suspicious processes on the agent
85
        3. Check network connections
86
        4. Assess security configuration
87
        5. Provide risk assessment
88
        6. Recommend immediate actions
89
        """
90

91
        return await self.analyze(investigation_prompt)

1
from langchain.chains import LLMChain, SequentialChain
2
from langchain.prompts import PromptTemplate
3
from langchain.output_parsers import PydanticOutputParser
4
from pydantic import BaseModel, Field
5
from typing import List
6

7
class ThreatAssessment(BaseModel):
8
    """Structured threat assessment output"""
9
    severity: int = Field(description="Severity score 1-10")
10
    threat_type: str = Field(description="Type of threat detected")
11
    affected_assets: List[str] = Field(description="List of affected assets")
12
    iocs: List[str] = Field(description="Indicators of compromise")
13
    remediation_steps: List[str] = Field(description="Recommended actions")
14
    confidence: float = Field(description="Confidence score 0-1")
15

16
class AdvancedOpenAIChains:
17
    """Advanced analysis chains using OpenAI"""
18

19
    def __init__(self, llm: ChatOpenAI, mcp_client: MCPClient):
20
        self.llm = llm
21
        self.mcp_client = mcp_client
22

23
    def create_threat_assessment_chain(self):
24
        """Chain for structured threat assessment"""
25

26
        parser = PydanticOutputParser(pydantic_object=ThreatAssessment)
27

28
        # First chain: Data collection
29
        collection_prompt = PromptTemplate(
30
            input_variables=["agent_id"],
31
            template="""
32
            Collect security data for agent {agent_id}:
33
            1. Get agent status and information
34
            2. List running processes
35
            3. Check network connections
36
            4. Review installed packages
37
            5. Assess security configuration
38

39
            Compile all findings.
40
            """
41
        )
42

43
        collection_chain = LLMChain(
44
            llm=self.llm,
45
            prompt=collection_prompt,
46
            output_key="raw_data"
47
        )
48

49
        # Second chain: Threat analysis
50
        analysis_prompt = PromptTemplate(
51
            input_variables=["raw_data"],
52
            template="""
53
            Analyze the following security data for threats:
54
            {raw_data}
55

56
            {format_instructions}
57

58
            Provide a comprehensive threat assessment.
59
            """,
60
            partial_variables={
61
                "format_instructions": parser.get_format_instructions()
62
            }
63
        )
64

65
        analysis_chain = LLMChain(
66
            llm=self.llm,
67
            prompt=analysis_prompt,
68
            output_key="threat_assessment"
69
        )
70

71
        # Sequential chain
72
        return SequentialChain(
73
            chains=[collection_chain, analysis_chain],
74
            input_variables=["agent_id"],
75
            output_variables=["threat_assessment"],
76
            verbose=True
77
        )
78

79
    def create_incident_response_chain(self):
80
        """Chain for automated incident response"""
81

82
        # Detection chain
83
        detection_prompt = PromptTemplate(
84
            input_variables=["indicators"],
85
            template="""
86
            Detect security incidents based on these indicators:
87
            {indicators}
88

89
            Search for:
90
            1. Matching patterns across agents
91
            2. Related suspicious activities
92
            3. Potential attack vectors
93
            4. Affected systems
94
            """
95
        )
96

97
        # Investigation chain
98
        investigation_prompt = PromptTemplate(
99
            input_variables=["detection_results"],
100
            template="""
101
            Investigate detected incidents:
102
            {detection_results}
103

104
            Perform:
105
            1. Root cause analysis
106
            2. Timeline reconstruction
107
            3. Impact assessment
108
            4. Attribution analysis
109
            """
110
        )
111

112
        # Response chain
113
        response_prompt = PromptTemplate(
114
            input_variables=["investigation_results"],
115
            template="""
116
            Generate incident response plan:
117
            {investigation_results}
118

119
            Include:
120
            1. Containment strategies
121
            2. Eradication steps
122
            3. Recovery procedures
123
            4. Lessons learned
124
            """
125
        )
126

127
        detection_chain = LLMChain(
128
            llm=self.llm,
129
            prompt=detection_prompt,
130
            output_key="detection_results"
131
        )
132

133
        investigation_chain = LLMChain(
134
            llm=self.llm,
135
            prompt=investigation_prompt,
136
            output_key="investigation_results"
137
        )
138

139
        response_chain = LLMChain(
140
            llm=self.llm,
141
            prompt=response_prompt,
142
            output_key="response_plan"
143
        )
144

145
        return SequentialChain(
146
            chains=[detection_chain, investigation_chain, response_chain],
147
            input_variables=["indicators"],
148
            output_variables=["response_plan"],
149
            verbose=True
150
        )

1
from langchain_anthropic import ChatAnthropic
2
from langchain.agents import create_structured_chat_agent
3
import os
4

5
class ClaudeSecurityAgent:
6
    """Claude-powered security analysis agent"""
7

8
    def __init__(self, mcp_client: MCPClient):
9
        self.mcp_client = mcp_client
10
        self.llm = ChatAnthropic(
11
            model="claude-3-opus-20240229",
12
            temperature=0,
13
            anthropic_api_key=os.getenv("ANTHROPIC_API_KEY"),
14
            max_tokens=4096
15
        )
16
        self.tools = WazuhToolRegistry.get_tools(mcp_client)
17
        self.agent = self._create_agent()
18

19
    def _create_agent(self):
20
        """Create Claude agent with enhanced reasoning"""
21

22
        system_prompt = """You are Claude, an advanced AI security analyst.
23

24
        Your capabilities include:
25
        - Deep security analysis with attention to detail
26
        - Pattern recognition across complex datasets
27
        - Hypothesis generation and testing
28
        - Comprehensive threat modeling
29

30
        Approach each analysis systematically:
31
        1. Gather comprehensive data
32
        2. Identify patterns and anomalies
33
        3. Form hypotheses about potential threats
34
        4. Test hypotheses with additional queries
35
        5. Provide evidence-based conclusions
36

37
        Use the Wazuh MCP tools to gather data and validate findings.
38
        Always explain your reasoning process.
39
        """
40

41
        prompt = ChatPromptTemplate.from_messages([
42
            ("system", system_prompt),
43
            ("human", "{input}"),
44
            MessagesPlaceholder(variable_name="agent_scratchpad")
45
        ])
46

47
        agent = create_structured_chat_agent(
48
            llm=self.llm,
49
            tools=self.tools,
50
            prompt=prompt
51
        )
52

53
        return AgentExecutor(
54
            agent=agent,
55
            tools=self.tools,
56
            verbose=True,
57
            max_iterations=15,  # Claude can handle more complex reasoning
58
            handle_parsing_errors=True,
59
            return_intermediate_steps=True
60
        )
61

62
    async def deep_analysis(self, context: Dict) -> Dict:
63
        """Perform deep security analysis with Claude"""
64

65
        analysis_prompt = f"""
66
        Perform a comprehensive security analysis:
67

68
        Context: {json.dumps(context, indent=2)}
69

70
        Your analysis should include:
71
        1. Threat landscape assessment
72
        2. Vulnerability identification
73
        3. Attack vector analysis
74
        4. Risk scoring with justification
75
        5. Detailed remediation plan
76
        6. Preventive measures
77

78
        Use multiple tools to validate findings.
79
        Provide confidence levels for each conclusion.
80
        """
81

82
        result = await self.agent.ainvoke({"input": analysis_prompt})
83

84
        return {
85
            "analysis": result["output"],
86
            "reasoning_steps": result.get("intermediate_steps", []),
87
            "tools_used": self._extract_tools_used(result)
88
        }
89

90
    def _extract_tools_used(self, result: Dict) -> List[str]:
91
        """Extract list of tools used in analysis"""
92
        tools_used = []
93
        for step in result.get("intermediate_steps", []):
94
            if hasattr(step[0], "tool"):
95
                tools_used.append(step[0].tool)
96
        return list(set(tools_used))

1
class ClaudeReasoningChains:
2
    """Advanced reasoning chains with Claude"""
3

4
    def __init__(self, claude_agent: ClaudeSecurityAgent):
5
        self.agent = claude_agent
6

7
    async def hypothesis_testing_analysis(self, initial_observation: str):
8
        """Hypothesis-driven security analysis"""
9

10
        # Generate hypotheses
11
        hypothesis_prompt = f"""
12
        Based on this observation: {initial_observation}
13

14
        Generate 5 possible security hypotheses that could explain this observation.
15
        For each hypothesis, list:
16
        1. The hypothesis statement
17
        2. What evidence would support it
18
        3. What evidence would refute it
19
        4. How to test it using Wazuh tools
20
        """
21

22
        hypotheses = await self.agent.agent.ainvoke({
23
            "input": hypothesis_prompt
24
        })
25

26
        # Test each hypothesis
27
        test_results = []
28
        for hypothesis in self._parse_hypotheses(hypotheses["output"]):
29
            test_prompt = f"""
30
            Test this security hypothesis:
31
            {hypothesis}
32

33
            Use Wazuh tools to:
34
            1. Gather supporting evidence
35
            2. Look for contradicting evidence
36
            3. Calculate probability of hypothesis being true
37
            """
38

39
            test_result = await self.agent.agent.ainvoke({
40
                "input": test_prompt
41
            })
42

43
            test_results.append({
44
                "hypothesis": hypothesis,
45
                "test_result": test_result["output"],
46
                "tools_used": self._extract_tools(test_result)
47
            })
48

49
        # Synthesize findings
50
        synthesis_prompt = f"""
51
        Synthesize the hypothesis testing results:
52
        {json.dumps(test_results, indent=2)}
53

54
        Provide:
55
        1. Most likely explanation
56
        2. Confidence level
57
        3. Supporting evidence
58
        4. Recommended actions
59
        5. Further investigation needed
60
        """
61

62
        synthesis = await self.agent.agent.ainvoke({
63
            "input": synthesis_prompt
64
        })
65

66
        return {
67
            "initial_observation": initial_observation,
68
            "hypotheses_tested": test_results,
69
            "conclusion": synthesis["output"]
70
        }
71

72
    async def threat_actor_attribution(self, indicators: List[str]):
73
        """Attribute activity to threat actors"""
74

75
        attribution_prompt = f"""
76
        Analyze these indicators for threat actor attribution:
77
        {json.dumps(indicators, indent=2)}
78

79
        Perform:
80
        1. TTP analysis and mapping to MITRE ATT&CK
81
        2. Infrastructure analysis
82
        3. Malware family identification
83
        4. Historical campaign correlation
84
        5. Attribution confidence assessment
85

86
        Use Wazuh tools to gather evidence from affected systems.
87
        Consider multiple attribution hypotheses.
88
        """
89

90
        attribution = await self.agent.agent.ainvoke({
91
            "input": attribution_prompt
92
        })
93

94
        return self._parse_attribution(attribution["output"])

1
from typing import List, Dict, Any
2
import asyncio
3

4
class SecurityAgentTeam:
5
    """Multi-agent collaborative security analysis"""
6

7
    def __init__(self, mcp_client: MCPClient):
8
        self.mcp_client = mcp_client
9

10
        # Initialize specialized agents
11
        self.agents = {
12
            "threat_hunter": self._create_threat_hunter(),
13
            "incident_responder": self._create_incident_responder(),
14
            "forensic_analyst": self._create_forensic_analyst(),
15
            "compliance_auditor": self._create_compliance_auditor(),
16
            "coordinator": self._create_coordinator()
17
        }
18

19
    def _create_threat_hunter(self):
20
        """Create specialized threat hunting agent"""
21
        llm = ChatOpenAI(model="gpt-4", temperature=0.2)
22

23
        prompt = """You are a specialized threat hunter.
24
        Focus on:
25
        - Identifying advanced persistent threats
26
        - Behavioral anomaly detection
27
        - Threat intelligence correlation
28
        - Proactive threat detection
29
        """
30

31
        return self._build_agent(llm, prompt, "threat_hunter")
32

33
    def _create_incident_responder(self):
34
        """Create specialized incident response agent"""
35
        llm = ChatOpenAI(model="gpt-4", temperature=0)
36

37
        prompt = """You are a specialized incident responder.
38
        Focus on:
39
        - Rapid incident triage
40
        - Containment strategies
41
        - Evidence preservation
42
        - Recovery procedures
43
        """
44

45
        return self._build_agent(llm, prompt, "incident_responder")
46

47
    def _create_forensic_analyst(self):
48
        """Create specialized forensic analysis agent"""
49
        llm = ChatAnthropic(model="claude-3-opus-20240229", temperature=0)
50

51
        prompt = """You are a specialized forensic analyst.
52
        Focus on:
53
        - Timeline reconstruction
54
        - Artifact analysis
55
        - Root cause analysis
56
        - Evidence chain preservation
57
        """
58

59
        return self._build_agent(llm, prompt, "forensic_analyst")
60

61
    def _create_compliance_auditor(self):
62
        """Create specialized compliance audit agent"""
63
        llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
64

65
        prompt = """You are a specialized compliance auditor.
66
        Focus on:
67
        - Regulatory compliance validation
68
        - Control effectiveness testing
69
        - Gap analysis
70
        - Audit evidence collection
71
        """
72

73
        return self._build_agent(llm, prompt, "compliance_auditor")
74

75
    def _create_coordinator(self):
76
        """Create coordinator agent to orchestrate team"""
77
        llm = ChatOpenAI(model="gpt-4", temperature=0.1)
78

79
        prompt = """You are the security team coordinator.
80
        Your role:
81
        - Assign tasks to specialized agents
82
        - Synthesize findings from multiple agents
83
        - Resolve conflicts in analysis
84
        - Provide unified recommendations
85

86
        Available agents:
87
        - threat_hunter: For APT and threat detection
88
        - incident_responder: For incident handling
89
        - forensic_analyst: For deep forensic analysis
90
        - compliance_auditor: For compliance checks
91
        """
92

93
        return self._build_agent(llm, prompt, "coordinator")
94

95
    def _build_agent(self, llm, prompt, role):
96
        """Build an agent with specific role"""
97
        tools = WazuhToolRegistry.get_tools(self.mcp_client)
98

99
        full_prompt = ChatPromptTemplate.from_messages([
100
            ("system", prompt),
101
            ("human", "{input}"),
102
            MessagesPlaceholder(variable_name="agent_scratchpad")
103
        ])
104

105
        agent = create_openai_tools_agent(
106
            llm=llm,
107
            tools=tools,
108
            prompt=full_prompt
109
        )
110

111
        return AgentExecutor(
112
            agent=agent,
113
            tools=tools,
114
            verbose=True,
115
            max_iterations=10
116
        )
117

118
    async def collaborative_analysis(self, scenario: str) -> Dict:
119
        """Perform collaborative multi-agent analysis"""
120

121
        # Coordinator assigns tasks
122
        coordination_prompt = f"""
123
        Analyze this scenario and assign tasks to appropriate agents:
124
        {scenario}
125

126
        For each agent, specify:
127
        1. What they should investigate
128
        2. Key questions to answer
129
        3. Tools they should use
130
        """
131

132
        coordination_plan = await self.agents["coordinator"].ainvoke({
133
            "input": coordination_prompt
134
        })
135

136
        # Execute parallel agent tasks
137
        tasks = []
138
        for agent_name, agent in self.agents.items():
139
            if agent_name != "coordinator":
140
                task_prompt = f"""
141
                Based on this scenario: {scenario}
142

143
                Perform your specialized analysis as {agent_name}.
144
                Focus on your area of expertise.
145
                """
146

147
                tasks.append(agent.ainvoke({"input": task_prompt}))
148

149
        # Gather results
150
        results = await asyncio.gather(*tasks)
151

152
        # Coordinator synthesizes findings
153
        synthesis_prompt = f"""
154
        Synthesize these findings from the team:
155

156
        {json.dumps([r["output"] for r in results], indent=2)}
157

158
        Provide:
159
        1. Unified threat assessment
160
        2. Prioritized findings
161
        3. Comprehensive action plan
162
        4. Risk score
163
        """
164

165
        final_analysis = await self.agents["coordinator"].ainvoke({
166
            "input": synthesis_prompt
167
        })
168

169
        return {
170
            "scenario": scenario,
171
            "coordination_plan": coordination_plan["output"],
172
            "agent_findings": {
173
                agent_name: results[i]["output"]
174
                for i, agent_name in enumerate(self.agents.keys())
175
                if agent_name != "coordinator"
176
            },
177
            "synthesis": final_analysis["output"]
178
        }

1
import asyncio
2
from asyncio import Queue
3
from dataclasses import dataclass
4
from enum import Enum
5

6
class MessageType(Enum):
7
    REQUEST = "request"
8
    RESPONSE = "response"
9
    BROADCAST = "broadcast"
10
    ALERT = "alert"
11

12
@dataclass
13
class AgentMessage:
14
    sender: str
15
    recipient: str
16
    message_type: MessageType
17
    content: Dict
18
    priority: int = 5
19

20
class AgentCommunicationHub:
21
    """Central communication hub for multi-agent system"""
22

23
    def __init__(self):
24
        self.message_queues = {}
25
        self.broadcast_queue = Queue()
26
        self.alert_handlers = []
27

28
    def register_agent(self, agent_name: str):
29
        """Register an agent with the communication hub"""
30
        self.message_queues[agent_name] = Queue()
31

32
    async def send_message(self, message: AgentMessage):
33
        """Send message between agents"""
34
        if message.message_type == MessageType.BROADCAST:
35
            await self.broadcast_queue.put(message)
36
            await self._handle_broadcast(message)
37
        elif message.message_type == MessageType.ALERT:
38
            await self._handle_alert(message)
39
        else:
40
            if message.recipient in self.message_queues:
41
                await self.message_queues[message.recipient].put(message)
42

43
    async def receive_message(self, agent_name: str) -> AgentMessage:
44
        """Receive message for specific agent"""
45
        if agent_name in self.message_queues:
46
            return await self.message_queues[agent_name].get()
47

48
    async def _handle_broadcast(self, message: AgentMessage):
49
        """Handle broadcast messages to all agents"""
50
        for agent_name, queue in self.message_queues.items():
51
            if agent_name != message.sender:
52
                await queue.put(message)
53

54
    async def _handle_alert(self, message: AgentMessage):
55
        """Handle high-priority alerts"""
56
        for handler in self.alert_handlers:
57
            await handler(message)
58

59
    def register_alert_handler(self, handler):
60
        """Register handler for alerts"""
61
        self.alert_handlers.append(handler)
62

63
class CollaborativeAgent:
64
    """Base class for collaborative agents"""
65

66
    def __init__(self, name: str, hub: AgentCommunicationHub):
67
        self.name = name
68
        self.hub = hub
69
        self.hub.register_agent(name)
70

71
    async def send_to(self, recipient: str, content: Dict):
72
        """Send message to another agent"""
73
        message = AgentMessage(
74
            sender=self.name,
75
            recipient=recipient,
76
            message_type=MessageType.REQUEST,
77
            content=content
78
        )
79
        await self.hub.send_message(message)
80

81
    async def broadcast(self, content: Dict):
82
        """Broadcast message to all agents"""
83
        message = AgentMessage(
84
            sender=self.name,
85
            recipient="all",
86
            message_type=MessageType.BROADCAST,
87
            content=content
88
        )
89
        await self.hub.send_message(message)
90

91
    async def alert(self, content: Dict):
92
        """Send high-priority alert"""
93
        message = AgentMessage(
94
            sender=self.name,
95
            recipient="all",
96
            message_type=MessageType.ALERT,
97
            content=content,
98
            priority=10
99
        )
100
        await self.hub.send_message(message)
101

102
    async def process_messages(self):
103
        """Process incoming messages"""
104
        while True:
105
            message = await self.hub.receive_message(self.name)
106
            await self.handle_message(message)
107

108
    async def handle_message(self, message: AgentMessage):
109
        """Override in subclass to handle messages"""
110
        pass

1
from langchain.embeddings import OpenAIEmbeddings
2
from langchain.vectorstores import Chroma
3
from langchain.text_splitter import RecursiveCharacterTextSplitter
4
from langchain.document_loaders import TextLoader, JSONLoader
5
import chromadb
6

7
class WazuhRAGSystem:
8
    """RAG system for Wazuh knowledge and context"""
9

10
    def __init__(self, persist_directory="./wazuh_vectorstore"):
11
        self.embeddings = OpenAIEmbeddings()
12
        self.persist_directory = persist_directory
13
        self.client = chromadb.PersistentClient(path=persist_directory)
14
        self.text_splitter = RecursiveCharacterTextSplitter(
15
            chunk_size=1000,
16
            chunk_overlap=200
17
        )
18
        self.collections = {
19
            "rules": self._create_collection("wazuh_rules"),
20
            "alerts": self._create_collection("wazuh_alerts"),
21
            "agents": self._create_collection("wazuh_agents"),
22
            "threats": self._create_collection("threat_intelligence")
23
        }
24

25
    def _create_collection(self, name: str):
26
        """Create or get a collection"""
27
        try:
28
            return self.client.get_collection(name)
29
        except:
30
            return self.client.create_collection(
31
                name=name,
32
                embedding_function=self.embeddings
33
            )
34

35
    async def index_wazuh_data(self, mcp_client: MCPClient):
36
        """Index Wazuh data for RAG"""
37

38
        # Index rules
39
        rules = await mcp_client.call_tool("ListRulesTool", {})
40
        for rule in rules.get("rules", []):
41
            self.collections["rules"].add(
42
                documents=[rule["description"]],
43
                metadatas=[{
44
                    "rule_id": rule["id"],
45
                    "level": rule["level"],
46
                    "groups": ",".join(rule.get("groups", []))
47
                }],
48
                ids=[f"rule_{rule['id']}"]
49
            )
50

51
        # Index agent information
52
        agents = await mcp_client.call_tool("GetAgentsTool", {})
53
        for agent in agents.get("agents", []):
54
            self.collections["agents"].add(
55
                documents=[json.dumps(agent)],
56
                metadatas=[{
57
                    "agent_id": agent["id"],
58
                    "name": agent["name"],
59
                    "os": agent.get("os", {}).get("platform", ""),
60
                    "status": agent["status"]
61
                }],
62
                ids=[f"agent_{agent['id']}"]
63
            )
64

65
    def similarity_search(self, query: str, collection: str, k: int = 5):
66
        """Search for similar documents"""
67
        results = self.collections[collection].query(
68
            query_texts=[query],
69
            n_results=k
70
        )
71
        return results
72

73
    async def get_relevant_context(self, query: str) -> str:
74
        """Get relevant context for a query"""
75
        contexts = []
76

77
        # Search across all collections
78
        for collection_name, collection in self.collections.items():
79
            results = collection.query(
80
                query_texts=[query],
81
                n_results=3
82
            )
83

84
            for doc, metadata in zip(
85
                results["documents"][0],
86
                results["metadatas"][0]
87
            ):
88
                contexts.append(f"[{collection_name}] {doc}")
89

90
        return "\n\n".join(contexts)

1
from langchain.chains import RetrievalQA
2
from langchain.memory import ConversationSummaryBufferMemory
3

4
class RAGEnhancedSecurityAgent:
5
    """Security agent enhanced with RAG capabilities"""
6

7
    def __init__(self, mcp_client: MCPClient, rag_system: WazuhRAGSystem):
8
        self.mcp_client = mcp_client
9
        self.rag_system = rag_system
10
        self.llm = ChatOpenAI(model="gpt-4", temperature=0)
11
        self.tools = WazuhToolRegistry.get_tools(mcp_client)
12
        self.memory = ConversationSummaryBufferMemory(
13
            llm=self.llm,
14
            max_token_limit=2000,
15
            return_messages=True
16
        )
17
        self.agent = self._create_rag_agent()
18

19
    def _create_rag_agent(self):
20
        """Create RAG-enhanced agent"""
21

22
        system_prompt = """You are a security analyst with access to:
23
        1. Wazuh MCP tools for real-time queries
24
        2. Historical knowledge base via RAG
25
        3. Threat intelligence database
26

27
        For each query:
28
        - Search the knowledge base for relevant context
29
        - Use MCP tools to get current data
30
        - Combine historical and real-time information
31
        - Provide comprehensive analysis
32
        """
33

34
        # Custom RAG tool
35
        rag_tool = Tool(
36
            name="search_knowledge_base",
37
            func=lambda q: asyncio.run(self.rag_system.get_relevant_context(q)),
38
            description="Search historical data and threat intelligence"
39
        )
40

41
        all_tools = self.tools + [rag_tool]
42

43
        prompt = ChatPromptTemplate.from_messages([
44
            ("system", system_prompt),
45
            MessagesPlaceholder(variable_name="chat_history"),
46
            ("human", "{input}"),
47
            MessagesPlaceholder(variable_name="agent_scratchpad")
48
        ])
49

50
        agent = create_openai_tools_agent(
51
            llm=self.llm,
52
            tools=all_tools,
53
            prompt=prompt
54
        )
55

56
        return AgentExecutor(
57
            agent=agent,
58
            tools=all_tools,
59
            memory=self.memory,
60
            verbose=True
61
        )
62

63
    async def analyze_with_context(self, query: str) -> Dict:
64
        """Analyze with RAG context"""
65

66
        # Get historical context
67
        context = await self.rag_system.get_relevant_context(query)
68

69
        enhanced_query = f"""
70
        Query: {query}
71

72
        Historical Context:
73
        {context}
74

75
        Use both the historical context and real-time Wazuh data to provide analysis.
76
        """
77

78
        result = await self.agent.ainvoke({"input": enhanced_query})
79

80
        return {
81
            "query": query,
82
            "context_used": context,
83
            "analysis": result["output"],
84
            "sources": self._extract_sources(result)
85
        }
86

87
    def _extract_sources(self, result: Dict) -> List[str]:
88
        """Extract data sources used"""
89
        sources = []
90
        for step in result.get("intermediate_steps", []):
91
            if hasattr(step[0], "tool"):
92
                sources.append(f"Tool: {step[0].tool}")
93
        return sources