• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/python
2# Copyright 2015 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""
7Mail the content of standard input.
8
9Example usage:
10  Use pipe:
11     $ echo "Some content" |./gmail_lib.py -s "subject" abc@bb.com xyz@gmail.com
12
13  Manually input:
14     $ ./gmail_lib.py -s "subject" abc@bb.com xyz@gmail.com
15     > Line 1
16     > Line 2
17     Ctrl-D to end standard input.
18"""
19import argparse
20import base64
21import httplib2
22import logging
23import sys
24import os
25import random
26from email.mime.text import MIMEText
27
28import common
29from autotest_lib.client.bin import utils
30from autotest_lib.client.common_lib import global_config
31from autotest_lib.server import site_utils
32
33try:
34  from apiclient.discovery import build as apiclient_build
35  from apiclient import errors as apiclient_errors
36  from oauth2client import file as oauth_client_fileio
37except ImportError as e:
38  apiclient_build = None
39  logging.debug("API client for gmail disabled. %s", e)
40
41# TODO(akeshet) These imports needs to come after the apiclient imports, because
42# of a sys.path war between chromite and autotest crbug.com/622988
43from autotest_lib.server import utils as server_utils
44from chromite.lib import retry_util
45
46try:
47    from chromite.lib import metrics
48except ImportError:
49    metrics = utils.metrics_mock
50
51
52DEFAULT_CREDS_FILE = global_config.global_config.get_config_value(
53        'NOTIFICATIONS', 'gmail_api_credentials', default=None)
54RETRY_DELAY = 5
55RETRY_BACKOFF_FACTOR = 1.5
56MAX_RETRY = 10
57RETRIABLE_MSGS = [
58        # User-rate limit exceeded
59        r'HttpError 429',]
60
61class GmailApiException(Exception):
62    """Exception raised in accessing Gmail API."""
63
64
65class Message():
66    """An email message."""
67
68    def __init__(self, to, subject, message_text):
69        """Initialize a message.
70
71        @param to: The recievers saperated by comma.
72                   e.g. 'abc@gmail.com,xyz@gmail.com'
73        @param subject: String, subject of the message
74        @param message_text: String, content of the message.
75        """
76        self.to = to
77        self.subject = subject
78        self.message_text = message_text
79
80
81    def get_payload(self):
82        """Get the payload that can be sent to the Gmail API.
83
84        @return: A dictionary representing the message.
85        """
86        message = MIMEText(self.message_text)
87        message['to'] = self.to
88        message['subject'] = self.subject
89        return {'raw': base64.urlsafe_b64encode(message.as_string())}
90
91
92class GmailApiClient():
93    """Client that talks to Gmail API."""
94
95    def __init__(self, oauth_credentials):
96        """Init Gmail API client
97
98        @param oauth_credentials: Path to the oauth credential token.
99        """
100        if not apiclient_build:
101            raise GmailApiException('Cannot get apiclient library.')
102
103        storage = oauth_client_fileio.Storage(oauth_credentials)
104        credentials = storage.get()
105        if not credentials or credentials.invalid:
106            raise GmailApiException('Invalid credentials for Gmail API, '
107                                    'could not send email.')
108        http = credentials.authorize(httplib2.Http())
109        self._service = apiclient_build('gmail', 'v1', http=http)
110
111
112    def send_message(self, message, ignore_error=True):
113        """Send an email message.
114
115        @param message: Message to be sent.
116        @param ignore_error: If True, will ignore any HttpError.
117        """
118        try:
119            # 'me' represents the default authorized user.
120            message = self._service.users().messages().send(
121                    userId='me', body=message.get_payload()).execute()
122            logging.debug('Email sent: %s' , message['id'])
123        except apiclient_errors.HttpError as error:
124            if ignore_error:
125                logging.error('Failed to send email: %s', error)
126            else:
127                raise
128
129
130def send_email(to, subject, message_text, retry=True, creds_path=None):
131    """Send email.
132
133    @param to: The recipients, separated by comma.
134    @param subject: Subject of the email.
135    @param message_text: Text to send.
136    @param retry: If retry on retriable failures as defined in RETRIABLE_MSGS.
137    @param creds_path: The credential path for gmail account, if None,
138                       will use DEFAULT_CREDS_FILE.
139    """
140    auth_creds = server_utils.get_creds_abspath(
141        creds_path or DEFAULT_CREDS_FILE)
142    if not auth_creds or not os.path.isfile(auth_creds):
143        logging.error('Failed to send email to %s: Credential file does not '
144                      'exist: %s. If this is a prod server, puppet should '
145                      'install it. If you need to be able to send email, '
146                      'find the credential file from chromeos-admin repo and '
147                      'copy it to %s', to, auth_creds, auth_creds)
148        return
149    client = GmailApiClient(oauth_credentials=auth_creds)
150    m = Message(to, subject, message_text)
151    retry_count = MAX_RETRY if retry else 0
152
153    def _run():
154        """Send the message."""
155        client.send_message(m, ignore_error=False)
156
157    def handler(exc):
158        """Check if exc is an HttpError and is retriable.
159
160        @param exc: An exception.
161
162        @return: True if is an retriable HttpError.
163        """
164        if not isinstance(exc, apiclient_errors.HttpError):
165            return False
166
167        error_msg = str(exc)
168        should_retry = any([msg in error_msg for msg in RETRIABLE_MSGS])
169        if should_retry:
170            logging.warning('Will retry error %s', exc)
171        return should_retry
172
173    success = False
174    try:
175        retry_util.GenericRetry(
176                handler, retry_count, _run, sleep=RETRY_DELAY,
177                backoff_factor=RETRY_BACKOFF_FACTOR)
178        success = True
179    finally:
180        metrics.Counter('chromeos/autotest/send_email/count').increment(
181                fields={'success': success})
182
183
184if __name__ == '__main__':
185    logging.basicConfig(level=logging.DEBUG)
186    parser = argparse.ArgumentParser(
187            description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
188    parser.add_argument('-s', '--subject', type=str, dest='subject',
189                        required=True, help='Subject of the mail')
190    parser.add_argument('-p', type=float, dest='probability',
191                        required=False, default=0,
192                        help='(optional) per-addressee probability '
193                             'with which to send email. If not specified '
194                             'all addressees will receive message.')
195    parser.add_argument('recipients', nargs='*',
196                        help='Email addresses separated by space.')
197    args = parser.parse_args()
198    if not args.recipients or not args.subject:
199        print 'Requires both recipients and subject.'
200        sys.exit(1)
201
202    message_text = sys.stdin.read()
203
204    if args.probability:
205        recipients = []
206        for r in args.recipients:
207            if random.random() < args.probability:
208                recipients.append(r)
209        if recipients:
210            print 'Randomly selected recipients %s' % recipients
211        else:
212            print 'Random filtering removed all recipients. Sending nothing.'
213            sys.exit(0)
214    else:
215        recipients = args.recipients
216
217
218    with site_utils.SetupTsMonGlobalState('gmail_lib', short_lived=True):
219        send_email(','.join(recipients), args.subject , message_text)
220