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