fix: harden email agent security (allowlist, pickle, exceptions)

end-to-end-use-cases/email_agent/email_agent.py

@@ -10,35 +10,34 @@
 import os
 import pytz
 import base64
-import pickle
 from datetime import datetime, timezone
 import json
 import ollama
 from pypdf import PdfReader
 from pathlib import Path
+from google.oauth2.credentials import Credentials
 
 SCOPES = ['https://www.googleapis.com/auth/gmail.readonly', 'https://www.googleapis.com/auth/gmail.compose']
 
 def authenticate_gmail(user_email):
     creds = None
-    token_file = f'token_{user_email}.pickle'  # Unique token file for each user
-    
+    token_file = f'token_{user_email}.json'
+
     # Load the user's token if it exists
     if os.path.exists(token_file):
-        with open(token_file, 'rb') as token:
-            creds = pickle.load(token)
-
+        creds = Credentials.from_authorized_user_file(token_file, SCOPES)
+
     # If no valid credentials, prompt the user to log in
     if not creds or not creds.valid:
         if creds and creds.expired and creds.refresh_token:
             creds.refresh(Request())
         else:
             flow = InstalledAppFlow.from_client_secrets_file('credentials.json', SCOPES)
             creds = flow.run_console()
-        
+
         # Save the new credentials to a user-specific token file
-        with open(token_file, 'wb') as token:
-            pickle.dump(creds, token)
+        with open(token_file, 'w') as token:
+            token.write(creds.to_json())
 
     # Build the Gmail API service
     service = build('gmail', 'v1', credentials=creds)
@@ -515,8 +514,14 @@ def __call__(self, user_prompt_or_tool_result, is_tool_call=False):
         self.messages.append({"role": ("ipython" if is_tool_call else "user"), "content": user_prompt_or_tool_result})
         result = self.llama()
         print(f"\nLlama returned: {result}.")
-        if type(result) == dict: # result is a dict only if it's a tool call spec
+        if isinstance(result, dict): # result is a dict only if it's a tool call spec
             function_name = result["function_name"]
+            ALLOWED_FUNCTIONS = {
+                "list_emails", "get_email_detail", "get_pdf_summary",
+                "send_email", "create_draft", "send_draft",
+            }
+            if function_name not in ALLOWED_FUNCTIONS:
+                return f"Unknown function: {function_name}. Please try again."
             func = globals()[function_name]
             parameters = result["parameters"]
             if function_name == "get_email_detail":
@@ -610,7 +615,7 @@ def llama(self):
           parameters = res['parameters']
           return {"function_name": function_name,
                   "parameters": parameters}
-        except:
+        except (json.JSONDecodeError, KeyError):
           return result
 
 

end-to-end-use-cases/email_agent/requirements.txt

@@ -1,9 +1,8 @@
-
-google-auth==2.27.0
-google-auth-oauthlib==0.4.6
-google-auth-httplib2==0.1.0
-google-api-python-client==2.34.0
+google-auth>=2.27.0
+google-auth-oauthlib>=1.0.0
+google-auth-httplib2>=0.1.0
+google-api-python-client>=2.100.0
 pytz
 beautifulsoup4
 ollama
-pypdf
+pypdf