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