Implement email callback, custom exceptions

jort/__init__.py

@@ -3,3 +3,4 @@
 from .tracker import Tracker, track
 from .track_cli import track_new, track_existing
 from .reporting_callbacks import EmailNotification, SMSNotification, PrintReport
+from .exceptions import JortException, JortCredentialException

jort/config.py

@@ -1,7 +1,18 @@
 import os
+import json
 from pathlib import Path
 
+
 # Create internal jort directory
 JORT_DIR = f"{os.path.expanduser('~')}/.jort"
 Path(f"{JORT_DIR}/").mkdir(mode=0o700, parents=True, exist_ok=True)
-Path(f"{JORT_DIR}/config").touch(mode=0o600, exist_ok=True)
+Path(f"{JORT_DIR}/config").touch(mode=0o600, exist_ok=True)
+
+
+def get_config_data():
+    with open(f"{JORT_DIR}/config", "r") as f:
+        try:
+            config_data = json.load(f)
+        except json.decoder.JSONDecodeError:
+            config_data = {}
+    return config_data

jort/exceptions.py

@@ -0,0 +1,9 @@
+class JortException(Exception):
+    pass
+
+
+class JortCredentialException(JortException):
+    """
+    Exception for missing credentials.
+    """
+    pass

jort/jort_exe.py

@@ -30,6 +30,12 @@ def main():
         help='PID of existing job to track',
     )
 
+    # Save stdout/stderr output
+    parser.add_argument('-o',
+                        '--output',
+                        action='store_true',
+                        help='save stdout/stderr output')
+
     # Send SMS at job completion
     parser.add_argument('-s',
                         '--sms',
@@ -57,23 +63,30 @@ def main():
     args = parser.parse_args()
 
     if args.init:
-        with open(f"{config.JORT_DIR}/config", "r") as f:
-            try:
-                config_data = json.load(f)
-            except json.decoder.JSONDecodeError:
-                config_data = {}
+        config_data = config.get_config_data()
         input_config_data = {
+            "machine": input('What name should this device go by? ({}) '
+                             .format(config_data.get("machine", ""))),
             "email": input('What email to use? ({}) '
-                           .format(config_data.get("email"))),
+                           .format(config_data.get("email", ""))),
+            "smtp_server": input('What SMTP server does your email use? ({}) '
+                                 .format(config_data.get("smtp_server", ""))),
+            "email_password": getpass.getpass('Email password? ({}) '
+                                              .format(("*"*16 
+                                                       if config_data.get("email_password", "") is not None 
+                                                       else ""))),
             "twilio_receive_number": input('What phone number to receive SMS? ({}) '
-                                           .format(config_data.get("twilio_receive_number"))),
+                                           .format(config_data.get("twilio_receive_number", ""))),
             "twilio_send_number": input('What Twilio number to send SMS? ({}) '
-                                        .format(config_data.get("twilio_send_number"))),
+                                        .format(config_data.get("twilio_send_number", ""))),
             "twilio_account_sid": input('Twilio Account SID? ({}) '
-                                        .format(config_data.get("twilio_account_sid"))),
+                                        .format(config_data.get("twilio_account_sid", ""))),
             "twilio_auth_token": getpass.getpass('Twilio Auth Token? ({}) '
-                                                 .format("*"*len(config_data.get("twilio_auth_token", "")))),
+                                                 .format(("*"*16 
+                                                         if config_data.get("twilio_auth_token", "") is not None 
+                                                         else "")))
         }
+        # Only save inputs if they aren't empty
         for key in input_config_data:
             if input_config_data[key] != "":
                 config_data[key] = input_config_data[key]
