feat: add encryption support #20
Conversation
WalkthroughThe changes introduce enhancements to the Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant API
participant Environment
participant Utils
Client->>API: Send Action
API->>Environment: Process Action
Environment->>Utils: Encrypt Action Data (if key present)
Utils-->>Environment: Encrypted Data
Environment->>API: Send Encrypted Data
API->>Client: Return Response
Client->>API: Request Response
API->>Environment: Process Response
Environment->>Utils: Decrypt Response Data (if key present)
Utils-->>Environment: Decrypted Data
Environment->>API: Send Decrypted Data
API->>Client: Return Decrypted Response
Poem
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Files ignored due to path filters (1)
poetry.lockis excluded by!**/*.lock
Files selected for processing (6)
- crab/core/environment.py (3 hunks)
- crab/server/api.py (1 hunks)
- crab/utils/init.py (1 hunks)
- crab/utils/encryption.py (1 hunks)
- pyproject.toml (1 hunks)
- test/core/test_utils.py (1 hunks)
Additional comments not posted (14)
test/core/test_utils.py (3)
1-13: LGTM!The copyright notice follows the Apache License, Version 2.0.
The code changes are approved.
16-18: LGTM!The import statements are necessary for the test function.
The code changes are approved.
21-26: LGTM!The
test_encrypt_decryptfunction is correctly implemented and tests the encryption and decryption process.The code changes are approved.
crab/utils/__init__.py (3)
1-13: LGTM!The copyright notice follows the Apache License, Version 2.0.
The code changes are approved.
15-25: LGTM!The import statements are necessary for the public API.
The code changes are approved.
27-35: LGTM!The
__all__list is correctly defined and includes all necessary functions.The code changes are approved.
crab/server/api.py (2)
14-24: LGTM!The import statements are necessary for the functionality of the
raw_actionendpoint.The code changes are approved.
29-29: LGTM!The
ENC_KEYvariable is necessary for encryption and decryption.The code changes are approved.
crab/utils/encryption.py (3)
26-40: Ensure key length and add error handling.The function is correctly implemented, but consider the following improvements:
- Verify that the key length is 256 bits (32 bytes).
- Add error handling for encryption failures.
def encrypt_message(plaintext: str, key: bytes) -> str: + if len(key) != 32: + raise ValueError("Key must be 256 bits (32 bytes) long.") nounce = os.urandom(12) cipher = Cipher(algorithms.AES(key), modes.GCM(nounce), backend=default_backend()) encryptor = cipher.encryptor() ciphertext = encryptor.update(plaintext.encode()) + encryptor.finalize() return base64.b64encode(nounce + ciphertext + encryptor.tag).decode("utf-8")
43-61: Ensure key length and add error handling.The function is correctly implemented, but consider the following improvements:
- Verify that the key length is 256 bits (32 bytes).
- Add error handling for decryption failures.
def decrypt_message(encrypted: str, key: bytes) -> str: + if len(key) != 32: + raise ValueError("Key must be 256 bits (32 bytes) long.") encrypted = base64.b64decode(encrypted) nounce = encrypted[:12] ciphertext = encrypted[12:-16] tag = encrypted[-16:] cipher = Cipher( algorithms.AES(key), modes.GCM(nounce, tag), backend=default_backend() ) decryptor = cipher.decryptor() return (decryptor.update(ciphertext) + decryptor.finalize()).decode("utf-8")
64-77: Ensure secure handling of environment variable and add error handling.The function is correctly implemented, but consider the following improvements:
- Ensure that the environment variable is securely handled.
- Add error handling for missing or invalid environment variables.
def generate_key_from_env() -> Optional[bytes]: enc_key = os.environ.get("CRAB_ENC_KEY") if not enc_key: logger.warning("CRAB_ENC_KEY is not set, connection will not be encrypted.") return None + if len(enc_key) < 32: + logger.warning("CRAB_ENC_KEY is too short, it should be at least 32 characters long.") + return None return hashlib.sha256(enc_key.encode("utf-8")).digest()pyproject.toml (1)
29-29: LGTM!The addition of the
cryptographydependency is appropriate for the encryption functionality.The code changes are approved.
crab/core/environment.py (2)
95-96: LGTM!The initialization of
_enc_keyusinggenerate_key_from_env()is appropriate for the encryption functionality.The code changes are approved.
218-241: Ensure correct application of encryption and decryption logic and add error handling.The function is correctly implemented, but consider the following improvements:
- Ensure that the encryption and decryption logic is correctly applied.
- Add error handling for encryption and decryption failures.
def _action_endpoint(self, action: Action, parameters: dict[str, Any]): if self._client is not None and not action.local: data = json.dumps( { "action": action.to_raw_action(), "parameters": action.parameters(**parameters).model_dump(), } ) content_type = "application/json" if self._enc_key is not None: data = encrypt_message(data, self._enc_key) content_type = "text/plain" # send action to remote machine response = self._client.post( "/raw_action", content=data, headers={"Content-Type": content_type}, ) resp_content = response.content.decode("utf-8") if self._enc_key is not None: resp_content = decrypt_message(resp_content, self._enc_key) resp_json = json.loads(resp_content) return resp_json["action_returns"] else: # or directly execute it action = action.set_kept_param(env=self) return action.run(**parameters)
|
Overall it looks very good. We can consider add an API test for encrypt mode. Ref: https://github.com/camel-ai/crab/blob/main/test/server/test_api.py |
|
close #18 |
There was a problem hiding this comment.
Actionable comments posted: 0
Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Files selected for processing (3)
- crab/server/api.py (2 hunks)
- pyproject.toml (2 hunks)
- test/server/test_api.py (2 hunks)
Files skipped from review as they are similar to previous changes (2)
- crab/server/api.py
- pyproject.toml
Additional comments not posted (4)
test/server/test_api.py (4)
15-18: LGTM!The imports are necessary and correctly included.
27-33: LGTM!The
mock_envfixture is correctly implemented and promotes code reuse.
36-41: LGTM!The
test_raw_action_unencryptedfunction is correctly implemented and tests various scenarios with unencrypted actions.
44-50: LGTM!The
test_raw_action_encryptedfunction is correctly implemented and tests various scenarios with encrypted actions.
There was a problem hiding this comment.
Actionable comments posted: 0
Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Files selected for processing (1)
- crab/agents/backend_models/camel_model.py (7 hunks)
Additional comments not posted (8)
crab/agents/backend_models/camel_model.py (8)
47-47: Type annotation updated to modern syntax.The update from
Union[ModelType, str]toModelType | straligns with the latest Python typing standards, enhancing readability and maintainability.
64-64: Updated type annotations in constructor.The change from
Optional[Dict[str, Any]]todict[str, Any] | Nonein the constructor's parameters argument improves clarity and aligns with modern Python typing practices.
73-73: Type annotation forclientattribute updated.The update from
Optional[ExternalToolAgent]toExternalToolAgent | Nonefor theclientattribute is a positive change, enhancing type clarity and consistency.
85-85: Type annotation updated inresetmethod.The change from
Optional[List[Action]]tolist[Action] | Nonefor theaction_spaceparameter in theresetmethod is a good practice, aligning with the newer Python typing conventions.
114-115: Updated return type annotation in_convert_action_to_schema.The update from
Optional[List[OpenAIFunction]]tolist[OpenAIFunction] | Noneimproves type clarity and aligns with modern Python typing standards.
121-121: Simplified list type annotation in_convert_tool_calls_to_action_list.The change from
List[ActionOutput]tolist[ActionOutput]simplifies the type annotation and aligns with the newer, more readable Python typing conventions.
133-133: Updated parameter type annotation inchatmethod.The change from
List[Tuple[str, MessageType]]tolist[tuple[str, MessageType]]in themessagesparameter of thechatmethod enhances readability and consistency with modern Python typing practices.
135-135: Updated variable type annotation withinchatmethod.The update from
List[Image.Image]tolist[Image.Image]for theimage_listvariable within thechatmethod aligns with the latest Python typing conventions, improving clarity.
This reverts commit 48b6f4e.
Add encryption support for communication between client and server. User can enable encryption by setting env variable
CRAB_ENC_KEY.Summary by CodeRabbit
New Features
Bug Fixes
Tests
Chores