# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
# Lint as: python3
"""Provides authentication support for TensorBoardUploader."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import errno
import json
import os
import sys
import webbrowser

import google_auth_oauthlib.flow
import grpc
import google.auth
import google.auth.transport.requests
import google.oauth2.credentials

from tensorboard.uploader import util
from tensorboard.util import tb_logging


logger = tb_logging.get_logger()


# OAuth2 scopes used for OpenID Connect:
# https://developers.google.com/identity/protocols/OpenIDConnect#scope-param
OPENID_CONNECT_SCOPES = (
    "openid",
    "https://www.googleapis.com/auth/userinfo.email",
)


# The client "secret" is public by design for installed apps. See
# https://developers.google.com/identity/protocols/OAuth2?csw=1#installed
OAUTH_CLIENT_CONFIG = u"""
{
  "installed": {
    "client_id": "373649185512-8v619h5kft38l4456nm2dj4ubeqsrvh6.apps.googleusercontent.com",
    "project_id": "hosted-tensorboard-prod",
    "auth_uri": "https://accounts.google.com/o/oauth2/auth",
    "token_uri": "https://oauth2.googleapis.com/token",
    "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
    "client_secret": "pOyAuU2yq2arsM98Bw5hwYtr",
    "redirect_uris": [
      "urn:ietf:wg:oauth:2.0:oob",
      "http://localhost"
    ]
  }
}
"""


# Components of the relative path (within the user settings directory) at which
# to store TensorBoard uploader credentials.
TENSORBOARD_CREDENTIALS_FILEPATH_PARTS = [
    "tensorboard", "credentials", "uploader-creds.json"]


class CredentialsStore(object):
  """Private file store for a `google.oauth2.credentials.Credentials`."""

  _DEFAULT_CONFIG_DIRECTORY = object()  # Sentinel value.

  def __init__(self, user_config_directory=_DEFAULT_CONFIG_DIRECTORY):
    """Creates a CredentialsStore.

    Args:
      user_config_directory: Optional absolute path to the root directory for
        storing user configs, under which to store the credentials file. If not
        set, defaults to a platform-specific location. If set to None, the
        store is disabled (reads return None; write and clear are no-ops).
    """
    if user_config_directory is CredentialsStore._DEFAULT_CONFIG_DIRECTORY:
      user_config_directory = util.get_user_config_directory()
      if user_config_directory is None:
        logger.warning(
            "Credentials caching disabled - no private config directory found")
    if user_config_directory is None:
      self._credentials_filepath = None
    else:
      self._credentials_filepath = os.path.join(
          user_config_directory, *TENSORBOARD_CREDENTIALS_FILEPATH_PARTS)

  def read_credentials(self):
    """Returns the current `google.oauth2.credentials.Credentials`, or None."""
    if self._credentials_filepath is None:
      return None
    if os.path.exists(self._credentials_filepath):
      return google.oauth2.credentials.Credentials.from_authorized_user_file(
          self._credentials_filepath)
    return None

  def write_credentials(self, credentials):
    """Writes a `google.oauth2.credentials.Credentials` to the store."""
    if not isinstance(credentials, google.oauth2.credentials.Credentials):
      raise TypeError("Cannot write credentials of type %s" % type(credentials))
    if self._credentials_filepath is None:
      return
    # Make the credential file private if not on Windows; on Windows we rely on
    # the default user config settings directory being private since we don't
    # have a straightforward way to make an individual file private.
    private = os.name != "nt"
    util.make_file_with_directories(self._credentials_filepath, private=private)
    data = {
        "refresh_token": credentials.refresh_token,
        "token_uri": credentials.token_uri,
        "client_id": credentials.client_id,
        "client_secret": credentials.client_secret,
        "scopes": credentials.scopes,
        "type": "authorized_user",
    }
    with open(self._credentials_filepath, "w") as f:
      json.dump(data, f)

  def clear(self):
    """Clears the store of any persisted credentials information."""
    if self._credentials_filepath is None:
      return
    try:
      os.remove(self._credentials_filepath)
    except OSError as e:
      if e.errno != errno.ENOENT:
        raise


def build_installed_app_flow(client_config):
  """Returns a `CustomInstalledAppFlow` for the given config.

  Args:
    client_config (Mapping[str, Any]): The client configuration in the Google
        client secrets format.

  Returns:
    CustomInstalledAppFlow: the constructed flow.
  """
  return CustomInstalledAppFlow.from_client_config(
      client_config, scopes=OPENID_CONNECT_SCOPES)


class CustomInstalledAppFlow(google_auth_oauthlib.flow.InstalledAppFlow):
  """Customized version of the Installed App OAuth2 flow."""

  def run(self, force_console=False):
    """Run the flow using a local server if possible, otherwise the console."""
    # TODO(b/141721828): make auto-detection smarter, especially for macOS.
    if not force_console and os.getenv("DISPLAY"):
      try:
        return self.run_local_server(port=0)
      except webbrowser.Error:
        sys.stderr.write("Falling back to console authentication flow...\n")
    return self.run_console()


class IdTokenAuthMetadataPlugin(grpc.AuthMetadataPlugin):
  """A `gRPC AuthMetadataPlugin` that uses ID tokens.

  This works like the existing `google.auth.transport.grpc.AuthMetadataPlugin`
  except that instead of always using access tokens, it preferentially uses the
  `Credentials.id_token` property if available (and logs an error otherwise).

  See http://www.grpc.io/grpc/python/grpc.html#grpc.AuthMetadataPlugin
  """

  def __init__(self, credentials, request):
    """Constructs an IdTokenAuthMetadataPlugin.

    Args:
      credentials (google.auth.credentials.Credentials): The credentials to
        add to requests.
      request (google.auth.transport.Request): A HTTP transport request object
        used to refresh credentials as needed.
    """
    super(IdTokenAuthMetadataPlugin, self).__init__()
    if not isinstance(credentials, google.oauth2.credentials.Credentials):
      raise TypeError(
          "Cannot get ID tokens from credentials type %s" % type(credentials))
    self._credentials = credentials
    self._request = request

  def __call__(self, context, callback):
    """Passes authorization metadata into the given callback.

    Args:
      context (grpc.AuthMetadataContext): The RPC context.
      callback (grpc.AuthMetadataPluginCallback): The callback that will
        be invoked to pass in the authorization metadata.
    """
    headers = {}
    self._credentials.before_request(
        self._request, context.method_name, context.service_url, headers)
    id_token = getattr(self._credentials, "id_token", None)
    if id_token:
      self._credentials.apply(headers, token=id_token)
    else:
      logger.error("Failed to find ID token credentials")
    # Pass headers as key-value pairs to match CallCredentials metadata.
    callback(list(headers.items()), None)


def id_token_call_credentials(credentials):
  """Constructs `grpc.CallCredentials` using `google.auth.Credentials.id_token`.

  Args:
    credentials (google.auth.credentials.Credentials): The credentials to use.

  Returns:
    grpc.CallCredentials: The call credentials.
  """
  request = google.auth.transport.requests.Request()
  return grpc.metadata_call_credentials(
      IdTokenAuthMetadataPlugin(credentials, request))
