From 6b71875cd0bce3f75edbd59865bdc463529e262d Mon Sep 17 00:00:00 2001 From: Zaytsev Dmitriy Date: Fri, 2 Sep 2016 13:39:52 +0300 Subject: [PATCH] email notifications --- pyprind/__init__.py | 1 + pyprind/email_notification.py | 76 +++++++++++++++++++++++++++++++++++ pyprind/prog_class.py | 75 ++++++++++++++++++++++++++++++---- pyprind/progbar.py | 5 ++- requirements.txt | 3 +- setup.py | 5 ++- 6 files changed, 154 insertions(+), 11 deletions(-) create mode 100644 pyprind/email_notification.py diff --git a/pyprind/__init__.py b/pyprind/__init__.py index 08af3c3..b541494 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' diff --git a/pyprind/email_notification.py b/pyprind/email_notification.py new file mode 100644 index 0000000..3934751 --- /dev/null +++ b/pyprind/email_notification.py @@ -0,0 +1,76 @@ +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(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', str(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()) diff --git a/pyprind/prog_class.py b/pyprind/prog_class.py index c0ce211..bebe99f 100644 --- a/pyprind/prog_class.py +++ b/pyprind/prog_class.py @@ -9,12 +9,25 @@ 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 +from email.mime.text import MIMEText +from .email_notification import AESCipher + + +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 +38,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 +67,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() + 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 +201,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 +248,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 diff --git a/pyprind/progbar.py b/pyprind/progbar.py index 853febd..a559880 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=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 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 diff --git a/setup.py b/setup.py index ebdee00..cbb4412 100644 --- a/setup.py +++ b/setup.py @@ -12,9 +12,12 @@ from setuptools import setup, find_packages +from pip.req import parse_requirements import pyprind VERSION = pyprind.__version__ +install_requirements = parse_requirements('requirements.txt', session=False) +requires = [str(i.req) for i in install_requirements] setup(name='PyPrind', version=VERSION, @@ -25,12 +28,12 @@ packages=find_packages(), package_data={'': ['LICENSE', 'README.md', - 'requirements.txt', 'CHANGELOG.md', 'CONTRIBUTING.md'], 'tests': ['tests/test_percentage_indicator.py', 'tests/test_progress_bar.py']}, include_package_data=True, + install_requires=requires, license='BSD 3-Clause', platforms='any', classifiers=[