email notifications

fx

pyprind/__init__.py

@@ -15,6 +15,7 @@
 from .progpercent import ProgPercent
 from .generator_factory import prog_percent
 from .generator_factory import prog_bar
+from .email_notification import setup_email
 
 
 __version__ = '2.9.8'

pyprind/email_notification.py

@@ -0,0 +1,66 @@
+import base64
+import os
+import ConfigParser
+
+from Crypto.Cipher import AES
+from Crypto import Random
+
+
+class AESCipher(object):
+
+    def __init__(self):
+        self.dir_path = os.path.dirname(os.path.abspath(__file__))
+        self.key = self.generate_key()
+        self.file = None
+        self.get_current_path()
+
+    @staticmethod
+    def pad(s):
+        return s + (16 - len(s) % 16) * chr(16 - len(s) % 16)
+
+    @staticmethod
+    def unpad(s):
+        return s[:-ord(s[len(s) - 1:])]
+
+    def get_current_path(self):
+        self.file = os.path.join(self.dir_path, 'email_settings.ini.enc')
+
+    def generate_key(self):
+        key_path = os.path.join(self.dir_path, 'pyprind.key')
+        if not os.path.exists(key_path):
+            with open(key_path, 'wb') as key_file:
+                key_file.write(os.urandom(16))
+        with open(key_path, 'rb') as f:
+            key = f.read()
+        return key
+
+    def encrypt(self, text):
+        text = self.pad(text)
+        iv = Random.new().read(AES.block_size)
+        cipher = AES.new(self.key, AES.MODE_CBC, iv)
+        encrypted_mes = base64.b64encode(iv + cipher.encrypt(text))
+        with open(self.file, 'wb') as f:
+            f.write(encrypted_mes)
+
+    def decrypt(self):
+        with open(self.file, 'rb') as f:
+            enc = base64.b64decode(f.read())
+        iv = enc[:16]
+        cipher = AES.new(self.key, AES.MODE_CBC, iv)
+        return self.unpad(cipher.decrypt(enc[16:]))
+
+
+def setup_email(smtp_server, smtp_port, username, password):
+    dir_path = os.path.dirname(os.path.abspath(__file__))
+    file_path = os.path.join(dir_path, 'email_settings.ini.enc')
+    cipher = AESCipher()
+    config = ConfigParser.ConfigParser()
+    config.add_section('Email')
+    config.set('Email', 'smtp_server', smtp_server)
+    config.set('Email', 'smtp_port', smtp_port)
+    config.set('Email', 'username', username)
+    config.set('Email', 'password', password)
+    with open(file_path, 'wb') as f:
+        config.write(f)
+    with open(file_path, 'rb') as af:
+        cipher.encrypt(af.read())

pyprind/prog_class.py

@@ -9,12 +9,20 @@
 Code Repository: https://github.com/rasbt/pyprind
 PyPI: https://pypi.python.org/pypi/PyPrind
 """
-
-
+import smtplib
+import socket
 import time
 import sys
 import os
 from io import UnsupportedOperation
+import ConfigParser
+from email.mime.text import MIMEText
+from email_notification import AESCipher
+
+try:
+    from StringIO import StringIO
+except ImportError:
+    from io import StringIO
 
 try:
     import psutil
@@ -25,7 +33,7 @@
 
 class Prog():
     def __init__(self, iterations, track_time, stream, title,
-                 monitor, update_interval=None):
+                 monitor, update_interval=None, email=False):
         """ Initializes tracking object. """
         self.cnt = 0
         self.title = title
@@ -54,6 +62,49 @@ def __init__(self, iterations, track_time, stream, title,
                 self.process = psutil.Process()
         if self.track:
             self.eta = 1
+        self.config = self.load_email_config() if email else False
+
+    def load_email_config(self):
+        dir_path = os.path.dirname(os.path.abspath(__file__))
+        file_path = os.path.join(dir_path, 'email_settings.ini.enc')
+        if not os.path.exists(file_path):
+            print('The email config cannot be found, please call'
+                  ' pyprind.setup_email function')
+            return False
+        return self.parse_email_config()
+
+    @staticmethod
+    def parse_email_config():
+        buf = StringIO.StringIO()
+        cipher = AESCipher()
+        raw_data = cipher.decrypt()
+        buf.write(raw_data)
+        buf.seek(0, os.SEEK_SET)
+        config = ConfigParser.ConfigParser()
+        config.readfp(buf)
+        return config
+
+    def send_email(self, message):
+        email_address = self.config.get('Email', 'username')
+        msg = MIMEText(message, 'plain')
+        msg['From'] = email_address
+        msg['To'] = email_address
+        msg['Subject'] = 'Your task has finished'
+        password = self.config.get('Email', 'password')
+        self.config.get('Email', 'smtp_port')
+        s = smtplib.SMTP_SSL()
+        s.connect(self.config.get('Email', 'smtp_server'),
+                  self.config.get('Email', 'smtp_port'))
+        try:
+            s.login(email_address, password)
+        except smtplib.SMTPAuthenticationError as e:
+            print('Error occurred while sending email: %s' % e)
+            return False
+        try:
+            s.sendmail(email_address, [email_address], msg.as_string())
+            s.quit()
+        except socket.error as e:
+            print('Error occurred while sending email: %s' % e)
 
     def update(self, iterations=1, item_id=None, force_flush=False):
         """
@@ -145,8 +196,9 @@ def _finish(self):
             self.last_progress -= 1  # to force a refreshed _print()
             self._print()
             if self.track:
-                self._stream_out('\nTotal time elapsed: ' +
-                                 self._get_time(self.total_time))
+                message = '\nTotal time elapsed: ' + \
+                      self._get_time(self.total_time)
+                self._stream_out(message)
             self._stream_out('\n')
             self.active = False
 
@@ -191,8 +243,12 @@ def __repr__(self):
 
             cpu_mem_info = '  CPU %: {:.2f}\n'\
                            '  Memory %: {:.2f}'.format(cpu_total, mem_total)
-
-            return time_info + '\n' + cpu_mem_info
+            time_elapsed = '\nTotal time elapsed: ' + \
+                           self._get_time(self.total_time)
+            body_message = time_info + '\n' + cpu_mem_info
+            if self.config:
+                self.send_email("{}\n{}".format(time_elapsed, body_message))
+            return body_message
         else:
             return time_info
 

pyprind/progbar.py

@@ -43,9 +43,10 @@ class ProgBar(Prog):
 
     """
     def __init__(self, iterations, track_time=True, width=30, bar_char='#',
-                 stream=2, title='', monitor=False, update_interval=None):
+                 stream=2, title='', monitor=False, update_interval=None,
+                 email=True):
         Prog.__init__(self, iterations, track_time,
-                      stream, title, monitor, update_interval)
+                      stream, title, monitor, update_interval, email)
         self.bar_width = width
         self._adjust_width()
         self.bar_char = bar_char

requirements.txt

@@ -1 +1,2 @@
-psutil>=3.2.0
+psutil>=3.2.0
+pycryptodome==3.4