modules/advanced-zero-click-deployment-interface/FlowSteering/ApplicationCode/EndUserCode/EndUserClient.py

import argparse
import email
import os
import random
import re
import socket
import time
import tkinter as tk
from email.mime.image import MIMEImage
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from tkinter import ttk

import pandas as pd
from PIL import Image, ImageTk
from utils.encryption import encrypt_data, decrypt_data

# Define global variables
SERVER_EMAIL_HOST = None
SERVER_EMAIL_PORT = None
SERVER_LLAVA_HOST = None
SERVER_LLAVA_PORT = None
MYEMAIL = None
MAILSERVER = None
saveMail_directory = None
MyEmails = None
CycleNewEmails = None
BaseEmails_directory = None
# Define the default image to be sent in case of network errors
default_image=''


def receive_complete_data(client_socket): # this function is used to receive the complete data from the client, adjust the parameters as needed based on your network conditions
    received_data = b""
    count = 0
    try:
        while True:
            chunk = client_socket.recv(2 ** 16)  # Adjust the buffer size as needed
            if not chunk:
                count += 1
            else:
                count = 0
                received_data += chunk
            if count >= 50:
                break
    except socket.timeout as e:
        print('timeout')
        print(e)
        pass

    return received_data


def parse_email_data(data):  # this function gets the data from the inbox and parse it to the email data
    msg = email.message_from_bytes(data)

    Command, subject, sender, recipient = msg['Command'], msg["Subject"], msg["From"], msg["To"]
    recipient_directory = f"{saveMail_directory}/{recipient}"
    os.makedirs(recipient_directory, exist_ok=True)

    if msg.is_multipart():
        for part in msg.get_payload():
            if part.get_content_type() == "text/plain":
                body = part.get_payload()
    else:
        print(msg.get_payload())
    for part in msg.walk():
        if part.get_content_maintype() == "multipart":
            continue
        if part.get("Content-Disposition") is None:
            continue

        filename = part.get_filename()
        #filename = filename.split("\\")[-1]
        filename = filename.split("/")[-1]

        # Save the image file
        with open(os.path.join(recipient_directory, filename), "wb") as f:
            f.write(part.get_payload(decode=True))
    print(f'\n Opened and parsed new email from {sender} to {recipient} with subject {subject}')
    print(f'Email body: {body}')
    print(f'Email attachment: {filename}')

    filepath = str(f"{recipient_directory}/{filename}")
    try: #We faced some network errors resulting in images being sent partially black. To address this issue, we implemented a try-except block to handle such occurrences. Now, if an image fails to send correctly, a default image is sent for that experiment.
        with open(filepath) as f: # TEST IF THE FILE IS A VALID IMAGE
            img = MIMEImage(f.read())
    except:  # network error
        if default_image=='':
            print('Network Error: No default image is set')
            return
        else:
            filepath = default_image

    # Decrypt email data
    decrypted_body = decrypt_data(body, msg['BodyKey'])
    decrypted_filepath = decrypt_data(filepath, msg['FilePathKey'])

    return (sender, recipient, subject, decrypted_body, decrypted_filepath)


def send_Email(Command, sender, recipient, subject, body, attachment_path, SERVER_HOST, SERVER_PORT,
               AdditionalQuery=['']):  # this function sends a new email to the email server
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as client_socket:
        client_socket.connect((SERVER_HOST, SERVER_PORT))

        # Create the message
        msg = MIMEMultipart()
        msg["Command"] = Command
        msg["Subject"] = subject
        msg["From"] = sender
        msg["To"] = recipient

        if AdditionalQuery != '':
            for i in range(len(AdditionalQuery)):
                msg["AdditionalQuery" + str(i)] = AdditionalQuery[i]
        msg["AdditionalQueryNum"] = str(len(AdditionalQuery))
        msg.attach(MIMEText(body, "plain"))

        filename = attachment_path
        with open(filename, "rb") as f:
            img = MIMEImage(f.read())
            img.add_header("Content-Disposition", "attachment", filename=filename)
            msg.attach(img)
        message = msg.as_string().encode('utf-8')

        # Encrypt email data
        encrypted_message, message_key = encrypt_data(message)

        client_socket.sendall(encrypted_message)  # send the message to the server
        response = receive_complete_data(client_socket)  # get the response from the server

    return response.decode('utf-8')


