edited app.py

app.py

@@ -1,22 +1,53 @@
-from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
+from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel,load_tool,tool
+import base64
 import datetime
+import hashlib
 import requests
 import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
-
 from Gradio_UI import GradioUI
 
-# Below is an example of a tool that does nothing. Amaze us with your creativity !
+# Message encryption/decryption tool
 @tool
-def my_custom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return type
-    #Keep this format for the description / args / args description but feel free to modify the tool
-    """A tool that does nothing yet 
+def my_custom_tool(message: str, action: str, key: str) -> str:
+    """
+    A tool that encrypts or decrypts messages using Base64 encoding for secure communication.
+
     Args:
-        arg1: the first argument
-        arg2: the second argument
+        message: The plaintext to encrypt, or the encrypted text to decrypt.
+        action: Either 'encrypt' or 'decrypt'.
+        key: The secret key used for both encryption and decryption.
     """
-    return "What magic will you build ?"
+    # Normalize action so "Encrypt" and " encrypt " both work
+    action = action.lower().strip()
+    if action not in ("encrypt", "decrypt"):
+        return "Error: action must be 'encrypt' or 'decrypt'."
+
+    if not key:
+        return "Error: key cannot be empty."
+
+    # Derive a fixed-length 32-byte key from the password via SHA-256
+    key_bytes = hashlib.sha256(key.encode()).digest()
+
+    if action == "encrypt":
+        # XOR each message byte with a repeating key byte, then encode as Base64
+        encrypted = bytes(
+            b ^ key_bytes[i % len(key_bytes)]
+            for i, b in enumerate(message.encode())
+        )
+        return base64.urlsafe_b64encode(encrypted).decode()
+
+    try:
+        # Reverse the process: decode Base64, XOR with the same key, recover plaintext
+        encrypted_bytes = base64.urlsafe_b64decode(message.encode())
+        decrypted = bytes(
+            b ^ key_bytes[i % len(key_bytes)]
+            for i, b in enumerate(encrypted_bytes)
+        )
+        return decrypted.decode()
+    except Exception as e:
+        return f"Error during decryption: {str(e)}"
 
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
@@ -55,7 +86,7 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer], ## add your tools here (don't remove final answer)
+    tools=[final_answer, my_custom_tool, get_current_time_in_timezone], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,