Spaces:

posix4e
/

brain

Runtime error

brain / Brain /src /rising_plugin /gmail /email_plugin.py

Move to a new way of bundling requirements

62bd6c6 over 3 years ago

13.1 kB

import email
import imaplib
import json
import mimetypes
import os
import re
import smtplib
import time
import base64
from email.header import decode_header
from email.message import EmailMessage
from socket import socket

from bs4 import BeautifulSoup

from Brain.src.common.utils import PROXY_IP, PROXY_PORT

# email variables
EMAIL_SMTP_HOST = "smtp.gmail.com"
EMAIL_SMTP_PORT = 587
EMAIL_IMAP_SERVER = "imap.gmail.com"
EMAIL_SIGNATURE = "This was sent by Rising Brain"


class EmailPlugin:
    def send_email(
        self, sender: str, pwd: str, to: str, subject: str, body: str, to_send: bool
    ) -> str:
        return self.send_email_with_attachment_internal(
            sender=sender,
            pwd=pwd,
            to=to,
            title=subject,
            message=body,
            attachment=None,
            attachment_path=None,
            to_send=to_send,
        )

    def send_email_with_attachment(
        self,
        sender: str,
        pwd: str,
        to: str,
        subject: str,
        body: str,
        filename: str,
        to_send: bool,
    ) -> str:
        attachment_path = filename
        attachment = os.path.basename(filename)
        return self.send_email_with_attachment_internal(
            sender=sender,
            pwd=pwd,
            to=to,
            title=subject,
            message=body,
            attachment_path=attachment_path,
            attachment=attachment,
            to_send=to_send,
        )

    def send_email_with_attachment_internal(
        self,
        sender: str,
        pwd: str,
        to: str,
        title: str,
        message: str,
        attachment_path: str | None,
        attachment: str | None,
        to_send: bool,
    ) -> str:
        """Send an email
        Args:
            sender (str): The email of the sender
            pwd (str): The password of the sender
            to (str): The email of the recipient
            title (str): The title of the email
            message (str): The message content of the email
        Returns:
            str: Any error messages
        """
        email_sender = sender
        email_password = pwd

        msg = EmailMessage()
        msg["Subject"] = title
        msg["From"] = email_sender
        msg["To"] = to

        signature = EMAIL_SIGNATURE
        if signature:
            message += f"\n{signature}"

        msg.set_content(message)

        if attachment_path:
            ctype, encoding = mimetypes.guess_type(attachment_path)
            if ctype is None or encoding is not None:
                # No guess could be made, or the file is encoded (compressed)
                ctype = "application/octet-stream"
            maintype, subtype = ctype.split("/", 1)
            with open(file=attachment_path, mode="rb") as fp:
                msg.add_attachment(
                    fp.read(), maintype=maintype, subtype=subtype, filename=attachment
                )

        if to_send:
            smtp_host = EMAIL_SMTP_HOST
            smtp_port = EMAIL_SMTP_PORT
            # send email
            with smtplib.SMTP(host=smtp_host, port=smtp_port) as smtp:
                smtp.ehlo()
                smtp.starttls()
                smtp.login(user=email_sender, password=email_password)
                smtp.send_message(msg)
                smtp.quit()
            return f"Email was sent to {to}!"
        else:
            conn = self.imap_open(
                imap_folder="[Gmail]/Drafts",
                email_sender=email_sender,
                email_password=email_password,
            )
            conn.append(
                mailbox="[Gmail]/Drafts",
                flags="",
                date_time=imaplib.Time2Internaldate(time.time()),
                message=str(msg).encode("UTF-8"),
            )
            return f"Email went to [Gmail]/Drafts!"

    def read_emails(
        self,
        sender: str,
        pwd: str,
        imap_folder: str = "inbox",
        imap_search_command: str = "UNSEEN",
        limit: int = 5,
        page: int = 1,
    ) -> str:
        """Read emails from an IMAP mailbox.
        This function reads emails from a specified IMAP folder, using a given IMAP search command, limits, and page numbers.
        It returns a list of emails with their details, including the sender, recipient, date, CC, subject, and message body.
        Args:
            sender (str): The email of the sender
            pwd (str): The password of the sender
            imap_folder (str, optional): The name of the IMAP folder to read emails from. Defaults to "inbox".
            imap_search_command (str, optional): The IMAP search command to filter emails. Defaults to "UNSEEN".
            limit (int, optional): Number of email's the function should return. Defaults to 5 emails.
            page (int, optional): The index of the page result the function should resturn. Defaults to 0, the first page.
        Returns:
            str: A list of dictionaries containing email details if there are any matching emails.
        """
        email_sender = sender
        imap_folder = self.adjust_imap_folder_for_gmail(
            imap_folder=imap_folder, email_sender=email_sender
        )
        imap_folder = self.enclose_with_quotes(imap_folder)
        imap_search_ar = self.split_imap_search_command(imap_search_command)
        email_password = pwd

        mark_as_seen = "False"
        if isinstance(mark_as_seen, str):
            mark_as_seen = json.loads(mark_as_seen.lower())

        conn = self.imap_open(
            imap_folder=imap_folder,
            email_sender=email_sender,
            email_password=email_password,
        )

        imap_keyword = imap_search_ar[0]
        if len(imap_search_ar) == 1:
            _, search_data = conn.search(None, imap_keyword)
        else:
            argument = self.enclose_with_quotes(imap_search_ar[1])
            _, search_data = conn.search(None, imap_keyword, argument)

        messages = []
        for num in search_data[0].split():
            if mark_as_seen:
                message_parts = "(RFC822)"
            else:
                message_parts = "(BODY.PEEK[])"
            _, msg_data = conn.fetch(message_set=num, message_parts=message_parts)
            for response_part in msg_data:
                if isinstance(response_part, tuple):
                    msg = email.message_from_bytes(response_part[1])

                    # If the subject has unknown encoding, return blank
                    if msg["Subject"] is not None:
                        subject, encoding = decode_header(msg["Subject"])[0]
                    else:
                        subject = ""
                        encoding = ""

                    if isinstance(subject, bytes):
                        try:
                            # If the subject has unknown encoding, return blank
                            if encoding is not None:
                                subject = subject.decode(encoding)
                            else:
                                subject = ""
                        except [LookupError] as e:
                            pass

                    body = self.get_email_body(msg)
                    # Clean email body
                    body = self.clean_email_body(body)

                    from_address = msg["From"]
                    to_address = msg["To"]
                    date = msg["Date"]
                    cc = msg["CC"] if msg["CC"] else ""

                    messages.append(
                        {
                            "from": from_address,
                            "to": to_address,
                            "date": date,
                            "cc": cc,
                            "subject": subject,
                            "body": body,
                        }
                    )

        conn.logout()
        if not messages:
            messages.append(
                {
                    "from": "",
                    "to": "",
                    "date": "",
                    "cc": "",
                    "subject": "",
                    "body": "There are no Emails",
                }
            )
            return json.dumps(messages)

        # Confirm that integer parameters are the right type
        limit = int(limit)
        page = int(page)

        # Validate parameter values
        if limit < 1:
            raise ValueError("Error: The message limit should be 1 or greater")

        page_count = len(messages) // limit + (len(messages) % limit > 0)

        if page < 1 or page > page_count:
            raise ValueError(
                "Error: The page value references a page that is not part of the results"
            )

        # Calculate paginated indexes
        start_index = len(messages) - (page * limit + 1)
        end_index = start_index + limit
        start_index = max(start_index, 0)

        # Return paginated indexes
        if start_index == end_index:
            return json.dumps([messages[start_index]])
        else:
            return json.dumps(messages[start_index:end_index])

    def adjust_imap_folder_for_gmail(self, imap_folder: str, email_sender: str) -> str:
        if "@gmail" in email_sender.lower() or "@googlemail" in email_sender.lower():
            if "sent" in imap_folder.lower():
                return '"[Gmail]/Sent Mail"'
            if "draft" in imap_folder.lower():
                return "[Gmail]/Drafts"
        return imap_folder

    def imap_open(
        self, imap_folder: str, email_sender: str, email_password: str
    ) -> imaplib.IMAP4_SSL:
        # Create a new socket object for later connections as a proxy

        # IMAP Server Connect
        imap_server = EMAIL_IMAP_SERVER
        conn = imaplib.IMAP4_SSL(imap_server)
        conn.login(user=email_sender, password=email_password)
        conn.select(imap_folder)
        return conn

    def get_email_body(self, msg: email.message.Message) -> str:
        if msg.is_multipart():
            for part in msg.walk():
                content_type = part.get_content_type()
                content_disposition = str(part.get("Content-Disposition"))
                if (
                    content_type == "text/plain"
                    and "attachment" not in content_disposition
                ):
                    # If the email body has unknown encoding, return null
                    try:
                        return part.get_payload(decode=True).decode()
                    except UnicodeDecodeError as e:
                        pass
        else:
            try:
                # If the email body has unknown encoding, return null
                return msg.get_payload(decode=True).decode()
            except UnicodeDecodeError as e:
                pass

    def enclose_with_quotes(self, s):
        # Check if string contains whitespace
        has_whitespace = bool(re.search(r"\s", s))

        # Check if string is already enclosed by quotes
        is_enclosed = s.startswith(("'", '"')) and s.endswith(("'", '"'))

        # If string has whitespace and is not enclosed by quotes, enclose it with double quotes
        if has_whitespace and not is_enclosed:
            return f'"{s}"'
        else:
            return s

    def split_imap_search_command(self, input_string):
        input_string = input_string.strip()
        parts = input_string.split(maxsplit=1)
        parts = [part.strip() for part in parts]

        return parts

    def clean_email_body(self, email_body):
        """Remove formating and URL's from an email's body
        Args:
            email_body (str, optional): The email's body
        Returns:
            str: The email's body without any formating or URL's
        """

        # If body is None, return an empty string
        if email_body is None:
            email_body = ""

        # Remove any HTML tags
        email_body = BeautifulSoup(email_body, "html.parser")
        email_body = email_body.get_text()

        # Remove return characters
        email_body = "".join(email_body.splitlines())

        # Remove extra spaces
        email_body = " ".join(email_body.split())

        # Remove unicode characters
        email_body = email_body.encode("ascii", "ignore")
        email_body = email_body.decode("utf-8", "ignore")

        # Remove any remaining URL's
        email_body = re.sub(r"http\S+", "", email_body)

        return email_body

    def write_attachment(self, filename: str, file_content: str) -> (str, str):
        # create folder for temporarily saving attached file
        milliseconds = int(time.time() * 1000)
        file_path = f"Brain/assets/{milliseconds}/{filename}"
        file_directory = f"Brain/assets/{milliseconds}"
        os.mkdir(file_directory)

        # file write
        file_content = base64.b64decode(file_content).decode("utf-8")
        file = open(file_path, "w")
        file.write(file_content)
        file.close()
        return file_path, file_directory

	import email
	import imaplib
	import json
	import mimetypes
	import os
	import re
	import smtplib
	import time
	import base64
	from email.header import decode_header
	from email.message import EmailMessage
	from socket import socket

	from bs4 import BeautifulSoup

	from Brain.src.common.utils import PROXY_IP, PROXY_PORT

	# email variables
	EMAIL_SMTP_HOST = "smtp.gmail.com"
	EMAIL_SMTP_PORT = 587
	EMAIL_IMAP_SERVER = "imap.gmail.com"
	EMAIL_SIGNATURE = "This was sent by Rising Brain"


	class EmailPlugin:
	def send_email(
	self, sender: str, pwd: str, to: str, subject: str, body: str, to_send: bool
	) -> str:
	return self.send_email_with_attachment_internal(
	sender=sender,
	pwd=pwd,
	to=to,
	title=subject,
	message=body,
	attachment=None,
	attachment_path=None,
	to_send=to_send,
	)

	def send_email_with_attachment(
	self,
	sender: str,
	pwd: str,
	to: str,
	subject: str,
	body: str,
	filename: str,
	to_send: bool,
	) -> str:
	attachment_path = filename
	attachment = os.path.basename(filename)
	return self.send_email_with_attachment_internal(
	sender=sender,
	pwd=pwd,
	to=to,
	title=subject,
	message=body,
	attachment_path=attachment_path,
	attachment=attachment,
	to_send=to_send,
	)

	def send_email_with_attachment_internal(
	self,
	sender: str,
	pwd: str,
	to: str,
	title: str,
	message: str,
	attachment_path: str \| None,
	attachment: str \| None,
	to_send: bool,
	) -> str:
	"""Send an email

	Args:
	sender (str): The email of the sender
	pwd (str): The password of the sender
	to (str): The email of the recipient
	title (str): The title of the email
	message (str): The message content of the email

	Returns:
	str: Any error messages
	"""
	email_sender = sender
	email_password = pwd

	msg = EmailMessage()
	msg["Subject"] = title
	msg["From"] = email_sender
	msg["To"] = to

	signature = EMAIL_SIGNATURE
	if signature:
	message += f"\n{signature}"

	msg.set_content(message)

	if attachment_path:
	ctype, encoding = mimetypes.guess_type(attachment_path)
	if ctype is None or encoding is not None:
	# No guess could be made, or the file is encoded (compressed)
	ctype = "application/octet-stream"
	maintype, subtype = ctype.split("/", 1)
	with open(file=attachment_path, mode="rb") as fp:
	msg.add_attachment(
	fp.read(), maintype=maintype, subtype=subtype, filename=attachment
	)

	if to_send:
	smtp_host = EMAIL_SMTP_HOST
	smtp_port = EMAIL_SMTP_PORT
	# send email
	with smtplib.SMTP(host=smtp_host, port=smtp_port) as smtp:
	smtp.ehlo()
	smtp.starttls()
	smtp.login(user=email_sender, password=email_password)
	smtp.send_message(msg)
	smtp.quit()
	return f"Email was sent to {to}!"
	else:
	conn = self.imap_open(
	imap_folder="[Gmail]/Drafts",
	email_sender=email_sender,
	email_password=email_password,
	)
	conn.append(
	mailbox="[Gmail]/Drafts",
	flags="",
	date_time=imaplib.Time2Internaldate(time.time()),
	message=str(msg).encode("UTF-8"),
	)
	return f"Email went to [Gmail]/Drafts!"

	def read_emails(
	self,
	sender: str,
	pwd: str,
	imap_folder: str = "inbox",
	imap_search_command: str = "UNSEEN",
	limit: int = 5,
	page: int = 1,
	) -> str:
	"""Read emails from an IMAP mailbox.

	This function reads emails from a specified IMAP folder, using a given IMAP search command, limits, and page numbers.
	It returns a list of emails with their details, including the sender, recipient, date, CC, subject, and message body.

	Args:
	sender (str): The email of the sender
	pwd (str): The password of the sender
	imap_folder (str, optional): The name of the IMAP folder to read emails from. Defaults to "inbox".
	imap_search_command (str, optional): The IMAP search command to filter emails. Defaults to "UNSEEN".
	limit (int, optional): Number of email's the function should return. Defaults to 5 emails.
	page (int, optional): The index of the page result the function should resturn. Defaults to 0, the first page.

	Returns:
	str: A list of dictionaries containing email details if there are any matching emails.
	"""
	email_sender = sender
	imap_folder = self.adjust_imap_folder_for_gmail(
	imap_folder=imap_folder, email_sender=email_sender
	)
	imap_folder = self.enclose_with_quotes(imap_folder)
	imap_search_ar = self.split_imap_search_command(imap_search_command)
	email_password = pwd

	mark_as_seen = "False"
	if isinstance(mark_as_seen, str):
	mark_as_seen = json.loads(mark_as_seen.lower())

	conn = self.imap_open(
	imap_folder=imap_folder,
	email_sender=email_sender,
	email_password=email_password,
	)

	imap_keyword = imap_search_ar[0]
	if len(imap_search_ar) == 1:
	_, search_data = conn.search(None, imap_keyword)
	else:
	argument = self.enclose_with_quotes(imap_search_ar[1])
	_, search_data = conn.search(None, imap_keyword, argument)

	messages = []
	for num in search_data[0].split():
	if mark_as_seen:
	message_parts = "(RFC822)"
	else:
	message_parts = "(BODY.PEEK[])"
	_, msg_data = conn.fetch(message_set=num, message_parts=message_parts)
	for response_part in msg_data:
	if isinstance(response_part, tuple):
	msg = email.message_from_bytes(response_part[1])

	# If the subject has unknown encoding, return blank
	if msg["Subject"] is not None:
	subject, encoding = decode_header(msg["Subject"])[0]
	else:
	subject = ""
	encoding = ""

	if isinstance(subject, bytes):
	try:
	# If the subject has unknown encoding, return blank
	if encoding is not None:
	subject = subject.decode(encoding)
	else:
	subject = ""
	except [LookupError] as e:
	pass

	body = self.get_email_body(msg)
	# Clean email body
	body = self.clean_email_body(body)

	from_address = msg["From"]
	to_address = msg["To"]
	date = msg["Date"]
	cc = msg["CC"] if msg["CC"] else ""

	messages.append(
	{
	"from": from_address,
	"to": to_address,
	"date": date,
	"cc": cc,
	"subject": subject,
	"body": body,
	}
	)

	conn.logout()
	if not messages:
	messages.append(
	{
	"from": "",
	"to": "",
	"date": "",
	"cc": "",
	"subject": "",
	"body": "There are no Emails",
	}
	)
	return json.dumps(messages)

	# Confirm that integer parameters are the right type
	limit = int(limit)
	page = int(page)

	# Validate parameter values
	if limit < 1:
	raise ValueError("Error: The message limit should be 1 or greater")

	page_count = len(messages) // limit + (len(messages) % limit > 0)

	if page < 1 or page > page_count:
	raise ValueError(
	"Error: The page value references a page that is not part of the results"
	)

	# Calculate paginated indexes
	start_index = len(messages) - (page * limit + 1)
	end_index = start_index + limit
	start_index = max(start_index, 0)

	# Return paginated indexes
	if start_index == end_index:
	return json.dumps([messages[start_index]])
	else:
	return json.dumps(messages[start_index:end_index])

	def adjust_imap_folder_for_gmail(self, imap_folder: str, email_sender: str) -> str:
	if "@gmail" in email_sender.lower() or "@googlemail" in email_sender.lower():
	if "sent" in imap_folder.lower():
	return '"[Gmail]/Sent Mail"'
	if "draft" in imap_folder.lower():
	return "[Gmail]/Drafts"
	return imap_folder

	def imap_open(
	self, imap_folder: str, email_sender: str, email_password: str
	) -> imaplib.IMAP4_SSL:
	# Create a new socket object for later connections as a proxy

	# IMAP Server Connect
	imap_server = EMAIL_IMAP_SERVER
	conn = imaplib.IMAP4_SSL(imap_server)
	conn.login(user=email_sender, password=email_password)
	conn.select(imap_folder)
	return conn

	def get_email_body(self, msg: email.message.Message) -> str:
	if msg.is_multipart():
	for part in msg.walk():
	content_type = part.get_content_type()
	content_disposition = str(part.get("Content-Disposition"))
	if (
	content_type == "text/plain"
	and "attachment" not in content_disposition
	):
	# If the email body has unknown encoding, return null
	try:
	return part.get_payload(decode=True).decode()
	except UnicodeDecodeError as e:
	pass
	else:
	try:
	# If the email body has unknown encoding, return null
	return msg.get_payload(decode=True).decode()
	except UnicodeDecodeError as e:
	pass

	def enclose_with_quotes(self, s):
	# Check if string contains whitespace
	has_whitespace = bool(re.search(r"\s", s))

	# Check if string is already enclosed by quotes
	is_enclosed = s.startswith(("'", '"')) and s.endswith(("'", '"'))

	# If string has whitespace and is not enclosed by quotes, enclose it with double quotes
	if has_whitespace and not is_enclosed:
	return f'"{s}"'
	else:
	return s

	def split_imap_search_command(self, input_string):
	input_string = input_string.strip()
	parts = input_string.split(maxsplit=1)
	parts = [part.strip() for part in parts]

	return parts

	def clean_email_body(self, email_body):
	"""Remove formating and URL's from an email's body

	Args:
	email_body (str, optional): The email's body

	Returns:
	str: The email's body without any formating or URL's
	"""

	# If body is None, return an empty string
	if email_body is None:
	email_body = ""

	# Remove any HTML tags
	email_body = BeautifulSoup(email_body, "html.parser")
	email_body = email_body.get_text()

	# Remove return characters
	email_body = "".join(email_body.splitlines())

	# Remove extra spaces
	email_body = " ".join(email_body.split())

	# Remove unicode characters
	email_body = email_body.encode("ascii", "ignore")
	email_body = email_body.decode("utf-8", "ignore")

	# Remove any remaining URL's
	email_body = re.sub(r"http\S+", "", email_body)

	return email_body

	def write_attachment(self, filename: str, file_content: str) -> (str, str):
	# create folder for temporarily saving attached file
	milliseconds = int(time.time() * 1000)
	file_path = f"Brain/assets/{milliseconds}/{filename}"
	file_directory = f"Brain/assets/{milliseconds}"
	os.mkdir(file_directory)

	# file write
	file_content = base64.b64decode(file_content).decode("utf-8")
	file = open(file_path, "w")
	file.write(file_content)
	file.close()
	return file_path, file_directory