@@ -90,7 +103,7 @@ def main():
         joined_command = ' '.join(args.command)
         print(f"Tracking command '{joined_command}'")
         track_cli.track_new(joined_command,
-                            store_stdout=False,
+                            store_stdout=args.output,
                             save_filename=None,
                             send_sms=args.sms,
                             send_email=args.email,

jort/reporting_callbacks.py

@@ -1,8 +1,17 @@
 import os
 import json
+import smtplib
+import ssl
+import email
+from email import encoders
+from email.mime.base import MIMEBase
+from email.mime.text import MIMEText
+from email.mime.multipart import MIMEMultipart
+from email.mime.application import MIMEApplication
 import twilio.rest
 import humanfriendly
 from . import config
+from . import exceptions
 
 
 class Callback(object):
@@ -22,7 +31,8 @@ def __init__(self):
 
     def format_message(self, payload):
         return (
-            f'The job \'{payload["short_name"]}\' finished running '
+            f'\n'
+            f'The job \'{payload["name"]}\' finished running '
             f'in {humanfriendly.format_timespan(payload["runtime"])}'
         )
 
@@ -31,34 +41,44 @@ def execute(self, payload):
 
 
 class SMSNotification(Callback):
+    """
+    Send SMS notifications to and from numbers managed by your Twilio account.
+    """
     def __init__(self, receive_number=None):
-        self.receive_number = receive_number
-        if receive_number is None:
-            with open(f"{config.JORT_DIR}/config", "r") as f:
-                config_data = json.load(f)
-                self.receive_number = config_data["twilio_receive_number"]
-                self.send_number = config_data["twilio_send_number"]
-                self.twilio_account_sid = config_data["twilio_account_sid"]
-                self.twilio_auth_token = config_data["twilio_auth_token"]
+        config_data = config.get_config_data()
+        self.receive_number = config_data.get("twilio_receive_number")
+        self.send_number = config_data.get("twilio_send_number")
+        self.twilio_account_sid = config_data.get("twilio_account_sid")
+        self.twilio_auth_token = config_data.get("twilio_auth_token")
+        if receive_number is not None:
+            self.receive_number = receive_number
+
+        if self.twilio_account_sid is None or self.twilio_auth_token is None:
+            raise exceptions.JortCredentialException("Missing Twilio credentials, add with `jort -i` command")
+        if self.send_number is None:
+            raise exceptions.JortException("Missing Twilio sending number, add with `jort -i` command")
+        if self.receive_number is None:
+            raise exceptions.JortException("Missing receiving number")
 
     def format_message(self, payload):
         if payload["status"] == "success":
             return (
-                f'Your job \'{payload["short_name"]}\' successfully completed '
+                f'Your job \'{payload["name"]}\' successfully completed '
                 f'in {humanfriendly.format_timespan(payload["runtime"])}'
             )
         elif payload["status"] == "error":
-            error_name = payload["error_message"].split(":")[0]
+            error_text = payload["error_message"].split(":")[0]
             return (
-                f'Your job \'{payload["short_name"]}\' exited in error ({error_name}) '
+                f'Your job \'{payload["name"]}\' exited in error ({error_text}) '
                 f'after {humanfriendly.format_timespan(payload["runtime"])}'
             )
-        else:
+        elif payload["status"] == "finished":
             return (
-                f'Your job \'{payload["short_name"]}\' finished running '
+                f'Your job \'{payload["name"]}\' finished running '
                 f'in {humanfriendly.format_timespan(payload["runtime"])}'
             )
-
+        else:
+            raise exceptions.JortException(f'Invalid status: {payload["status"]}')
 
     def execute(self, payload):
         client = twilio.rest.Client(self.twilio_account_sid,
@@ -69,17 +89,126 @@ def execute(self, payload):
 
 
 class EmailNotification(Callback):
+    """
+    Send email notifications to and from your email account.
+    """
     def __init__(self, email=None):
-        self.email = email
-        if email is None:
-            with open(f"{config.JORT_DIR}/config", "r") as f:
-                config_data = json.load(f)
-                self.email = config_data["email"]
+        config_data = config.get_config_data()
+        self.email = config_data.get("email")
+        self.smtp_server = config_data.get("smtp_server")
+        self.email_password = config_data.get("email_password")
+        if email is not None:
+            self.email = email
+
+        if self.email_password is None:
+            raise exceptions.JortCredentialException("Missing email password, add with `jort -i` command")
+        if self.smtp_server is None:
+            raise exceptions.JortException("Missing SMTP server, add with `jort -i` command")
+        if self.email is None:
+            raise exceptions.JortException("Missing email")
 
     def format_message(self, payload):
-        return ""
+        if payload["machine"] is not None:
+            machine_text = f' on machine {payload["machine"]}'
+            html_machine_text = f' on machine <strong>{payload["machine"]}</strong>'
+        else:
+            machine_text = ''
+            html_machine_text = ''
+
+        if payload["status"] == "success":
+            subject = "[jort] Your job finished successfully!"
+
+            summary_text = (
+                f'Your job `{payload["name"]}` completed at {payload["date_modified"]} (UTC)'
+                f'{machine_text} with no errors.'
+            )
+            html_summary_text = (
+                f'Your job <strong>{payload["name"]}</strong> completed at <strong>{payload["date_modified"]}</strong> (UTC)'
+                f'{html_machine_text} with no errors.'
+            )
+        elif payload["status"] == "error":
+            subject = "[jort] Your job exited with an error"
+
+            error_text = payload["error_message"].split(":")[0]
+
+            summary_text = (
+                f'Your job `{payload["name"]}` exited at {payload["date_modified"]} (UTC)'
+                f'{machine_text} with error `{error_text}`.'
+            )
+            html_summary_text = (
+                f'Your job <strong>{payload["name"]}</strong> exited at <strong>{payload["date_modified"]}</strong> (UTC)'
+                f'{html_machine_text} with error <strong>{error_text}</strong>.'
+            )
+        elif payload["status"] == "finished":
+            subject = "[jort] Your job has finished!"
+
+            summary_text = (
+                f'Your job `{payload["name"]}` finished running at {payload["date_modified"]} (UTC)'
+                f'{machine_text}.'
+            )
+            html_summary_text = (
+                f'Your job <strong>{payload["name"]}</strong> finished running at <strong>{payload["date_modified"]}</strong> (UTC)'
+                f'{html_machine_text}.'
+            )
+        else:
+            raise exceptions.JortException(f'Invalid status: {payload["status"]}')
+
+        runtime_text = f'The job\'s total runtime was {humanfriendly.format_timespan(payload["runtime"])}.'
+        html_runtime_text = f'The job\'s total runtime was <strong>{humanfriendly.format_timespan(payload["runtime"])}</strong>.'
+
+        body = (
+            f'{summary_text}\r\n'
+            f'{runtime_text}\r\n'
+            f'--\r\n'
+            f'jort'
+        )
+        html_body = (
+            f'<html>'
+            f'<head></head>'
+            f'<body>'
+            f'  <p>{html_summary_text}</p>'
+            f'  <p>{html_runtime_text}</p>'
+            f'  <p>--<br>jort</p>'
+            f'</body>'
+            f'</html>'
+        )
+        email_data = {
+            "subject": subject,
+            "body": body,
+            "html_body": html_body,
+        }
+        return email_data
 
     def execute(self, payload):
-        pass
+        email_data = self.format_message(payload)
+
+        message = MIMEMultipart("alternative")
+        message.attach(MIMEText(email_data["body"], "plain"))
+        message.attach(MIMEText(email_data["html_body"], "html"))
+
+        if payload["stdout_fn"] is not None:
+            stdout_path = f'{config.JORT_DIR}/{payload["stdout_fn"]}'
+            with open(stdout_path, "r") as f:
+                attachment = MIMEApplication(f.read(), _subtype="txt")
+            attachment.add_header("Content-Disposition", "attachment", filename="output.txt")
+
+            message_mix = MIMEMultipart("mixed")
+            message_mix.attach(message)
+            message_mix.attach(attachment)
+            message = message_mix
+
+        message["Subject"] = email_data["subject"]
+        message["From"] = self.email
+        message["To"] = self.email
+
+        # Secure connection
+        context = ssl.create_default_context()
+        with smtplib.SMTP_SSL(self.smtp_server, port=465, context=context) as server:
+            server.login(self.email, self.email_password)
+            server.sendmail(message["From"], message["To"], message.as_string())
+
+
+
+
 
 

jort/track_cli.py

@@ -63,9 +63,9 @@ def track_new(command,
     if verbose:
         pprint(payload)
     if save_filename or store_stdout:
-        f = open(stdout_path, "a+")
-        f.write(f"$ {command}\n")
-        f.close()
+        with open(stdout_path, "a+") as f:
+            f.write(f"{command}\n")
+            f.write(f"--\n")
 
     buffer = ""
     temp_start = time.time()
@@ -74,9 +74,8 @@ def track_new(command,
             if verbose:
                 print("Buffered! (Not sent)", [buffer])
             if save_filename or store_stdout:
-                f = open(stdout_path, "a+")
-                f.write(buffer)
-                f.close()
+                with open(stdout_path, "a+") as f:
+                    f.write(buffer)
 
             payload['status'] = 'running'
             datetime_utils.update_payload_times(payload)
@@ -92,14 +91,15 @@ def track_new(command,
     if verbose:
         print("Buffered!", [buffer])
     if save_filename or store_stdout:
-        f = open(stdout_path, "a+")
-        f.write(buffer)
-        f.close()
+        with open(stdout_path, "a+") as f:
+            f.write(buffer)
+
+    p.wait()
 
     if verbose:
-        print("Exit code:", p.poll())
+        print(f"Exit code: {p.returncode}")
 
-    if p.returncode in [0, None]:
+    if p.returncode == 0:
         payload["status"] = "success"
     else:
         payload["status"] = "error"

jort/tracker.py

@@ -3,6 +3,7 @@
 import logging
 import functools
 
+from . import config
 from . import checkpoint
 from . import datetime_utils
 
@@ -14,6 +15,7 @@ def __init__(self, logname="tracker.log", verbose=0):
         1 for INFO, and 2 for DEBUG.
         """
         self.date_created = datetime_utils.get_iso_date()
+        self.machine = config.get_config_data().get("machine")
         self.checkpoints = {}
         self.open_checkpoint_payloads = {}
         self.logname = logname
@@ -46,10 +48,10 @@ def start(self, name=None, date_created=None):
         self.open_checkpoint_payloads[name] = {
             "user_id": None,
             "job_id": None,
-            "short_name": name,
+            "name": name,
             "long_name": name,
             "status": "running",
-            "machine": None,
+            "machine": self.machine,
             "date_created": start,
             "date_modified": now,
             "runtime": datetime_utils.get_runtime(start, now),