def show_email_popup(email_data):  # this function shows a popup with the email data
    popup = tk.Tk()
    popup.title("New Email")
    text_sub_font = ("Helvetica", 12, "bold")
    text_font = ("Helvetica", 10)
    title_style = ttk.Style()
    title_font = ("Helvetica", 16, "bold")
    title_style.configure("Title.TLabel", font=title_font)
    ttk.Label(popup, text="NEW EMAIL!", style="Title.TLabel").pack()
    separator = ttk.Separator(popup, orient='horizontal')
    separator.pack(fill='x')
    email_text = tk.Text(popup, height=10, width=40, wrap=tk.WORD, spacing2=5, bg="#f0f0f0", relief=tk.FLAT)
    email_text.configure(state=tk.DISABLED)
    email_text.tag_configure("bold", font=text_sub_font)
    email_text.tag_configure("normal", font=text_font)
    email_text.configure(state=tk.NORMAL)
    email_text.insert(tk.END, "From: ", "bold")
    email_text.insert(tk.END, email_data[0] + "\n", "normal")
    email_text.insert(tk.END, "To: ", "bold")
    email_text.insert(tk.END, email_data[1] + "\n", "normal")
    email_text.insert(tk.END, "Subject: ", "bold")
    email_text.insert(tk.END, email_data[2] + "\n\n", "normal")
    separator = ttk.Separator(popup, orient='horizontal')
    separator.pack(fill='x')
    email_text.insert(tk.END, email_data[3] + "\n", "normal")
    email_text.configure(state=tk.DISABLED)
    email_text.pack(pady=10)
    image_path = email_data[4]
    image = Image.open(image_path)
    image.thumbnail((200, 200))  # Adjust the size as needed
    tk_image = ImageTk.PhotoImage(image)
    label = tk.Label(popup, image=tk_image, bg="#f0f0f0")
    label.image = tk_image
    label.pack()
    popup.after(5000, popup.destroy)  # destroy the popup after 5 seconds
    popup.mainloop()  # Show the popup


def check_email_inbox():  # this function checks the inbox for new emails from the server, if there are new emails it shows a popup with the email data and then calls the Handle_New_Inbox_Email function
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as client_socket:
        client_socket.connect((SERVER_EMAIL_HOST, SERVER_EMAIL_PORT))
        msg = MIMEMultipart()
        msg["Command"] = "CHECK_INBOX"
        msg["Subject"] = "CHECK_INBOX"
        msg["From"] = MYEMAIL
        msg["To"] = MAILSERVER
        msg.attach(MIMEText("Check Inbox", "plain"))
        message = msg.as_bytes()

        client_socket.sendall(message)
        inbox_data = receive_complete_data(client_socket)
        time.sleep(2)

        if inbox_data == b'No Emails':
            print(f'there are no new Emails in the inbox for you')
            return
        client_socket.close()
        try:
            email_data = parse_email_data(inbox_data)
            show_email_popup(email_data)
            Handle_New_Inbox_Email(email_data)
        except:
            pass


def read_emails_from_file():  # this function reads 5 emails from the Email csv file and returns them as a list

    Emails_df = pd.read_csv(BaseEmails_directory)
    """Tailor this function to meet your requirements; it serves as an illustration. 
    In our experiment, we utilized a single CSV file containing all the emails, located under "EndUserClientBaseEmails."
     We provided each end-user client with the same CSV file, ensuring that each client reads distinct emails from the file.
    Feel free to re-implement this function, but ensure that the construction of MyEmails follows the same pattern as the example provided below.
    """
    digit = ''

    for char in MYEMAIL:
        if char.isdigit():
            digit += char
    MyEmails = []

    for i in range(0, 6):
        Index = i + int(digit) * 6
        Email = Emails_df.iloc[Index]

        if Email['Sender'] != MYEMAIL:
            Body = Email['Body']
            EmailString = f'sender email address:{Email["Sender"]} \n email content: {Body}'
            MyEmails.append(EmailString)

    return MyEmails


