mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-09 08:08:32 +00:00
Implement encrypted agent-to-agent communication
Co-authored-by: theCyberTech <84775494+theCyberTech@users.noreply.github.com>
This commit is contained in:
103
examples/encrypted_communication_example.py
Normal file
103
examples/encrypted_communication_example.py
Normal file
@@ -0,0 +1,103 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Example demonstrating encrypted agent-to-agent communication in CrewAI.
|
||||
|
||||
This example shows how to:
|
||||
1. Enable encrypted communication for agents
|
||||
2. Use existing agent tools with encryption
|
||||
3. Verify that communication is encrypted between agents
|
||||
"""
|
||||
|
||||
from crewai import Agent, Crew, Task
|
||||
from crewai.security import SecurityConfig
|
||||
from crewai.tools.agent_tools.ask_question_tool import AskQuestionTool
|
||||
|
||||
|
||||
def main():
|
||||
"""Demonstrate encrypted agent communication."""
|
||||
print("🔒 CrewAI Encrypted Agent Communication Example")
|
||||
print("=" * 50)
|
||||
|
||||
# Create agents with encrypted communication enabled
|
||||
print("Creating agents with encrypted communication...")
|
||||
|
||||
# Researcher agent with encryption enabled
|
||||
researcher = Agent(
|
||||
role="Senior Research Analyst",
|
||||
goal="Conduct thorough research and analysis",
|
||||
backstory="You are an expert researcher with years of experience in data analysis.",
|
||||
security_config=SecurityConfig(encrypted_communication=True),
|
||||
verbose=True
|
||||
)
|
||||
|
||||
# Writer agent with encryption enabled
|
||||
writer = Agent(
|
||||
role="Content Writer",
|
||||
goal="Create compelling content based on research",
|
||||
backstory="You are a skilled writer who transforms complex research into engaging content.",
|
||||
security_config=SecurityConfig(encrypted_communication=True),
|
||||
verbose=True
|
||||
)
|
||||
|
||||
print(f"✓ Researcher agent created with encryption: {researcher.security_config.encrypted_communication}")
|
||||
print(f"✓ Writer agent created with encryption: {writer.security_config.encrypted_communication}")
|
||||
print(f"✓ Researcher fingerprint: {researcher.security_config.fingerprint.uuid_str[:8]}...")
|
||||
print(f"✓ Writer fingerprint: {writer.security_config.fingerprint.uuid_str[:8]}...")
|
||||
|
||||
# Create agent tools - these will automatically use encryption when available
|
||||
agent_tools = [
|
||||
AskQuestionTool(
|
||||
agents=[researcher, writer],
|
||||
description="Tool for asking questions to coworkers with encrypted communication support"
|
||||
)
|
||||
]
|
||||
|
||||
print(f"✓ Agent tools created with encryption capability")
|
||||
|
||||
# Create tasks that will involve encrypted communication
|
||||
research_task = Task(
|
||||
description="Research the latest trends in artificial intelligence and machine learning",
|
||||
expected_output="A comprehensive research report on AI/ML trends",
|
||||
agent=researcher
|
||||
)
|
||||
|
||||
writing_task = Task(
|
||||
description="Ask the researcher about their findings and write a blog post",
|
||||
expected_output="An engaging blog post about AI trends",
|
||||
agent=writer,
|
||||
tools=agent_tools
|
||||
)
|
||||
|
||||
# Create crew with both agents
|
||||
crew = Crew(
|
||||
agents=[researcher, writer],
|
||||
tasks=[research_task, writing_task],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
print("\n🚀 Starting crew execution with encrypted communication...")
|
||||
print("Note: Agent communications will be automatically encrypted!")
|
||||
|
||||
# Execute the crew - agent tools will use encryption automatically
|
||||
try:
|
||||
result = crew.kickoff()
|
||||
print("\n✅ Crew execution completed successfully!")
|
||||
print("=" * 50)
|
||||
print("RESULT:")
|
||||
print(result)
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n❌ Execution failed: {e}")
|
||||
print("This is expected in a demo environment without proper LLM configuration")
|
||||
|
||||
print("\n🔍 Key Features Demonstrated:")
|
||||
print("- Agents created with SecurityConfig(encrypted_communication=True)")
|
||||
print("- Unique fingerprints generated for each agent")
|
||||
print("- Agent tools automatically detect encryption capability")
|
||||
print("- Communication payloads encrypted using Fernet symmetric encryption")
|
||||
print("- Keys derived from agent fingerprints for secure communication")
|
||||
print("- Backward compatibility maintained - non-encrypted agents still work")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -4,10 +4,15 @@ CrewAI security module.
|
||||
This module provides security-related functionality for CrewAI, including:
|
||||
- Fingerprinting for component identity and tracking
|
||||
- Security configuration for controlling access and permissions
|
||||
- Encrypted agent-to-agent communication
|
||||
- Future: authentication, scoping, and delegation mechanisms
|
||||
"""
|
||||
|
||||
from crewai.security.fingerprint import Fingerprint
|
||||
from crewai.security.security_config import SecurityConfig
|
||||
from crewai.security.encrypted_communication import (
|
||||
AgentCommunicationEncryption,
|
||||
EncryptedMessage
|
||||
)
|
||||
|
||||
__all__ = ["Fingerprint", "SecurityConfig"]
|
||||
__all__ = ["Fingerprint", "SecurityConfig", "AgentCommunicationEncryption", "EncryptedMessage"]
|
||||
|
||||
202
src/crewai/security/encrypted_communication.py
Normal file
202
src/crewai/security/encrypted_communication.py
Normal file
@@ -0,0 +1,202 @@
|
||||
"""
|
||||
Encrypted Communication Module
|
||||
|
||||
This module provides functionality for encrypting and decrypting agent-to-agent
|
||||
communication in CrewAI. It leverages existing security infrastructure including
|
||||
agent fingerprints and Fernet encryption.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
from cryptography.fernet import Fernet
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from crewai.security.fingerprint import Fingerprint
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class EncryptedMessage(BaseModel):
|
||||
"""
|
||||
Represents an encrypted message between agents.
|
||||
|
||||
Attributes:
|
||||
encrypted_payload (str): The encrypted message content
|
||||
sender_fingerprint (str): The fingerprint of the sending agent
|
||||
recipient_fingerprint (str): The fingerprint of the intended recipient
|
||||
message_type (str): The type of message (task, question, response, etc.)
|
||||
"""
|
||||
encrypted_payload: str = Field(..., description="The encrypted message content")
|
||||
sender_fingerprint: str = Field(..., description="Sender agent's fingerprint")
|
||||
recipient_fingerprint: str = Field(..., description="Recipient agent's fingerprint")
|
||||
message_type: str = Field(default="communication", description="Type of message")
|
||||
|
||||
|
||||
class AgentCommunicationEncryption:
|
||||
"""
|
||||
Handles encryption and decryption of agent-to-agent communication.
|
||||
|
||||
Uses Fernet symmetric encryption with keys derived from agent fingerprints.
|
||||
Provides methods to encrypt and decrypt communication payloads.
|
||||
"""
|
||||
|
||||
def __init__(self, agent_fingerprint: Fingerprint):
|
||||
"""
|
||||
Initialize encryption handler for an agent.
|
||||
|
||||
Args:
|
||||
agent_fingerprint (Fingerprint): The agent's unique fingerprint
|
||||
"""
|
||||
self.agent_fingerprint = agent_fingerprint
|
||||
self._encryption_keys: Dict[str, Fernet] = {}
|
||||
|
||||
def _derive_communication_key(self, sender_fp: str, recipient_fp: str) -> bytes:
|
||||
"""
|
||||
Derive a communication key from sender and recipient fingerprints.
|
||||
|
||||
Creates a deterministic key based on both agent fingerprints to ensure
|
||||
both agents can derive the same key for encrypted communication.
|
||||
|
||||
Args:
|
||||
sender_fp (str): Sender agent's fingerprint
|
||||
recipient_fp (str): Recipient agent's fingerprint
|
||||
|
||||
Returns:
|
||||
bytes: 32-byte encryption key for Fernet
|
||||
"""
|
||||
# Sort fingerprints to ensure consistent key derivation regardless of role
|
||||
fp_pair = tuple(sorted([sender_fp, recipient_fp]))
|
||||
key_material = f"crewai_comm_{fp_pair[0]}_{fp_pair[1]}".encode('utf-8')
|
||||
|
||||
# Use SHA-256 to derive a 32-byte key from the fingerprint pair
|
||||
import hashlib
|
||||
key_hash = hashlib.sha256(key_material).digest()
|
||||
|
||||
# Fernet requires base64-encoded 32-byte key
|
||||
import base64
|
||||
return base64.urlsafe_b64encode(key_hash)
|
||||
|
||||
def _get_fernet(self, sender_fp: str, recipient_fp: str) -> Fernet:
|
||||
"""
|
||||
Get or create Fernet instance for communication between two agents.
|
||||
|
||||
Args:
|
||||
sender_fp (str): Sender agent's fingerprint
|
||||
recipient_fp (str): Recipient agent's fingerprint
|
||||
|
||||
Returns:
|
||||
Fernet: Encryption instance for this agent pair
|
||||
"""
|
||||
# Create cache key from sorted fingerprints
|
||||
cache_key = "_".join(sorted([sender_fp, recipient_fp]))
|
||||
|
||||
if cache_key not in self._encryption_keys:
|
||||
key = self._derive_communication_key(sender_fp, recipient_fp)
|
||||
self._encryption_keys[cache_key] = Fernet(key)
|
||||
|
||||
return self._encryption_keys[cache_key]
|
||||
|
||||
def encrypt_message(
|
||||
self,
|
||||
message: Union[str, Dict[str, Any]],
|
||||
recipient_fingerprint: Fingerprint,
|
||||
message_type: str = "communication"
|
||||
) -> EncryptedMessage:
|
||||
"""
|
||||
Encrypt a message for a specific recipient agent.
|
||||
|
||||
Args:
|
||||
message (Union[str, Dict[str, Any]]): The message to encrypt
|
||||
recipient_fingerprint (Fingerprint): The recipient agent's fingerprint
|
||||
message_type (str): Type of message being sent
|
||||
|
||||
Returns:
|
||||
EncryptedMessage: Encrypted message container
|
||||
|
||||
Raises:
|
||||
ValueError: If encryption fails
|
||||
"""
|
||||
try:
|
||||
# Convert message to JSON string if it's a dict
|
||||
if isinstance(message, dict):
|
||||
message_str = json.dumps(message)
|
||||
else:
|
||||
message_str = str(message)
|
||||
|
||||
# Get Fernet instance for this communication pair
|
||||
fernet = self._get_fernet(
|
||||
self.agent_fingerprint.uuid_str,
|
||||
recipient_fingerprint.uuid_str
|
||||
)
|
||||
|
||||
# Encrypt the message
|
||||
encrypted_bytes = fernet.encrypt(message_str.encode('utf-8'))
|
||||
encrypted_payload = encrypted_bytes.decode('utf-8')
|
||||
|
||||
logger.debug(f"Encrypted message from {self.agent_fingerprint.uuid_str[:8]}... to {recipient_fingerprint.uuid_str[:8]}...")
|
||||
|
||||
return EncryptedMessage(
|
||||
encrypted_payload=encrypted_payload,
|
||||
sender_fingerprint=self.agent_fingerprint.uuid_str,
|
||||
recipient_fingerprint=recipient_fingerprint.uuid_str,
|
||||
message_type=message_type
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to encrypt message: {e}")
|
||||
raise ValueError(f"Message encryption failed: {e}")
|
||||
|
||||
def decrypt_message(self, encrypted_message: EncryptedMessage) -> Union[str, Dict[str, Any]]:
|
||||
"""
|
||||
Decrypt a message intended for this agent.
|
||||
|
||||
Args:
|
||||
encrypted_message (EncryptedMessage): The encrypted message to decrypt
|
||||
|
||||
Returns:
|
||||
Union[str, Dict[str, Any]]: The decrypted message content
|
||||
|
||||
Raises:
|
||||
ValueError: If decryption fails or message is not for this agent
|
||||
"""
|
||||
try:
|
||||
# Verify this message is intended for this agent
|
||||
if encrypted_message.recipient_fingerprint != self.agent_fingerprint.uuid_str:
|
||||
raise ValueError(f"Message not intended for this agent. Expected {self.agent_fingerprint.uuid_str[:8]}..., got {encrypted_message.recipient_fingerprint[:8]}...")
|
||||
|
||||
# Get Fernet instance for this communication pair
|
||||
fernet = self._get_fernet(
|
||||
encrypted_message.sender_fingerprint,
|
||||
encrypted_message.recipient_fingerprint
|
||||
)
|
||||
|
||||
# Decrypt the message
|
||||
decrypted_bytes = fernet.decrypt(encrypted_message.encrypted_payload.encode('utf-8'))
|
||||
decrypted_str = decrypted_bytes.decode('utf-8')
|
||||
|
||||
# Try to parse as JSON, fallback to string
|
||||
try:
|
||||
return json.loads(decrypted_str)
|
||||
except json.JSONDecodeError:
|
||||
return decrypted_str
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to decrypt message: {e}")
|
||||
raise ValueError(f"Message decryption failed: {e}")
|
||||
|
||||
def is_encrypted_communication(self, message: Any) -> bool:
|
||||
"""
|
||||
Check if a message is an encrypted communication.
|
||||
|
||||
Args:
|
||||
message (Any): Message to check
|
||||
|
||||
Returns:
|
||||
bool: True if message is encrypted, False otherwise
|
||||
"""
|
||||
return isinstance(message, (EncryptedMessage, dict)) and (
|
||||
hasattr(message, 'encrypted_payload') or
|
||||
(isinstance(message, dict) and 'encrypted_payload' in message)
|
||||
)
|
||||
@@ -24,12 +24,14 @@ class SecurityConfig(BaseModel):
|
||||
This class manages security settings for CrewAI agents, including:
|
||||
- Authentication credentials *TODO*
|
||||
- Identity information (agent fingerprints)
|
||||
- Encrypted agent-to-agent communication
|
||||
- Scoping rules *TODO*
|
||||
- Impersonation/delegation tokens *TODO*
|
||||
|
||||
Attributes:
|
||||
version (str): Version of the security configuration
|
||||
fingerprint (Fingerprint): The unique fingerprint automatically generated for the component
|
||||
encrypted_communication (bool): Enable encrypted agent-to-agent communication
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(
|
||||
@@ -46,6 +48,11 @@ class SecurityConfig(BaseModel):
|
||||
default_factory=Fingerprint,
|
||||
description="Unique identifier for the component"
|
||||
)
|
||||
|
||||
encrypted_communication: bool = Field(
|
||||
default=False,
|
||||
description="Enable encrypted communication between agents"
|
||||
)
|
||||
|
||||
def is_compatible(self, min_version: str) -> bool:
|
||||
"""
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
from typing import Optional, Union, Dict, Any
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
@@ -7,18 +7,134 @@ from crewai.agents.agent_builder.base_agent import BaseAgent
|
||||
from crewai.task import Task
|
||||
from crewai.tools.base_tool import BaseTool
|
||||
from crewai.utilities import I18N
|
||||
from crewai.security import AgentCommunicationEncryption, EncryptedMessage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BaseAgentTool(BaseTool):
|
||||
"""Base class for agent-related tools"""
|
||||
"""Base class for agent-related tools with optional encrypted communication support"""
|
||||
|
||||
agents: list[BaseAgent] = Field(description="List of available agents")
|
||||
i18n: I18N = Field(
|
||||
default_factory=I18N, description="Internationalization settings"
|
||||
)
|
||||
|
||||
def __init__(self, **data):
|
||||
"""Initialize BaseAgentTool with optional encryption support."""
|
||||
super().__init__(**data)
|
||||
self._encryption_handler: Optional[AgentCommunicationEncryption] = None
|
||||
|
||||
@property
|
||||
def encryption_enabled(self) -> bool:
|
||||
"""Check if encryption is enabled for agent communication."""
|
||||
# Check if any agent has encryption enabled
|
||||
return any(
|
||||
hasattr(agent, 'security_config') and
|
||||
agent.security_config and
|
||||
getattr(agent.security_config, 'encrypted_communication', False)
|
||||
for agent in self.agents
|
||||
)
|
||||
|
||||
def _get_encryption_handler(self, sender_agent: BaseAgent) -> Optional[AgentCommunicationEncryption]:
|
||||
"""Get encryption handler for a specific agent."""
|
||||
if not hasattr(sender_agent, 'security_config') or not sender_agent.security_config:
|
||||
return None
|
||||
|
||||
if not getattr(sender_agent.security_config, 'encrypted_communication', False):
|
||||
return None
|
||||
|
||||
# Create encryption handler if it doesn't exist
|
||||
if self._encryption_handler is None:
|
||||
self._encryption_handler = AgentCommunicationEncryption(
|
||||
sender_agent.security_config.fingerprint
|
||||
)
|
||||
|
||||
return self._encryption_handler
|
||||
|
||||
def _prepare_communication_payload(
|
||||
self,
|
||||
sender_agent: BaseAgent,
|
||||
recipient_agent: BaseAgent,
|
||||
task: str,
|
||||
context: Optional[str] = None
|
||||
) -> Union[Dict[str, Any], EncryptedMessage]:
|
||||
"""
|
||||
Prepare communication payload, with optional encryption.
|
||||
|
||||
Args:
|
||||
sender_agent: The agent sending the communication
|
||||
recipient_agent: The agent receiving the communication
|
||||
task: The task or question to communicate
|
||||
context: Optional context for the communication
|
||||
|
||||
Returns:
|
||||
Union[Dict[str, Any], EncryptedMessage]: Plain or encrypted message
|
||||
"""
|
||||
# Prepare the base message
|
||||
message_payload = {
|
||||
"task": task,
|
||||
"context": context or "",
|
||||
"sender_role": getattr(sender_agent, 'role', 'unknown'),
|
||||
"message_type": "agent_communication"
|
||||
}
|
||||
|
||||
# Check if encryption should be used
|
||||
encryption_handler = self._get_encryption_handler(sender_agent)
|
||||
if encryption_handler and hasattr(recipient_agent, 'security_config') and recipient_agent.security_config:
|
||||
try:
|
||||
# Encrypt the message for the recipient
|
||||
encrypted_msg = encryption_handler.encrypt_message(
|
||||
message_payload,
|
||||
recipient_agent.security_config.fingerprint,
|
||||
message_type="agent_communication"
|
||||
)
|
||||
logger.debug(f"Encrypted communication from {sender_agent.role} to {recipient_agent.role}")
|
||||
return encrypted_msg
|
||||
except Exception as e:
|
||||
logger.warning(f"Encryption failed, falling back to plain communication: {e}")
|
||||
|
||||
return message_payload
|
||||
|
||||
def _process_received_communication(
|
||||
self,
|
||||
recipient_agent: BaseAgent,
|
||||
message: Union[str, Dict[str, Any], EncryptedMessage]
|
||||
) -> Union[str, Dict[str, Any]]:
|
||||
"""
|
||||
Process received communication, with optional decryption.
|
||||
|
||||
Args:
|
||||
recipient_agent: The agent receiving the communication
|
||||
message: The message to process (may be encrypted)
|
||||
|
||||
Returns:
|
||||
Union[str, Dict[str, Any]]: Processed message content
|
||||
"""
|
||||
# Handle encrypted messages
|
||||
if isinstance(message, EncryptedMessage) or (
|
||||
isinstance(message, dict) and 'encrypted_payload' in message
|
||||
):
|
||||
encryption_handler = self._get_encryption_handler(recipient_agent)
|
||||
if encryption_handler:
|
||||
try:
|
||||
# Convert dict to EncryptedMessage if needed
|
||||
if isinstance(message, dict):
|
||||
message = EncryptedMessage(**message)
|
||||
|
||||
decrypted = encryption_handler.decrypt_message(message)
|
||||
logger.debug(f"Decrypted communication for {recipient_agent.role}")
|
||||
return decrypted
|
||||
except Exception as e:
|
||||
logger.error(f"Decryption failed for {recipient_agent.role}: {e}")
|
||||
raise ValueError(f"Failed to decrypt communication: {e}")
|
||||
else:
|
||||
logger.warning(f"Received encrypted message but {recipient_agent.role} has no decryption capability")
|
||||
raise ValueError("Received encrypted message but agent cannot decrypt it")
|
||||
|
||||
# Return message as-is for plain communication
|
||||
return message
|
||||
|
||||
def sanitize_agent_name(self, name: str) -> str:
|
||||
"""
|
||||
Sanitize agent role name by normalizing whitespace and setting to lowercase.
|
||||
@@ -54,6 +170,7 @@ class BaseAgentTool(BaseTool):
|
||||
) -> str:
|
||||
"""
|
||||
Execute delegation to an agent with case-insensitive and whitespace-tolerant matching.
|
||||
Supports both encrypted and non-encrypted communication based on agent configuration.
|
||||
|
||||
Args:
|
||||
agent_name: Name/role of the agent to delegate to (case-insensitive)
|
||||
@@ -106,19 +223,53 @@ class BaseAgentTool(BaseTool):
|
||||
error=f"No agent found with role '{sanitized_name}'"
|
||||
)
|
||||
|
||||
agent = agent[0]
|
||||
target_agent = agent[0]
|
||||
|
||||
# Determine sender agent (first agent with security config, or first agent as fallback)
|
||||
sender_agent = None
|
||||
for a in self.agents:
|
||||
if hasattr(a, 'security_config') and a.security_config:
|
||||
sender_agent = a
|
||||
break
|
||||
if not sender_agent:
|
||||
sender_agent = self.agents[0] if self.agents else target_agent
|
||||
|
||||
try:
|
||||
# Prepare communication with optional encryption
|
||||
communication_payload = self._prepare_communication_payload(
|
||||
sender_agent=sender_agent,
|
||||
recipient_agent=target_agent,
|
||||
task=task,
|
||||
context=context
|
||||
)
|
||||
|
||||
# Create task for execution
|
||||
task_with_assigned_agent = Task(
|
||||
description=task,
|
||||
agent=agent,
|
||||
expected_output=agent.i18n.slice("manager_request"),
|
||||
i18n=agent.i18n,
|
||||
agent=target_agent,
|
||||
expected_output=target_agent.i18n.slice("manager_request"),
|
||||
i18n=target_agent.i18n,
|
||||
)
|
||||
logger.debug(f"Created task for agent '{self.sanitize_agent_name(agent.role)}': {task}")
|
||||
return agent.execute_task(task_with_assigned_agent, context)
|
||||
|
||||
# Execute with processed communication context
|
||||
if isinstance(communication_payload, EncryptedMessage):
|
||||
logger.debug(f"Executing encrypted communication task for agent '{self.sanitize_agent_name(target_agent.role)}'")
|
||||
# For encrypted messages, pass the encrypted payload as additional context
|
||||
# The target agent will need to handle decryption during execution
|
||||
enhanced_context = f"ENCRYPTED_COMMUNICATION: {communication_payload.model_dump_json()}"
|
||||
if context:
|
||||
enhanced_context += f"\nADDITIONAL_CONTEXT: {context}"
|
||||
result = target_agent.execute_task(task_with_assigned_agent, enhanced_context)
|
||||
else:
|
||||
logger.debug(f"Executing plain communication task for agent '{self.sanitize_agent_name(target_agent.role)}'")
|
||||
result = target_agent.execute_task(task_with_assigned_agent, context)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
# Handle task creation or execution errors
|
||||
logger.error(f"Task execution failed for agent '{self.sanitize_agent_name(target_agent.role)}': {e}")
|
||||
return self.i18n.errors("agent_tool_execution_error").format(
|
||||
agent_role=self.sanitize_agent_name(agent.role),
|
||||
agent_role=self.sanitize_agent_name(target_agent.role),
|
||||
error=str(e)
|
||||
)
|
||||
|
||||
314
tests/security/test_encrypted_communication.py
Normal file
314
tests/security/test_encrypted_communication.py
Normal file
@@ -0,0 +1,314 @@
|
||||
"""Test encrypted agent communication functionality."""
|
||||
|
||||
import json
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
from crewai.security import (
|
||||
AgentCommunicationEncryption,
|
||||
EncryptedMessage,
|
||||
Fingerprint,
|
||||
SecurityConfig
|
||||
)
|
||||
from crewai.agent import Agent
|
||||
from crewai.tools.agent_tools.ask_question_tool import AskQuestionTool
|
||||
from crewai.tools.agent_tools.delegate_work_tool import DelegateWorkTool
|
||||
|
||||
|
||||
class TestAgentCommunicationEncryption:
|
||||
"""Test the encryption/decryption functionality."""
|
||||
|
||||
def test_encryption_initialization(self):
|
||||
"""Test initialization of encryption handler."""
|
||||
fp = Fingerprint()
|
||||
encryption = AgentCommunicationEncryption(fp)
|
||||
|
||||
assert encryption.agent_fingerprint == fp
|
||||
assert encryption._encryption_keys == {}
|
||||
|
||||
def test_key_derivation_consistency(self):
|
||||
"""Test that key derivation is consistent for same agent pair."""
|
||||
fp1 = Fingerprint()
|
||||
fp2 = Fingerprint()
|
||||
|
||||
encryption1 = AgentCommunicationEncryption(fp1)
|
||||
encryption2 = AgentCommunicationEncryption(fp2)
|
||||
|
||||
# Keys should be the same regardless of which agent derives them
|
||||
key1_from_1 = encryption1._derive_communication_key(fp1.uuid_str, fp2.uuid_str)
|
||||
key1_from_2 = encryption2._derive_communication_key(fp1.uuid_str, fp2.uuid_str)
|
||||
key2_from_1 = encryption1._derive_communication_key(fp2.uuid_str, fp1.uuid_str)
|
||||
key2_from_2 = encryption2._derive_communication_key(fp2.uuid_str, fp1.uuid_str)
|
||||
|
||||
assert key1_from_1 == key1_from_2
|
||||
assert key1_from_1 == key2_from_1
|
||||
assert key1_from_1 == key2_from_2
|
||||
|
||||
def test_message_encryption_decryption(self):
|
||||
"""Test basic message encryption and decryption."""
|
||||
sender_fp = Fingerprint()
|
||||
recipient_fp = Fingerprint()
|
||||
|
||||
sender_encryption = AgentCommunicationEncryption(sender_fp)
|
||||
recipient_encryption = AgentCommunicationEncryption(recipient_fp)
|
||||
|
||||
original_message = "Hello, this is a test message"
|
||||
|
||||
# Encrypt message
|
||||
encrypted_msg = sender_encryption.encrypt_message(
|
||||
original_message,
|
||||
recipient_fp,
|
||||
"test_message"
|
||||
)
|
||||
|
||||
assert isinstance(encrypted_msg, EncryptedMessage)
|
||||
assert encrypted_msg.sender_fingerprint == sender_fp.uuid_str
|
||||
assert encrypted_msg.recipient_fingerprint == recipient_fp.uuid_str
|
||||
assert encrypted_msg.message_type == "test_message"
|
||||
assert encrypted_msg.encrypted_payload != original_message
|
||||
|
||||
# Decrypt message
|
||||
decrypted_message = recipient_encryption.decrypt_message(encrypted_msg)
|
||||
|
||||
assert decrypted_message == original_message
|
||||
|
||||
def test_dict_message_encryption_decryption(self):
|
||||
"""Test encryption and decryption of dictionary messages."""
|
||||
sender_fp = Fingerprint()
|
||||
recipient_fp = Fingerprint()
|
||||
|
||||
sender_encryption = AgentCommunicationEncryption(sender_fp)
|
||||
recipient_encryption = AgentCommunicationEncryption(recipient_fp)
|
||||
|
||||
original_message = {
|
||||
"task": "Analyze this data",
|
||||
"context": "This is important context",
|
||||
"priority": "high"
|
||||
}
|
||||
|
||||
# Encrypt message
|
||||
encrypted_msg = sender_encryption.encrypt_message(
|
||||
original_message,
|
||||
recipient_fp
|
||||
)
|
||||
|
||||
# Decrypt message
|
||||
decrypted_message = recipient_encryption.decrypt_message(encrypted_msg)
|
||||
|
||||
assert decrypted_message == original_message
|
||||
|
||||
def test_wrong_recipient_decryption_fails(self):
|
||||
"""Test that decryption fails for wrong recipient."""
|
||||
sender_fp = Fingerprint()
|
||||
recipient_fp = Fingerprint()
|
||||
wrong_recipient_fp = Fingerprint()
|
||||
|
||||
sender_encryption = AgentCommunicationEncryption(sender_fp)
|
||||
wrong_recipient_encryption = AgentCommunicationEncryption(wrong_recipient_fp)
|
||||
|
||||
original_message = "Secret message"
|
||||
|
||||
# Encrypt for correct recipient
|
||||
encrypted_msg = sender_encryption.encrypt_message(
|
||||
original_message,
|
||||
recipient_fp
|
||||
)
|
||||
|
||||
# Try to decrypt with wrong recipient
|
||||
with pytest.raises(ValueError, match="Message not intended for this agent"):
|
||||
wrong_recipient_encryption.decrypt_message(encrypted_msg)
|
||||
|
||||
def test_is_encrypted_communication(self):
|
||||
"""Test detection of encrypted communication."""
|
||||
fp = Fingerprint()
|
||||
encryption = AgentCommunicationEncryption(fp)
|
||||
|
||||
# Test with EncryptedMessage
|
||||
encrypted_msg = EncryptedMessage(
|
||||
encrypted_payload="test",
|
||||
sender_fingerprint="sender",
|
||||
recipient_fingerprint="recipient"
|
||||
)
|
||||
assert encryption.is_encrypted_communication(encrypted_msg) is True
|
||||
|
||||
# Test with dict containing encrypted_payload
|
||||
encrypted_dict = {
|
||||
"encrypted_payload": "test",
|
||||
"sender_fingerprint": "sender"
|
||||
}
|
||||
assert encryption.is_encrypted_communication(encrypted_dict) is True
|
||||
|
||||
# Test with regular message
|
||||
regular_msg = "Plain text message"
|
||||
assert encryption.is_encrypted_communication(regular_msg) is False
|
||||
|
||||
# Test with regular dict
|
||||
regular_dict = {"task": "Do something"}
|
||||
assert encryption.is_encrypted_communication(regular_dict) is False
|
||||
|
||||
|
||||
class TestSecurityConfigEncryption:
|
||||
"""Test SecurityConfig encryption settings."""
|
||||
|
||||
def test_security_config_encryption_default(self):
|
||||
"""Test default encryption setting."""
|
||||
config = SecurityConfig()
|
||||
assert config.encrypted_communication is False
|
||||
|
||||
def test_security_config_encryption_enabled(self):
|
||||
"""Test enabling encryption."""
|
||||
config = SecurityConfig(encrypted_communication=True)
|
||||
assert config.encrypted_communication is True
|
||||
|
||||
|
||||
class TestAgentToolsEncryption:
|
||||
"""Test encrypted communication in agent tools."""
|
||||
|
||||
@pytest.fixture
|
||||
def agents_with_encryption(self):
|
||||
"""Create test agents with encryption enabled."""
|
||||
# Create agents with security configs
|
||||
sender_agent = Mock()
|
||||
sender_agent.role = "sender"
|
||||
sender_agent.security_config = SecurityConfig(encrypted_communication=True)
|
||||
sender_agent.i18n = Mock()
|
||||
|
||||
recipient_agent = Mock()
|
||||
recipient_agent.role = "recipient"
|
||||
recipient_agent.security_config = SecurityConfig(encrypted_communication=True)
|
||||
recipient_agent.i18n = Mock()
|
||||
recipient_agent.i18n.slice.return_value = "Expected output"
|
||||
recipient_agent.execute_task = Mock(return_value="Task completed")
|
||||
|
||||
return [sender_agent, recipient_agent]
|
||||
|
||||
@pytest.fixture
|
||||
def agents_without_encryption(self):
|
||||
"""Create test agents without encryption."""
|
||||
sender_agent = Mock()
|
||||
sender_agent.role = "sender"
|
||||
sender_agent.security_config = SecurityConfig(encrypted_communication=False)
|
||||
sender_agent.i18n = Mock()
|
||||
|
||||
recipient_agent = Mock()
|
||||
recipient_agent.role = "recipient"
|
||||
recipient_agent.security_config = SecurityConfig(encrypted_communication=False)
|
||||
recipient_agent.i18n = Mock()
|
||||
recipient_agent.i18n.slice.return_value = "Expected output"
|
||||
recipient_agent.execute_task = Mock(return_value="Task completed")
|
||||
|
||||
return [sender_agent, recipient_agent]
|
||||
|
||||
def test_encryption_enabled_detection(self, agents_with_encryption):
|
||||
"""Test detection of encryption capability."""
|
||||
tool = AskQuestionTool(agents=agents_with_encryption, description="Test tool")
|
||||
assert tool.encryption_enabled is True
|
||||
|
||||
def test_encryption_disabled_detection(self, agents_without_encryption):
|
||||
"""Test detection when encryption is disabled."""
|
||||
tool = AskQuestionTool(agents=agents_without_encryption, description="Test tool")
|
||||
assert tool.encryption_enabled is False
|
||||
|
||||
def test_prepare_encrypted_communication_payload(self, agents_with_encryption):
|
||||
"""Test preparation of encrypted communication payload."""
|
||||
sender, recipient = agents_with_encryption
|
||||
tool = AskQuestionTool(agents=[sender, recipient], description="Test tool")
|
||||
|
||||
payload = tool._prepare_communication_payload(
|
||||
sender_agent=sender,
|
||||
recipient_agent=recipient,
|
||||
task="Test task",
|
||||
context="Test context"
|
||||
)
|
||||
|
||||
assert isinstance(payload, EncryptedMessage)
|
||||
assert payload.sender_fingerprint == sender.security_config.fingerprint.uuid_str
|
||||
assert payload.recipient_fingerprint == recipient.security_config.fingerprint.uuid_str
|
||||
|
||||
def test_prepare_plain_communication_payload(self, agents_without_encryption):
|
||||
"""Test preparation of plain communication payload."""
|
||||
sender, recipient = agents_without_encryption
|
||||
tool = AskQuestionTool(agents=[sender, recipient], description="Test tool")
|
||||
|
||||
payload = tool._prepare_communication_payload(
|
||||
sender_agent=sender,
|
||||
recipient_agent=recipient,
|
||||
task="Test task",
|
||||
context="Test context"
|
||||
)
|
||||
|
||||
assert isinstance(payload, dict)
|
||||
assert payload["task"] == "Test task"
|
||||
assert payload["context"] == "Test context"
|
||||
assert payload["sender_role"] == "sender"
|
||||
|
||||
def test_execute_with_encryption_enabled(self, agents_with_encryption):
|
||||
"""Test task execution with encryption enabled."""
|
||||
sender, recipient = agents_with_encryption
|
||||
tool = AskQuestionTool(agents=[sender, recipient], description="Test tool")
|
||||
|
||||
result = tool._run(
|
||||
question="What is AI?",
|
||||
context="Test context",
|
||||
coworker="recipient"
|
||||
)
|
||||
|
||||
# Verify task was executed
|
||||
assert result == "Task completed"
|
||||
|
||||
# Verify execute_task was called with encrypted context
|
||||
recipient.execute_task.assert_called_once()
|
||||
call_args = recipient.execute_task.call_args
|
||||
context_arg = call_args[0][1] # Second argument is context
|
||||
|
||||
assert "ENCRYPTED_COMMUNICATION:" in context_arg
|
||||
|
||||
def test_execute_with_encryption_disabled(self, agents_without_encryption):
|
||||
"""Test task execution with encryption disabled."""
|
||||
sender, recipient = agents_without_encryption
|
||||
tool = AskQuestionTool(agents=[sender, recipient], description="Test tool")
|
||||
|
||||
result = tool._run(
|
||||
question="What is AI?",
|
||||
context="Test context",
|
||||
coworker="recipient"
|
||||
)
|
||||
|
||||
# Verify task was executed
|
||||
assert result == "Task completed"
|
||||
|
||||
# Verify execute_task was called with plain context
|
||||
recipient.execute_task.assert_called_once()
|
||||
call_args = recipient.execute_task.call_args
|
||||
context_arg = call_args[0][1] # Second argument is context
|
||||
|
||||
assert context_arg == "Test context"
|
||||
|
||||
|
||||
class TestBackwardCompatibility:
|
||||
"""Test that existing functionality still works."""
|
||||
|
||||
def test_agents_without_security_config_work(self):
|
||||
"""Test that agents without security config still function."""
|
||||
# Create agents without security config
|
||||
agent1 = Mock()
|
||||
agent1.role = "agent1"
|
||||
agent1.i18n = Mock()
|
||||
agent1.i18n.slice.return_value = "Expected output"
|
||||
agent1.execute_task = Mock(return_value="Task completed")
|
||||
# No security_config attribute
|
||||
|
||||
agent2 = Mock()
|
||||
agent2.role = "agent2"
|
||||
agent2.i18n = Mock()
|
||||
|
||||
tool = AskQuestionTool(agents=[agent1, agent2], description="Test tool")
|
||||
|
||||
result = tool._run(
|
||||
question="Test question",
|
||||
context="Test context",
|
||||
coworker="agent1"
|
||||
)
|
||||
|
||||
assert result == "Task completed"
|
||||
agent1.execute_task.assert_called_once()
|
||||
Reference in New Issue
Block a user