lib/alerts.py - third_party/chromite - Git at Google

 # Copyright 2012 The ChromiumOS Authors
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.

 """Chromite email utility functions."""

 import base64
 import collections
 from email.mime.application import MIMEApplication
 from email.mime.multipart import MIMEMultipart
 from email.mime.text import MIMEText
 import gzip
 import http.client
 import io
 import json
 import logging
 import os
 import smtplib
 import socket
 import sys
 import traceback

 from chromite.third_party import httplib2
 from chromite.third_party.googleapiclient import errors as apiclient_errors
 from chromite.third_party.googleapiclient.discovery import (
     build as apiclient_build,
 )
 from chromite.third_party.oauth2client import client
 from chromite.third_party.oauth2client import file as oauth_client_fileio

 from chromite.lib import retry_util


 class MailServer:
     """Base class for servers."""

     def Send(self, message) -> None:
         """Send the message.

         Override by sub-classes.

         Args:
             message: A MIMEMultipart() object containing the message body.

         Returns:
             True if the email was sent, else False.
         """
         raise NotImplementedError("Should be implemented by sub-classes.")


 class AuthenticationError(Exception):
     """Error raised when authenticating via oauth2."""


 # Represent token in oauth2 token file.
 RefreshToken = collections.namedtuple(
     "RefreshToken",
     (
         "client_id",
         "client_secret",
         "refresh_token",
     ),
 )


 def ReadRefreshTokenJson(path):
     """Returns RefreshToken by reading it from the JSON file.

     Args:
         path: Path to the json file that contains the credential tokens.

     Returns:
         A RefreshToken object.

     Raises:
         AuthenticationError if failed to read from json file.
     """
     try:
         with open(path, "rb") as f:
             data = json.load(f)
             return RefreshToken(
                 client_id=str(data["client_id"]),
                 client_secret=str(data["client_secret"]),
                 refresh_token=str(data["refresh_token"]),
             )
     except (IOError, ValueError) as e:
         raise AuthenticationError(
             "Failed to read refresh token from %s: %s" % (path, e)
         )
     except KeyError as e:
         raise AuthenticationError(
             "Failed to read refresh token from %s: missing key %s" % (path, e)
         )


 class GmailServer(MailServer):
     """Gmail server."""

     TOKEN_URI = "https://accounts.google.com/o/oauth2/token"

     def __init__(self, token_cache_file, token_json_file=None) -> None:
         """Initialize GmailServer.

         If token_cache_file contains valid credentials, it will be used.
         If not or the file doesn't exist, will try to load tokens
         from token_json_file. The loaded credentials will be stored to the
         cache file.

         Args:
             token_cache_file: Absolute path to gmail credentials cache file.
             token_json_file: Absolute path to a json file that contains
                 refresh token for gmail.
         """
         self._token_cache_file = token_cache_file
         self._token_json_file = token_json_file

     def _GetCachedCredentials(self):
         """Get credentials from cached file or json file.

         Returns:
             OAuth2Credentials object.

         Raises:
             AuthenticationError on failure to read json file.
         """
         storage = oauth_client_fileio.Storage(self._token_cache_file)
         # Try loading credentials from existing token cache file.
         if os.path.isfile(self._token_cache_file):
             credentials = storage.get()
             if credentials and not credentials.invalid:
                 return credentials

         if self._token_json_file is None:
             raise AuthenticationError("Gmail token file path is not provided.")

         # Create new credentials if cache file doesn't exist or not valid.
         refresh_token_json = ReadRefreshTokenJson(self._token_json_file)
         credentials = client.OAuth2Credentials(
             access_token=None,
             client_id=refresh_token_json.client_id,
             client_secret=refresh_token_json.client_secret,
             refresh_token=refresh_token_json.refresh_token,
             token_expiry=None,
             token_uri=self.TOKEN_URI,
             user_agent=None,
             revoke_uri=None,
         )
         credentials.set_store(storage)
         storage.put(credentials)
         return credentials

     def Send(self, message):
         """Send an e-mail via Gmail API.

         Args:
             message: A MIMEMultipart() object containing the message body.

         Returns:
             True if the email was sent, else False.
         """
         try:
             credentials = self._GetCachedCredentials()
         except AuthenticationError as e:
             logging.warning("Could not get gmail credentials: %s", e)
             return False

         httpcreds = credentials.authorize(httplib2.Http())
         service = apiclient_build("gmail", "v1", http=httpcreds)
         try:
             # 'me' represents the default authorized user.
             payload = {
                 "raw": base64.urlsafe_b64encode(
                     message.as_string().encode("utf-8")
                 ).decode("utf-8"),
             }
             service.users().messages().send(userId="me", body=payload).execute()
             return True
         except (
             apiclient_errors.HttpError,
             http.client.HTTPException,
             client.Error,
         ) as error:
             logging.warning("Could not send email: %s", error)
             return False


 class SmtpServer(MailServer):
     """Smtp server."""

     # Note: When importing this module from cbuildbot code that will run on
     # a builder in the golo, set this to constants.GOLO_SMTP_SERVER
     DEFAULT_SERVER = "localhost"
     # Retry parameters for the actual smtp connection.
     SMTP_RETRY_COUNT = 3
     SMTP_RETRY_DELAY = 30

     def __init__(self, smtp_server=None) -> None:
         """Initialize SmtpServer.

         Args:
             smtp_server: The server with which to send the message.
         """
         self._smtp_server = smtp_server or self.DEFAULT_SERVER

     def Send(self, message):
         """Send an email via SMTP

         If we get a socket error (e.g. the SMTP server is not listening or
         timesout), we will retry a few times.  All socket errors will be
         caught here.

         Args:
             message: A MIMEMultipart() object containing the message body.

         Returns:
             True if the email was sent, else False.
         """

         def _Send() -> None:
             smtp_client = smtplib.SMTP(self._smtp_server)
             recipients = [s.strip() for s in message["To"].split(",")]
             smtp_client.sendmail(
                 message["From"], recipients, message.as_string()
             )
             smtp_client.quit()

         try:
             retry_util.RetryException(
                 socket.error,
                 self.SMTP_RETRY_COUNT,
                 _Send,
                 sleep=self.SMTP_RETRY_DELAY,
             )
             return True
         except socket.error as e:
             logging.warning(
                 "Could not send e-mail from %s to %s via %r: %s",
                 message["From"],
                 message["To"],
                 self._smtp_server,
                 e,
             )
             return False


 def CreateEmail(
     subject, recipients, message="", attachment=None, extra_fields=None
 ):
     """Create an email message object.

     Args:
         subject: E-mail subject.
         recipients: List of e-mail recipients.
         message: Message to put in the e-mail body.
         attachment: text to attach.
         extra_fields: A dictionary of additional message header
             fields to be added to the message. Custom field names should begin
             with the prefix 'X-'.

     Returns:
         A MIMEMultipart object, or None if recipients is empty.
     """
     # Ignore if the list of recipients is empty.
     if not recipients:
         logging.warning("Could not create email: recipient list is emtpy.")
         return None

     extra_fields = extra_fields or {}
     sender = socket.getfqdn()
     msg = MIMEMultipart()
     for key, val in extra_fields.items():
         msg[key] = val
     msg["From"] = sender
     msg["Subject"] = subject
     msg["To"] = ", ".join(recipients)

     msg.attach(MIMEText(message))
     if attachment:
         if isinstance(attachment, str):
             attachment = attachment.encode()
         s = io.BytesIO()
         with gzip.GzipFile(fileobj=s, mode="wb") as f:
             f.write(attachment)
         part = MIMEApplication(s.getvalue(), _subtype="x-gzip")
         s.close()
         part.add_header(
             "Content-Disposition", "attachment", filename="logs.txt.gz"
         )
         msg.attach(part)

     return msg


 def SendEmail(
     subject,
     recipients,
     server=None,
     message="",
     attachment=None,
     extra_fields=None,
 ) -> None:
     """Send an e-mail job notification with the given message in the body.

     Args:
         subject: E-mail subject.
         recipients: List of e-mail recipients.
         server: A MailServer instance. Default to local SmtpServer.
         message: Message to put in the e-mail body.
         attachment: Text to attach.
         extra_fields: A dictionary of additional message header fields
             to be added to the message. Custom field names should begin
             with the prefix 'X-'.
     """
     if server is None:
         server = SmtpServer()
     msg = CreateEmail(subject, recipients, message, attachment, extra_fields)
     if not msg:
         return
     server.Send(msg)


 def SendEmailLog(
     subject,
     recipients,
     server=None,
     message="",
     inc_trace=True,
     log=None,
     extra_fields=None,
 ) -> None:
     """Send an e-mail with a stack trace and log snippets.

     Args:
         subject: E-mail subject.
         recipients: list of e-mail recipients.
         server: A MailServer instance. Default to local SmtpServer.
         inc_trace: Append a backtrace of the current stack.
         message: Message to put at the top of the e-mail body.
         log: List of lines (log data) to include in the notice.
         extra_fields: A dictionary of additional message header
             fields to be added to the message. Custom fields names should begin
             with the prefix 'X-'.
     """
     if server is None:
         server = SmtpServer()
     if not message:
         message = subject
     message = message[:]

     if inc_trace:
         if sys.exc_info() != (None, None, None):
             trace = traceback.format_exc()
             message += "\n\n" + trace

     attachment = None
     if log:
         message += (
             "\n\n"
             + "***************************\n"
             + "Last log messages:\n"
             + "***************************\n"
             + "".join(log[-50:])
         )
         attachment = "".join(log)

     SendEmail(
         subject,
         recipients,
         server,
         message=message,
         attachment=attachment,
         extra_fields=extra_fields,
     )
	# Copyright 2012 The ChromiumOS Authors
	# Use of this source code is governed by a BSD-style license that can be
	# found in the LICENSE file.

	"""Chromite email utility functions."""

	import base64
	import collections
	from email.mime.application import MIMEApplication
	from email.mime.multipart import MIMEMultipart
	from email.mime.text import MIMEText
	import gzip
	import http.client
	import io
	import json
	import logging
	import os
	import smtplib
	import socket
	import sys
	import traceback

	from chromite.third_party import httplib2
	from chromite.third_party.googleapiclient import errors as apiclient_errors
	from chromite.third_party.googleapiclient.discovery import (
	build as apiclient_build,
	)
	from chromite.third_party.oauth2client import client
	from chromite.third_party.oauth2client import file as oauth_client_fileio

	from chromite.lib import retry_util


	class MailServer:
	"""Base class for servers."""

	def Send(self, message) -> None:
	"""Send the message.

	Override by sub-classes.

	Args:
	message: A MIMEMultipart() object containing the message body.

	Returns:
	True if the email was sent, else False.
	"""
	raise NotImplementedError("Should be implemented by sub-classes.")


	class AuthenticationError(Exception):
	"""Error raised when authenticating via oauth2."""


	# Represent token in oauth2 token file.
	RefreshToken = collections.namedtuple(
	"RefreshToken",
	(
	"client_id",
	"client_secret",
	"refresh_token",
	),
	)


	def ReadRefreshTokenJson(path):
	"""Returns RefreshToken by reading it from the JSON file.

	Args:
	path: Path to the json file that contains the credential tokens.

	Returns:
	A RefreshToken object.

	Raises:
	AuthenticationError if failed to read from json file.
	"""
	try:
	with open(path, "rb") as f:
	data = json.load(f)
	return RefreshToken(
	client_id=str(data["client_id"]),
	client_secret=str(data["client_secret"]),
	refresh_token=str(data["refresh_token"]),
	)
	except (IOError, ValueError) as e:
	raise AuthenticationError(
	"Failed to read refresh token from %s: %s" % (path, e)
	)
	except KeyError as e:
	raise AuthenticationError(
	"Failed to read refresh token from %s: missing key %s" % (path, e)
	)


	class GmailServer(MailServer):
	"""Gmail server."""

	TOKEN_URI = "https://accounts.google.com/o/oauth2/token"

	def __init__(self, token_cache_file, token_json_file=None) -> None:
	"""Initialize GmailServer.

	If token_cache_file contains valid credentials, it will be used.
	If not or the file doesn't exist, will try to load tokens
	from token_json_file. The loaded credentials will be stored to the
	cache file.

	Args:
	token_cache_file: Absolute path to gmail credentials cache file.
	token_json_file: Absolute path to a json file that contains
	refresh token for gmail.
	"""
	self._token_cache_file = token_cache_file
	self._token_json_file = token_json_file

	def _GetCachedCredentials(self):
	"""Get credentials from cached file or json file.

	Returns:
	OAuth2Credentials object.

	Raises:
	AuthenticationError on failure to read json file.
	"""
	storage = oauth_client_fileio.Storage(self._token_cache_file)
	# Try loading credentials from existing token cache file.
	if os.path.isfile(self._token_cache_file):
	credentials = storage.get()
	if credentials and not credentials.invalid:
	return credentials

	if self._token_json_file is None:
	raise AuthenticationError("Gmail token file path is not provided.")

	# Create new credentials if cache file doesn't exist or not valid.
	refresh_token_json = ReadRefreshTokenJson(self._token_json_file)
	credentials = client.OAuth2Credentials(
	access_token=None,
	client_id=refresh_token_json.client_id,
	client_secret=refresh_token_json.client_secret,
	refresh_token=refresh_token_json.refresh_token,
	token_expiry=None,
	token_uri=self.TOKEN_URI,
	user_agent=None,
	revoke_uri=None,
	)
	credentials.set_store(storage)
	storage.put(credentials)
	return credentials

	def Send(self, message):
	"""Send an e-mail via Gmail API.

	Args:
	message: A MIMEMultipart() object containing the message body.

	Returns:
	True if the email was sent, else False.
	"""
	try:
	credentials = self._GetCachedCredentials()
	except AuthenticationError as e:
	logging.warning("Could not get gmail credentials: %s", e)
	return False

	httpcreds = credentials.authorize(httplib2.Http())
	service = apiclient_build("gmail", "v1", http=httpcreds)
	try:
	# 'me' represents the default authorized user.
	payload = {
	"raw": base64.urlsafe_b64encode(
	message.as_string().encode("utf-8")
	).decode("utf-8"),
	}
	service.users().messages().send(userId="me", body=payload).execute()
	return True
	except (
	apiclient_errors.HttpError,
	http.client.HTTPException,
	client.Error,
	) as error:
	logging.warning("Could not send email: %s", error)
	return False


	class SmtpServer(MailServer):
	"""Smtp server."""

	# Note: When importing this module from cbuildbot code that will run on
	# a builder in the golo, set this to constants.GOLO_SMTP_SERVER
	DEFAULT_SERVER = "localhost"
	# Retry parameters for the actual smtp connection.
	SMTP_RETRY_COUNT = 3
	SMTP_RETRY_DELAY = 30

	def __init__(self, smtp_server=None) -> None:
	"""Initialize SmtpServer.

	Args:
	smtp_server: The server with which to send the message.
	"""
	self._smtp_server = smtp_server or self.DEFAULT_SERVER

	def Send(self, message):
	"""Send an email via SMTP

	If we get a socket error (e.g. the SMTP server is not listening or
	timesout), we will retry a few times. All socket errors will be
	caught here.

	Args:
	message: A MIMEMultipart() object containing the message body.

	Returns:
	True if the email was sent, else False.
	"""

	def _Send() -> None:
	smtp_client = smtplib.SMTP(self._smtp_server)
	recipients = [s.strip() for s in message["To"].split(",")]
	smtp_client.sendmail(
	message["From"], recipients, message.as_string()
	)
	smtp_client.quit()

	try:
	retry_util.RetryException(
	socket.error,
	self.SMTP_RETRY_COUNT,
	_Send,
	sleep=self.SMTP_RETRY_DELAY,
	)
	return True
	except socket.error as e:
	logging.warning(
	"Could not send e-mail from %s to %s via %r: %s",
	message["From"],
	message["To"],
	self._smtp_server,
	e,
	)
	return False


	def CreateEmail(
	subject, recipients, message="", attachment=None, extra_fields=None
	):
	"""Create an email message object.

	Args:
	subject: E-mail subject.
	recipients: List of e-mail recipients.
	message: Message to put in the e-mail body.
	attachment: text to attach.
	extra_fields: A dictionary of additional message header
	fields to be added to the message. Custom field names should begin
	with the prefix 'X-'.

	Returns:
	A MIMEMultipart object, or None if recipients is empty.
	"""
	# Ignore if the list of recipients is empty.
	if not recipients:
	logging.warning("Could not create email: recipient list is emtpy.")
	return None

	extra_fields = extra_fields or {}
	sender = socket.getfqdn()
	msg = MIMEMultipart()
	for key, val in extra_fields.items():
	msg[key] = val
	msg["From"] = sender
	msg["Subject"] = subject
	msg["To"] = ", ".join(recipients)

	msg.attach(MIMEText(message))
	if attachment:
	if isinstance(attachment, str):
	attachment = attachment.encode()
	s = io.BytesIO()
	with gzip.GzipFile(fileobj=s, mode="wb") as f:
	f.write(attachment)
	part = MIMEApplication(s.getvalue(), _subtype="x-gzip")
	s.close()
	part.add_header(
	"Content-Disposition", "attachment", filename="logs.txt.gz"
	)
	msg.attach(part)

	return msg


	def SendEmail(
	subject,
	recipients,
	server=None,
	message="",
	attachment=None,
	extra_fields=None,
	) -> None:
	"""Send an e-mail job notification with the given message in the body.

	Args:
	subject: E-mail subject.
	recipients: List of e-mail recipients.
	server: A MailServer instance. Default to local SmtpServer.
	message: Message to put in the e-mail body.
	attachment: Text to attach.
	extra_fields: A dictionary of additional message header fields
	to be added to the message. Custom field names should begin
	with the prefix 'X-'.
	"""
	if server is None:
	server = SmtpServer()
	msg = CreateEmail(subject, recipients, message, attachment, extra_fields)
	if not msg:
	return
	server.Send(msg)


	def SendEmailLog(
	subject,
	recipients,
	server=None,
	message="",
	inc_trace=True,
	log=None,
	extra_fields=None,
	) -> None:
	"""Send an e-mail with a stack trace and log snippets.

	Args:
	subject: E-mail subject.
	recipients: list of e-mail recipients.
	server: A MailServer instance. Default to local SmtpServer.
	inc_trace: Append a backtrace of the current stack.
	message: Message to put at the top of the e-mail body.
	log: List of lines (log data) to include in the notice.
	extra_fields: A dictionary of additional message header
	fields to be added to the message. Custom fields names should begin
	with the prefix 'X-'.
	"""
	if server is None:
	server = SmtpServer()
	if not message:
	message = subject
	message = message[:]

	if inc_trace:
	if sys.exc_info() != (None, None, None):
	trace = traceback.format_exc()
	message += "\n\n" + trace

	attachment = None
	if log:
	message += (
	"\n\n"
	+ "***************************\n"
	+ "Last log messages:\n"
	+ "***************************\n"
	+ "".join(log[-50:])
	)
	attachment = "".join(log)

	SendEmail(
	subject,
	recipients,
	server,
	message=message,
	attachment=attachment,
	extra_fields=extra_fields,
	)