mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-05-01 23:32:39 +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:
|
This module provides security-related functionality for CrewAI, including:
|
||||||
- Fingerprinting for component identity and tracking
|
- Fingerprinting for component identity and tracking
|
||||||
- Security configuration for controlling access and permissions
|
- Security configuration for controlling access and permissions
|
||||||
|
- Encrypted agent-to-agent communication
|
||||||
- Future: authentication, scoping, and delegation mechanisms
|
- Future: authentication, scoping, and delegation mechanisms
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from crewai.security.fingerprint import Fingerprint
|
from crewai.security.fingerprint import Fingerprint
|
||||||
from crewai.security.security_config import SecurityConfig
|
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:
|
This class manages security settings for CrewAI agents, including:
|
||||||
- Authentication credentials *TODO*
|
- Authentication credentials *TODO*
|
||||||
- Identity information (agent fingerprints)
|
- Identity information (agent fingerprints)
|
||||||
|
- Encrypted agent-to-agent communication
|
||||||
- Scoping rules *TODO*
|
- Scoping rules *TODO*
|
||||||
- Impersonation/delegation tokens *TODO*
|
- Impersonation/delegation tokens *TODO*
|
||||||
|
|
||||||
Attributes:
|
Attributes:
|
||||||
version (str): Version of the security configuration
|
version (str): Version of the security configuration
|
||||||
fingerprint (Fingerprint): The unique fingerprint automatically generated for the component
|
fingerprint (Fingerprint): The unique fingerprint automatically generated for the component
|
||||||
|
encrypted_communication (bool): Enable encrypted agent-to-agent communication
|
||||||
"""
|
"""
|
||||||
|
|
||||||
model_config = ConfigDict(
|
model_config = ConfigDict(
|
||||||
@@ -46,6 +48,11 @@ class SecurityConfig(BaseModel):
|
|||||||
default_factory=Fingerprint,
|
default_factory=Fingerprint,
|
||||||
description="Unique identifier for the component"
|
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:
|
def is_compatible(self, min_version: str) -> bool:
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -1,5 +1,5 @@
|
|||||||
import logging
|
import logging
|
||||||
from typing import Optional
|
from typing import Optional, Union, Dict, Any
|
||||||
|
|
||||||
from pydantic import Field
|
from pydantic import Field
|
||||||
|
|
||||||
@@ -7,18 +7,134 @@ from crewai.agents.agent_builder.base_agent import BaseAgent
|
|||||||
from crewai.task import Task
|
from crewai.task import Task
|
||||||
from crewai.tools.base_tool import BaseTool
|
from crewai.tools.base_tool import BaseTool
|
||||||
from crewai.utilities import I18N
|
from crewai.utilities import I18N
|
||||||
|
from crewai.security import AgentCommunicationEncryption, EncryptedMessage
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
class BaseAgentTool(BaseTool):
|
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")
|
agents: list[BaseAgent] = Field(description="List of available agents")
|
||||||
i18n: I18N = Field(
|
i18n: I18N = Field(
|
||||||
default_factory=I18N, description="Internationalization settings"
|
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:
|
def sanitize_agent_name(self, name: str) -> str:
|
||||||
"""
|
"""
|
||||||
Sanitize agent role name by normalizing whitespace and setting to lowercase.
|
Sanitize agent role name by normalizing whitespace and setting to lowercase.
|
||||||
@@ -54,6 +170,7 @@ class BaseAgentTool(BaseTool):
|
|||||||
) -> str:
|
) -> str:
|
||||||
"""
|
"""
|
||||||
Execute delegation to an agent with case-insensitive and whitespace-tolerant matching.
|
Execute delegation to an agent with case-insensitive and whitespace-tolerant matching.
|
||||||
|
Supports both encrypted and non-encrypted communication based on agent configuration.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
agent_name: Name/role of the agent to delegate to (case-insensitive)
|
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}'"
|
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:
|
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(
|
task_with_assigned_agent = Task(
|
||||||
description=task,
|
description=task,
|
||||||
agent=agent,
|
agent=target_agent,
|
||||||
expected_output=agent.i18n.slice("manager_request"),
|
expected_output=target_agent.i18n.slice("manager_request"),
|
||||||
i18n=agent.i18n,
|
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:
|
except Exception as e:
|
||||||
# Handle task creation or execution errors
|
# 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(
|
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)
|
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