def Handle_New_Inbox_Email(email_data):  # this function handles the new email from the server, it parses the email data and then sends it to the LLaVa server, then it parses the LLaVa response and sends the email to the two recipients that the LLaVa server chose if the LLaVa server chose to forward the email
    sender, recipient, subject, body, filepath = email_data

    bodyWrap = 'classify the email into one of the following options based on its content and return the JSON format.  \n  The options are \n 1."forward"- send to someone in case the content is relevant to additional contacts, Forward JSON format: "{"Classification": "forward","Subject": "forwarded email", "Body": "Begin Forwarded Email"}" \n 2."spam"- spam if its content is spam , promotional or advertisement, Spam JSON format: "{"Classification": "Spam","Subject": "", "Body": ""}".\n  3."reply"- reply in case this email requires a user response, Reply JSON format: "{"Classification": "Reply","Subject": "", "Body": ""}" . \n\n  ' \
                 f'Email content text: {body}.'

    query2 = 'Return the JSON object'

    global MyEmails

    processed_emails = []
    for i, email in enumerate(MyEmails):
        # Replace all non-ASCII characters and remove new lines and replace them with -@- so the utf-8 encoding will work
        processed_email = email.encode('ascii', 'ignore').decode('ascii').replace('\n', '-@-')
        processed_emails.append(processed_email)

    MyEmails = processed_emails

    query3=' here are five recent emails I have received: '
    for email in MyEmails:
        query3 += email + '\n'

    query3 += 'read them and say ok.'
    query4 = 'Based on the content of the five emails, choose the email addresses that might be interested in the content of this new email. return their email address'
    query5 = 'add their email address to the JSON object you built earlier, add them under { "Emails": " , "}'

    AdditionalQuery = [query2, query3, query4, query5]
    print('\n ---Sending the new email from the inbox to LLaVa for classification ---\n')
    response = send_Email('CHECK_EMAIL', MYEMAIL, sender, subject, bodyWrap, filepath, SERVER_LLAVA_HOST,
                          SERVER_LLAVA_PORT, AdditionalQuery) # Send the email to the LLaVa server for classification
    ResponsesDict = {}
    for i in range(len(AdditionalQuery) + 1):
        start_index = response.find(f'Response {i}:')
        end_index = response.find(f'Response {i + 1}:')
        if end_index == -1:
            end_index = len(response)
        ResponsesDict[i] = response[start_index:end_index]

    ClassificaionRespond1 = ResponsesDict[1]
    data = ResponsesDict[4]

    ClassificationAll = data.lower()

    start_index = ClassificationAll.find('forward')
    if start_index != -1:
        Classification = 'forward'
    else:
        start_index = ClassificationAll.find('reply')
        if start_index != -1:
            Classification = 'reply'
        else:
            start_index = ClassificationAll.find('spam')
            if start_index != -1:
                Classification = 'spam'
            else:
                Classification = 'Error'

    ClassificaionRespond1 = ClassificaionRespond1.lower() # if something goes wrong in the last JSON, the classification will be error and we will use the first JSON to classify the email
    if Classification == 'Error':
        start_index = ClassificaionRespond1.find('forward')
        if start_index != -1:
            Classification = 'forward'
        else:
            start_index = ClassificaionRespond1.find('reply')
            if start_index != -1:
                Classification = 'reply'
            else:
                start_index = ClassificaionRespond1.find('spam')
                if start_index != -1:
                    Classification = 'spam'
                else:
                    Classification = 'Error2'

    print('Classification from LLaVa is:', Classification)

    if Classification == 'reply': # if the LLaVa server chose to reply to the email, we will move the email to the Manual Folder
        print('Manual action is required for replying to this email, so it will be transferred to the Manual Folder.')
        pass
    elif Classification == 'forward':
        print('Starting to forward the emails to the correspondents')
        EmailAddresses = re.findall(r'[\w\.-]+@[\w\.-]+', data)
        Command = "SEND_EMAIL"
        EmailAddresses = list(set(EmailAddresses))
        for Email in EmailAddresses:
            recipient = Email
            response = send_Email(Command, MYEMAIL, recipient, subject, body, filepath, SERVER_EMAIL_HOST,
                                  SERVER_EMAIL_PORT)
            print(f'{response} to: {recipient}')

    elif Classification == 'spam':# if the LLaVa server chose to move the email to the spam folder, we will move the email to the Spam Folder
        print('Moving the email to the Spam Folder')
        pass
    else:
        print('Error in classification')
        pass

    # remove the first email from the list with pop and append the new email to the end of the list
    if CycleNewEmails: #this allows us to decide if we want to cycle the new emails or use the same base emails (in our experiment, we cycled the emails)
        MyEmails.pop(0)
        NewEmailString = f'sender email address:{sender} \n email content: {body}'
        MyEmails.append(NewEmailString)
    else:
        pass



def main():
    global MAILSERVER, SERVER_EMAIL_HOST, SERVER_EMAIL_PORT, SERVER_LLAVA_HOST, SERVER_LLAVA_PORT, MYEMAIL, BaseEmails_directory, saveMail_directory, MyEmails, CycleNewEmails, default_image

    MAILSERVER = 'MailServer@example.com'
    parser = argparse.ArgumentParser(description='Description of your program')
    parser.add_argument('--SERVER_EMAIL_HOST', type=str, help='Server Email IP')
    parser.add_argument('--SERVER_EMAIL_PORT', type=int, help='Server Email Port')
    parser.add_argument('--SERVER_LLAVA_HOST', type=str, help='Server LLaVa IP')
    parser.add_argument('--SERVER_LLAVA_PORT', type=int, help='Server LLaVa Port')
    parser.add_argument('--MYEMAIL', type=str, help='PersonX@example.com Email')
    parser.add_argument('--saveMail_directory', type=str, help='Directory to save the emails')
    parser.add_argument('--BaseEmails_directory', type=str, help='Directory to save the base emails')
    parser.add_argument('--CycleNewEmails', type=bool,
                        help='True if you want to cycle the new emails, False if you want to use the same base emails')
    parser.add_argument('--default_image', type=str, help='Path to the default image, if you do not want to use the default image, leave it empty')

    args = parser.parse_args()
    SERVER_EMAIL_HOST = args.SERVER_EMAIL_HOST
    SERVER_EMAIL_PORT = args.SERVER_EMAIL_PORT
    SERVER_LLAVA_HOST = args.SERVER_LLAVA_HOST
    SERVER_LLAVA_PORT = args.SERVER_LLAVA_PORT
    MYEMAIL = args.MYEMAIL
    saveMail_directory = args.saveMail_directory
    BaseEmails_directory = args.BaseEmails_directory
    CycleNewEmails = args.CycleNewEmails
    default_image = args.default_image
    MyEmails = read_emails_from_file()

    print(f'Starting the Client for Email {MYEMAIL}')

    while True:
        print('-' * 50)
        time.sleep(random.randint(10, 20))
        print('Checking the inbox for new emails')
        check_email_inbox()
        print('-' * 50)


if __name__ == '__main__':
    main()