diff --git a/README.md b/README.md index dcc80b5..83b794e 100644 --- a/README.md +++ b/README.md @@ -127,7 +127,7 @@ Alternatively, the progress can be tracked via the equivalent generator function ##### ProgBar *`ProgBar(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=False))`* - iterations : `int` Number of iterations for the iterative computation. @@ -146,11 +146,13 @@ Alternatively, the progress can be tracked via the equivalent generator function The update_interval in seconds controls how often the progress is flushed to the screen. Automatic mode if `update_interval=None`. +- email : `bool` (default: False) + If `True` sends an email notification after finishing the task ##### ProgPercent *`ProgPercent(iterations, track_time=True, - stream=2, title='', monitor=False, update_interval=None)`* + stream=2, title='', monitor=False, update_interval=None, email=False)`* - iterations : `int` Number of iterations for the iterative computation. @@ -167,6 +169,8 @@ Alternatively, the progress can be tracked via the equivalent generator function The update_interval in seconds controls how often the progress is flushed to the screen. Automatic mode if `update_interval=None`. +- email : `bool` (default: False) + If `True` sends an email notification after finishing the task ##### update method @@ -182,6 +186,21 @@ Alternatively, the progress can be tracked via the equivalent generator function If True, flushes the progress indicator to the output screen in each iteration. + + +##### Enable email notifications + +*`pyprind.setup_email(smtp_server, smtp_port, username, password)`* + +- smtp_server : str +- smtp_port : int +- username : str +your full email username example (pyprind@pyprind.com) +- password : str +your password + +If you want to use email notifications you can call function *`pyprind.setup_email`* only once and it will create an encrypted file with your email config and will be using it, if you want to change the email config just call *`pyprind.setup_email`* with new parameters and it will rewrite the email config file. +
diff --git a/examples/pyprind_demo.ipynb b/examples/pyprind_demo.ipynb index edf1c3d..76cbef4 100644 --- a/examples/pyprind_demo.ipynb +++ b/examples/pyprind_demo.ipynb @@ -78,7 +78,8 @@ "- [Progress Bar/Percentage Indicator - Changing the output stream](#Progress-Bar/Percentage-Indicator---Setting-a-title)\n", "- [Stopping the Progress Bar/Percentage Indicator early](#Stopping-the-Progress-Bar/Percentage-Indicator-early)\n", "- [Choosing your own progress bar style](#Choosing-your-own-progress-bar-style)\n", - "- [Controlling the update frequency](#Controlling-the-update-frequency)" + "- [Controlling the update frequency](#Controlling-the-update-frequency)\n", + "- [Setting up email notification](#Setting-up-email-notification)" ] }, { @@ -1038,6 +1039,40 @@ " time.sleep(0.2) # do some computation\n", " bar.update()" ] + }, + { + "cell_type": "markdown", + "metadata": { + "collapsed": false + }, + "source": [ + "## Setting up email notification" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": { + "collapsed": false + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "[100 %] Time elapsed: 00:00:20 | ETA: 00:00:00\n", + "Total time elapsed: 00:00:20\n" + ] + } + ], + "source": [ + "n = 100\n", + "pyprind.setup_email('smtp.example.ru', 465, 'example@example.com', 'password')\n", + "bar = pyprind.ProgPercent(n, update_interval=4, email=True)\n", + "for i in range(n):\n", + " time.sleep(0.2) # do some computation\n", + " bar.update()" + ] } ], "metadata": { diff --git a/pyprind/__init__.py b/pyprind/__init__.py index 08af3c3..e36094f 100644 --- a/pyprind/__init__.py +++ b/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' +__version__ = '2.9.9dev0' diff --git a/pyprind/email_notification.py b/pyprind/email_notification.py new file mode 100644 index 0000000..d066e65 --- /dev/null +++ b/pyprind/email_notification.py @@ -0,0 +1,86 @@ +import base64 +import os +try: + import configparser +except ImportError: + import ConfigParser as configparser + +try: + from Crypto.Cipher import AES + from Crypto import Random + crypto_import = True +except ImportError: + crypto_import = False + + +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() + if not crypto_import: + raise ValueError('crypto package is required when using' + ' email notifications.') + + @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(get_pyprind_config_dir(), + 'email_settings.ini.enc') + + def generate_key(self): + key_path = os.path.join(get_pyprind_config_dir(), '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 = str.encode(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, 'r') 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): + """Create and encrypt email config file""" + pyprind_dir = get_pyprind_config_dir() + if not os.path.exists(pyprind_dir): + os.makedirs(pyprind_dir) + file_path = os.path.join(pyprind_dir, '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', str(smtp_port)) + config.set('Email', 'username', username) + config.set('Email', 'password', password) + with open(file_path, 'w') as f: + config.write(f) + with open(file_path, 'r') as af: + cipher.encrypt(af.read()) + + +def get_pyprind_config_dir(): + home = os.path.expanduser("~") + config_path = os.path.join(home, '.pyprind') + return config_path diff --git a/pyprind/prog_class.py b/pyprind/prog_class.py index c0ce211..cb8efba 100644 --- a/pyprind/prog_class.py +++ b/pyprind/prog_class.py @@ -11,10 +11,25 @@ """ +import smtplib +import socket import time import sys import os from io import UnsupportedOperation +from email.mime.text import MIMEText +from .email_notification import AESCipher, get_pyprind_config_dir + + +try: + from StringIO import StringIO +except ImportError: + from io import StringIO + +try: + import configparser +except ImportError: + import ConfigParser as configparser try: import psutil @@ -25,7 +40,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 +69,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 = get_pyprind_config_dir() + 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() + cipher = AESCipher() + raw_data = cipher.decrypt() + buf.write(raw_data.decode()) + 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.getint('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 +203,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,9 +250,15 @@ 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: + if self.config: + self.send_email(time_info) return time_info def __str__(self): diff --git a/pyprind/progbar.py b/pyprind/progbar.py index 853febd..a9b5bb7 100755 --- a/pyprind/progbar.py +++ b/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=False): 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 diff --git a/requirements.txt b/requirements.txt index ed9e437..21da7ac 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,2 @@ -psutil>=3.2.0 \ No newline at end of file +psutil>=3.2.0 +pycryptodome==3.4