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