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 import random 26 from email.mime.text import MIMEText 27 28 import common 29 from autotest_lib.client.bin import utils 30 from autotest_lib.client.common_lib import global_config 31 from autotest_lib.server import site_utils 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 # 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 43 from autotest_lib.server import utils as server_utils 44 from chromite.lib import retry_util 45 46 try: 47 from chromite.lib import metrics 48 except ImportError: 49 metrics = utils.metrics_mock 50 51 52 DEFAULT_CREDS_FILE = global_config.global_config.get_config_value( 53 'NOTIFICATIONS', 'gmail_api_credentials', default=None) 54 RETRY_DELAY = 5 55 RETRY_BACKOFF_FACTOR = 1.5 56 MAX_RETRY = 10 57 RETRIABLE_MSGS = [ 58 # User-rate limit exceeded 59 r'HttpError 429',] 60 61 class GmailApiException(Exception): 62 """Exception raised in accessing Gmail API.""" 63 64 65 class 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 (at] gmail.com,xyz (at] 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 92 class 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 130 def 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 184 if __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
