Update get_gmail_data.py

get_gmail_data.py

@@ -3,6 +3,9 @@ import base64
 from bs4 import BeautifulSoup
 import re
 import jwt
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from cryptography.hazmat.backends import default_backend
+import os
 class GmailDataExtractor:
 
     def __init__(self,jwt:str  , user_input: str = None) -> None:
@@ -133,6 +136,7 @@ class GmailDataExtractor:
         """
         print("process_messages")
         message_id = message.get("id")
+        encrypted_message_id = self.encrypt_message_id(message_id)
         if not message_id:
             return None, None, [], False
         subject=''
@@ -178,9 +182,27 @@ class GmailDataExtractor:
 
                     if data:
                         # Save only the first 10 characters of the attachment data
-                        return subject,text ,{"filename":filename , "data":data} , company_from_gmail
+                        return subject,text ,{"filename":filename , "data":data} , company_from_gmail , encrypted_message_id
 
-        return subject, text,None , company_from_gmail
+        return subject, text,None , company_from_gmail , encrypted_message_id
+
+    def encrypt_message_id(self,message_id:str):
+        key = os.getenv('AES_KEY')
+
+        # Initialize AES cipher with the key and CBC mode
+        cipher = Cipher(algorithms.AES(key), modes.CBC(), backend=default_backend())
+        
+        # Create a encryptor object
+        encryptor = cipher.encryptor()
+        
+        # Pad the message_id to be a multiple of 16 bytes (AES block size)
+        # This is necessary for AES encryption
+        message_id_padded = message_id.encode().rjust(32, b'\0')
+        
+        # Encrypt the padded message_id
+        ciphertext = encryptor.update(message_id_padded) + encryptor.finalize()
+        return ciphertext
+        
 
     def extract_domain_from_email(self,email_string):
         # Extracting the email address using regex
@@ -241,14 +263,15 @@ class GmailDataExtractor:
         messages = self.__fetch_messages()
         results = []
         for message in messages:
-            subject, body, attachment_data , company_name = self.__process_message(message)
+            subject, body, attachment_data , company_name , encrypt_mssg_id = self.__process_message(message)
 
             """ Handling None values """
             body = body if body is not None else None
             attachment_data = attachment_data if attachment_data is not None else {}
             company_associated = company_name if company_name is not None else None
+            en_msg_id = encrypt_mssg_id if encrypt_mssg_id is not None else None
 
-            results.append({"body": body, "attachment_data": [attachment_data] ,'company_associated':company_associated})
+            results.append({"body": body, "attachment_data": [attachment_data] ,'company_associated':company_associated , "message_id":en_msg_id})
 
         return {"results": results}