1# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4import datetime
5import logging
6import logging.handlers
7import os
8import socket
9import time
10
11from config import rpm_config
12
13import common
14from autotest_lib.site_utils import log_socket_server
15from autotest_lib.site_utils.rpm_control_system import rpm_infrastructure_exception
16
17LOGGING_FORMAT = rpm_config.get('GENERAL', 'logging_format')
18RECEIVERS = rpm_config.get('RPM_INFRASTRUCTURE',
19                           'email_notification_recipients').split(',')
20SUBJECT_LINE = (rpm_config.get('GENERAL', 'email_subject_line_format') %
21                socket.gethostname())
22
23
24class SuspendableSMTPHandler(logging.handlers.SMTPHandler):
25    """SMTPHandler that can have it's emails suspended."""
26    _suspend_start_time = datetime.datetime.now()
27    _suspend_time_hrs = 0
28
29
30    def suspend_emails(self, hours):
31        """Suspend email notifications.
32
33        @param hours: How many hours to suspend email notifications.
34        """
35        self._suspend_start_time = datetime.datetime.now()
36        self._suspend_time_hrs = int(hours, 0)
37
38
39    def resume_emails(self):
40        """Resume email notifications."""
41        self._suspend_time_hrs = 0
42
43
44    def emit(self, record):
45        """Emit a log record.
46
47        This subclassed version only emits the log record if emails are not
48        suspended.
49
50        @param record: Log record object we want to emit/record.
51        """
52        if datetime.datetime.now() < (self._suspend_start_time +
53                datetime.timedelta(hours=self._suspend_time_hrs)):
54            return
55        record.msg += ('\n\nTo disable these emails use rpm_client from your '
56                       'local checkout. For a 12 hour suspension run: '
57                       'site_utils/rpm_control_system/rpm_client.py -d 12')
58        return super(SuspendableSMTPHandler, self).emit(record)
59
60
61def get_log_filename(log_filename_format):
62    """Get file name of log based on given log_filename_format.
63
64    @param log_filename_format: Format to use to create the log file.
65
66    @raise Exception: If log_filename_format is None.
67    """
68    if not log_filename_format:
69            raise Exception('log_filename_format must be set.')
70
71    log_filename = os.path.abspath(time.strftime(log_filename_format))
72    log_dir = os.path.dirname(log_filename)
73    if not os.path.isdir(log_dir):
74        os.makedirs(log_dir)
75    return log_filename
76
77
78def set_up_logging(log_filename_format=None, use_log_server=False):
79    """
80    Correctly set up logging to have the correct format/level, log to a file,
81    and send out email notifications in case of error level messages.
82
83    @param log_filename_format: Format to use to create the log file.
84    @param use_log_server: True if log to a TCP server.
85
86    @returns email_handler: Logging handler used to send out email alerts.
87    """
88    if log_socket_server.LogSocketServer.port is None and use_log_server:
89        # Port is unknown, can't log to the server.
90        raise rpm_infrastructure_exception.RPMLoggingSetupError(
91                'set_up_logging failed: Log server port is unknown.')
92    if use_log_server:
93        socketHandler = logging.handlers.SocketHandler(
94                'localhost', log_socket_server.LogSocketServer.port)
95        logging.getLogger().addHandler(socketHandler)
96    else:
97        log_filename = get_log_filename(log_filename_format)
98        logging.basicConfig(filename=log_filename, level=logging.INFO,
99                            format=LOGGING_FORMAT)
100
101    if rpm_config.getboolean('GENERAL', 'debug'):
102        logging.getLogger().setLevel(logging.DEBUG)
103
104    email_handler = SuspendableSMTPHandler('localhost', 'rpm@google.com',
105                                           RECEIVERS, SUBJECT_LINE, None)
106    email_handler.setLevel(logging.ERROR)
107    email_handler.setFormatter(logging.Formatter(LOGGING_FORMAT))
108    logging.getLogger('').addHandler(email_handler)
109    return email_handler
110
111
112def start_log_server(log_filename_format):
113    """Start log server to accept logging through a TCP server.
114
115    @param log_filename_format: Format to use to create the log file.
116    """
117    log_filename = get_log_filename(log_filename_format)
118    log_socket_server.LogSocketServer.start(filename=log_filename,
119                                            level=logging.INFO,
120                                            format=LOGGING_FORMAT)
121