1# -*- coding: utf-8 -*-
2import sys
3import traceback
4
5from django.conf import settings
6from django.core.mail import send_mail
7from django.core.management import BaseCommand
8
9
10class EmailNotificationCommand(BaseCommand):
11    """
12    A BaseCommand subclass which adds sending email fuctionality.
13
14    Subclasses will have an extra command line option ``--email-notification``
15    and will be able to send emails by calling ``send_email_notification()``
16    if SMTP host and port are specified in settings. The handling of the
17    command line option is left to the management command implementation.
18    Configuration is done in settings.EMAIL_NOTIFICATIONS dict.
19
20    Configuration example::
21
22        EMAIL_NOTIFICATIONS = {
23            'scripts.my_script': {
24                'subject': 'my_script subject',
25                'body': 'my_script body',
26                'from_email': 'from_email@example.com',
27                'recipients': ('recipient0@example.com',),
28                'no_admins': False,
29                'no_traceback': False,
30                'notification_level': 0,
31                'fail_silently': False
32            },
33            'scripts.another_script': {
34                ...
35            },
36            ...
37        }
38
39    Configuration explained:
40        subject:            Email subject.
41        body:               Email body.
42        from_email:         Email from address.
43        recipients:         Sequence of email recipient addresses.
44        no_admins:          When True do not include ADMINS to recipients.
45        no_traceback:       When True do not include traceback to email body.
46        notification_level: 0: send email on fail, 1: send email always.
47        fail_silently:      Parameter passed to django's send_mail().
48    """
49
50    def add_arguments(self, parser):
51        parser.add_argument('--email-notifications',
52                            action='store_true',
53                            default=False,
54                            dest='email_notifications',
55                            help='Send email notifications for command.')
56        parser.add_argument('--email-exception',
57                            action='store_true',
58                            default=False,
59                            dest='email_exception',
60                            help='Send email for command exceptions.')
61
62    def run_from_argv(self, argv):
63        """Overriden in order to access the command line arguments."""
64        self.argv_string = ' '.join(argv)
65        super().run_from_argv(argv)
66
67    def execute(self, *args, **options):
68        """
69        Overriden in order to send emails on unhandled exception.
70
71        If an unhandled exception in ``def handle(self, *args, **options)``
72        occurs and `--email-exception` is set or `self.email_exception` is
73        set to True send an email to ADMINS with the traceback and then
74        reraise the exception.
75        """
76        try:
77            super().execute(*args, **options)
78        except Exception:
79            if options['email_exception'] or getattr(self, 'email_exception', False):
80                self.send_email_notification(include_traceback=True)
81            raise
82
83    def send_email_notification(self, notification_id=None, include_traceback=False, verbosity=1):
84        """
85        Send email notifications.
86
87        Reads settings from settings.EMAIL_NOTIFICATIONS dict, if available,
88        using ``notification_id`` as a key or else provides reasonable
89        defaults.
90        """
91        # Load email notification settings if available
92        if notification_id is not None:
93            try:
94                email_settings = settings.EMAIL_NOTIFICATIONS.get(notification_id, {})
95            except AttributeError:
96                email_settings = {}
97        else:
98            email_settings = {}
99
100        # Exit if no traceback found and not in 'notify always' mode
101        if not include_traceback and not email_settings.get('notification_level', 0):
102            print(self.style.ERROR("Exiting, not in 'notify always' mode."))
103            return
104
105        # Set email fields.
106        subject = email_settings.get('subject', "Django extensions email notification.")
107
108        command_name = self.__module__.split('.')[-1]
109
110        body = email_settings.get(
111            'body',
112            "Reporting execution of command: '%s'" % command_name
113        )
114
115        # Include traceback
116        if include_traceback and not email_settings.get('no_traceback', False):
117            try:
118                exc_type, exc_value, exc_traceback = sys.exc_info()
119                trb = ''.join(traceback.format_tb(exc_traceback))
120                body += "\n\nTraceback:\n\n%s\n" % trb
121            finally:
122                del exc_traceback
123
124        # Set from address
125        from_email = email_settings.get('from_email', settings.DEFAULT_FROM_EMAIL)
126
127        # Calculate recipients
128        recipients = list(email_settings.get('recipients', []))
129
130        if not email_settings.get('no_admins', False):
131            recipients.extend(settings.ADMINS)
132
133        if not recipients:
134            if verbosity > 0:
135                print(self.style.ERROR("No email recipients available."))
136            return
137
138        # Send email...
139        send_mail(subject, body, from_email, recipients,
140                  fail_silently=email_settings.get('fail_silently', True))
141