utils/ui.py

import os
import smtplib
from datetime import datetime
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.utils import formatdate, make_msgid

import pytz
from markdown import markdown


def get_timezone_by_ip(ip, session):
    try:
        data = session.get(f'https://worldtimeapi.org/api/ip/{ip}').json()
        return data['timezone']
    except Exception:
        return 'UTC'


def ts_to_str(timestamp, timezone):
    # Create a timezone-aware datetime object from the UNIX timestamp
    dt = datetime.fromtimestamp(timestamp, pytz.utc)

    # Convert the timezone-aware datetime object to the target timezone
    target_timezone = pytz.timezone(timezone)
    localized_dt = dt.astimezone(target_timezone)

    # Format the datetime object to the specified string format
    return localized_dt.strftime('%Y-%m-%d %H:%M:%S (%Z%z)')


def send_email(job_info):
    if job_info.get('email'):
        try:
            email_info = job_info.copy()
            email_serv = os.getenv('EMAIL_SERV')
            email_port = os.getenv('EMAIL_PORT')
            email_addr = os.getenv('EMAIL_ADDR')
            email_pass = os.getenv('EMAIL_PASS')
            email_form = os.getenv('EMAIL_FORM')
            email_subj = os.getenv('EMAIL_SUBJ')

            for key, value in email_info.items():
                if key.endswith("time") and value:
                    email_info[key] = ts_to_str(value, get_timezone_by_ip(email_info['ip']))

            server = smtplib.SMTP(email_serv, int(email_port))
            # server.starttls()

            server.login(email_addr, email_pass)
            msg = MIMEMultipart("alternative")
            msg["From"] = email_addr
            msg["To"] = email_info['email']
            msg["Subject"] = email_subj.format(**email_info)
            msg["Date"] = formatdate(localtime=True)
            msg["Message-ID"] = make_msgid()

            msg.attach(MIMEText(markdown(email_form.format(**email_info)), 'html'))
            msg.attach(MIMEText(email_form.format(**email_info), 'plain'))

            server.sendmail(email_addr, email_info['email'], msg.as_string())
            server.quit()
            gr.Info('Email notification sent.')
        except Exception as e:
            gr.Warning('Failed to send email notification due to error: ' + str(e))