# Copyright 2017 The LUCI Authors. All rights reserved.
# Use of this source code is governed under the Apache License, Version 2.0
# that can be found in the LICENSE file.

"""Common gRPC implementation for Swarming and Isolate"""

import logging
import os
import re
import time
import urlparse
from utils import net

# gRPC may not be installed on the worker machine. This is fine, as long as
# the bot doesn't attempt to use gRPC (checked in IsolateServerGrpc.__init__).
# Full external requirements are: grpcio, certifi.
try:
  import grpc
  from google import auth as google_auth
  from google.auth.transport import grpc as google_auth_transport_grpc
  from google.auth.transport import requests as google_auth_transport_requests
except ImportError as err:
  grpc = None


# If gRPC was successfully imported, try to import certifi as well.  This is not
# actually used anywhere in this module, but if certifi is missing,
# google.auth.transport will fail (see
# https://stackoverflow.com/questions/24973326). So checking it here allows us
# to print out a somewhat-sane error message.
certifi = None
if grpc is not None:
  try:
    import certifi
  except ImportError:
    # Will print out error messages later (ie when we have a logger)
    pass


# How many times to retry a gRPC call
MAX_GRPC_ATTEMPTS = 30


# Longest time to sleep between gRPC calls
MAX_GRPC_SLEEP = 10.


# Start the timeout at three minutes.
GRPC_TIMEOUT_SEC = 3 * 60


def available():
  """Returns true if gRPC can be used on this host."""
  return grpc != None


class Proxy(object):
  """Represents a gRPC proxy.

  If the proxy begins with 'https', the returned channel will be secure and
  authorized using default application credentials - see
  https://developers.google.com/identity/protocols/application-default-credentials.
  Currently, we're using Cloud Container Builder scopes for testing; this may
  change in the future to allow different scopes to be passed in for different
  channels.

  To use the returned channel to call methods directly, say:

    proxy = grpc_proxy.Proxy('https://grpc.luci.org/resource/prefix',
                             myapi_pb2.MyApiStub)

  To make a unary call with retries (recommended):

    proto_output = proxy.call_unary('MyMethod', proto_input)

  To make a unary call without retries, or to pass in a client side stream
  (proto_input can be an iterator here):

    proto_output = proxy.call_no_retries('MyMethod', proto_input)

  You can also call the stub directly (not recommended, since no errors will be
  caught or logged):

    proto_output = proxy.stub.MyMethod(proto_input)

  To make a call to a server-side streaming call (these are not retried):

    for response in proxy.get_stream('MyStreaminingMethod', proto_input):
      <process response>

  To retrieve the prefix:

    prefix = proxy.prefix # returns "prefix/for/resource/names"

  All exceptions are logged using "logging.warning."
  """

  def __init__(self, proxy, stub_class):
    self._verbose = os.environ.get('LUCI_GRPC_PROXY_VERBOSE')
    if self._verbose:
      logging.info('Enabled verbose mode for %s with stub %s',
                   proxy, stub_class.__name__)
    # NB: everything in url is unicode; convert to strings where
    # needed.
    url = urlparse.urlparse(proxy)
    if self._verbose:
      logging.info('Parsed URL for proxy is %r', url)
    if url.scheme == 'http':
      self._secure = False
    elif url.scheme == 'https':
      self._secure = True
    else:
      raise ValueError('gRPC proxy %s must use http[s], not %s' % (
          proxy, url.scheme))
    if url.netloc == '':
      raise ValueError('gRPC proxy is missing hostname: %s' % proxy)
    self._host = url.netloc
    self._prefix = url.path
    if self._prefix.endswith('/'):
      self._prefix = self._prefix[:-1]
    if self._prefix.startswith('/'):
      self._prefix = self._prefix[1:]
    if url.params != '' or url.fragment != '':
      raise ValueError('gRPC proxy may not contain params or fragments: %s' %
                       proxy)
    self._debug_info = ['full proxy name: ' + proxy]
    self._channel = self._create_channel()
    self._stub = stub_class(self._channel)
    logging.info('%s: initialized', self.name)
    if self._verbose:
      self._dump_proxy_info()

  @property
  def prefix(self):
    return self._prefix

  @property
  def channel(self):
    return self._channel

  @property
  def stub(self):
    return self._stub

  @property
  def name(self):
    security = 'insecure'
    if self._secure:
      security = 'secure'
    return 'gRPC %s proxy %s/%s' % (
        security, self._host, self._stub.__class__.__name__)

  def call_unary(self, name, request):
    """Calls a method, waiting if the service is not available.

    Usage: proto_output = proxy.call_unary('MyMethod', proto_input)
    """
    for attempt in range(1, MAX_GRPC_ATTEMPTS+1):
      try:
        return self.call_no_retries(name, request)
      except grpc.RpcError as g:
        if g.code() is not grpc.StatusCode.UNAVAILABLE:
          raise
        logging.warning('%s: call_grpc - proxy is unavailable (attempt %d/%d)',
                        self.name, attempt, MAX_GRPC_ATTEMPTS)
        # Save the error in case we need to return it
        grpc_error = g
        time.sleep(net.calculate_sleep_before_retry(attempt, MAX_GRPC_SLEEP))
    # If we get here, it must be because we got (and saved) an error
    assert grpc_error is not None
    raise grpc_error

  def get_stream(self, name, request):
    """Calls a server-side streaming method, returning an iterator.

    Usage: for resp in proxy.get_stream('MyMethod', proto_input'):
    """
    stream = self.call_no_retries(name, request)
    while True:
      # The lambda "next(stream, 1)" will return a protobuf on success, or the
      # integer 1 if the stream has ended. This allows us to avoid attempting
      # to catch StopIteration, which gets logged by _wrap_grpc_operation.
      response = self._wrap_grpc_operation(name + ' pull from stream',
                                           lambda: next(stream, 1))
      if isinstance(response, int):
        # Iteration is finished
        return
      yield response

  def call_no_retries(self, name, request):
    """Calls a method without any retries.

    Recommended for client-side streaming or nonidempotent unary calls.
    """
    method = getattr(self._stub, name)
    if method is None:
      raise NameError('%s: "%s" is not a valid method name', self.name, name)
    return self._wrap_grpc_operation(
        name, lambda: method(request, timeout=GRPC_TIMEOUT_SEC))

  def _wrap_grpc_operation(self, name, fn):
    """Wraps a gRPC operation (call or iterator increment) for logging."""
    if self._verbose:
      logging.info('%s/%s - starting gRPC operation', self.name, name)
    try:
      return fn()
    except grpc.RpcError as g:
      logging.warning('\n\nFailure in %s/%s: gRPC error %s', self.name, name, g)
      self._dump_proxy_info()
      raise g
    except Exception as e:
      logging.warning('\n\nFailure in %s/%s: exception %s', self.name, name, e)
      self._dump_proxy_info()
      raise e

  def _dump_proxy_info(self):
    logging.warning('DETAILED PROXY INFO')
    logging.warning('prefix = %s', self.prefix)
    logging.warning('debug info:\n\t%s\n\n',
                    '\n\t'.join(self._debug_info))

  def _create_channel(self):
    # Make sure grpc was successfully imported
    assert available()

    if not self._secure:
      return grpc.insecure_channel(self._host)

    # Authenticate the host.
    #
    # You're allowed to override the root certs and server if necessary. For
    # example, if you're running your proxy on localhost, you'll need to set
    # GRPC_PROXY_TLS_ROOTS to the "roots.crt" file specifying the certificate
    # for the root CA that your localhost server has used to certify itself, and
    # the GRPC_PROXY_TLS_OVERRIDE to the name that your server is using to
    # identify itself. For example, the ROOTS env var might be
    # "/path/to/roots.crt" while the OVERRIDE env var might be "test_server," if
    # this is what's used by the server you're running.
    #
    # If you're connecting to a real server with real SSL, none of this should
    # be used.
    if not certifi:
      self._debug_info.append('CERTIFI IS NOT PRESENT;' +
                              ' gRPC HTTPS CONNECTIONS MAY FAIL')
    root_certs = None
    roots = os.environ.get('LUCI_GRPC_PROXY_TLS_ROOTS')
    if roots:
      self._debug_info.append('Overridden root CA: %s' % roots)
      with open(roots) as f:
        root_certs = f.read()
    else:
      self._debug_info.append('Using default root CAs from certifi')
    overd = os.environ.get('LUCI_GRPC_PROXY_TLS_OVERRIDE')
    options = ()
    if overd:
      options=(('grpc.ssl_target_name_override', overd),)
    ssl_creds = grpc.ssl_channel_credentials(root_certificates=root_certs)

    # Authenticate the user.
    scopes = ('https://www.googleapis.com/auth/cloud-source-tools',)
    self._debug_info.append('Scopes are: %r' % scopes)
    user_creds, _ = google_auth.default(scopes=scopes)

    # Create the channel.
    request = google_auth_transport_requests.Request()
    self._debug_info.append('Options are: %r' % options)
    return google_auth_transport_grpc.secure_authorized_channel(
        user_creds, request, self._host, ssl_creds, options=options)