# -*- coding: utf-8 -*- #
# Copyright 2018 Google LLC. 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.
"""cloud-shell ssh command."""

from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals

import datetime
import threading
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.cloud_shell import util
from googlecloudsdk.command_lib.util.ssh import ssh
from googlecloudsdk.core import log
from googlecloudsdk.core import properties
import six


@base.ReleaseTracks(base.ReleaseTrack.BETA)
class SshBeta(base.Command):
  """Allows you to establish an interactive SSH session with Cloud Shell."""

  detailed_help = {
      'DESCRIPTION':
          """\
        *{command}* lets you remotely log in to Cloud Shell. If your Cloud Shell
        is not currently running, this will cause it to be started before
        establishing the SSH session.
        """,
      'EXAMPLES':
          """\
        To SSH into your Cloud Shell, run:

            $ {command}

        To run a remote command in your Cloud Shell, run:

            $ {command} --command=ls
        """,
  }

  @staticmethod
  def Args(parser):
    util.ParseCommonArgs(parser)
    parser.add_argument(
        '--command',
        help="""\
        A command to run in Cloud Shell.

        Runs the command in Cloud Shell and then exits.
        """)
    parser.add_argument(
        '--dry-run',
        help="""\
        If provided, prints the command that would be run to standard out
        instead of executing it.
        """,
        action='store_true')
    parser.add_argument(
        '--ssh-flag',
        help='Additional flags to be passed to *ssh(1)*.',
        action='append')
    parser.add_argument(
        '--authorize-session',
        help="""\
        If provided, sends OAuth credentials to the current Cloud Shell session
        on behalf of the user. When this completes, the session will be
        authorized to run various Google Cloud command-line tools without
        requiring the user to manually authenticate.
        """,
        action='store_true')

  def Run(self, args):
    if not args.authorize_session:
      log.Print('Automatic authentication with GCP CLI tools in Cloud Shell is '
                'disabled. To enable, please rerun command with '
                '`--authorize-session` flag.')
    command_list = args.command.split(' ') if args.command else ['bash -l']
    project = properties.VALUES.core.project.Get()
    connection_info = util.PrepareV1Environment(args)
    if args.authorize_session:
      util.AuthorizeV1Environment()
    command = ssh.SSHCommand(
        remote=ssh.Remote(host=connection_info.host, user=connection_info.user),
        port=six.text_type(connection_info.port),
        identity_file=connection_info.key,
        remote_command=(['DEVSHELL_PROJECT_ID=' + project] if project else []) +
        command_list,
        extra_flags=args.ssh_flag,
        tty=not args.command,
        options={'StrictHostKeyChecking': 'no'},
    )

    if args.dry_run:
      elems = command.Build(connection_info.ssh_env)
      log.Print(' '.join([six.moves.shlex_quote(elem) for elem in elems]))
    elif args.authorize_session:
      self.done = threading.Event()
      thread = threading.Thread(target=self.Reauthorize, args=())
      thread.daemon = True
      thread.start()
      command.Run(connection_info.ssh_env)
      self.done.set()
    else:
      command.Run(connection_info.ssh_env)

  def Reauthorize(self):
    while not self.done.is_set():
      self.done.wait(
          (util.MIN_CREDS_EXPIRY - datetime.timedelta(minutes=2)).seconds)
      if not self.done.is_set():
        util.AuthorizeV1Environment()


@base.ReleaseTracks(base.ReleaseTrack.ALPHA)
class SshAlpha(base.Command):
  """Allows you to establish an interactive SSH session with Cloud Shell."""

  detailed_help = {
      'DESCRIPTION':
          """\
        *{command}* lets you remotely log in to Cloud Shell. If your Cloud Shell
        is not currently running, this will cause it to be started before
        establishing the SSH session.
        """,
      'EXAMPLES':
          """\
        To SSH into your Cloud Shell, run:

            $ {command}

        To run a remote command in your Cloud Shell, run:

            $ {command} --command=ls
        """,
  }

  @staticmethod
  def Args(parser):
    util.ParseCommonArgs(parser)
    util.AddSshArgFlag(parser)
    parser.add_argument(
        '--command',
        help="""\
        A command to run in Cloud Shell.

        Runs the command in Cloud Shell and then exits.
        """)
    parser.add_argument(
        '--dry-run',
        help="""\
        If provided, prints the command that would be run to standard out
        instead of executing it.
        """,
        action='store_true')
    parser.add_argument(
        '--ssh-flag',
        help='Additional flags to be passed to *ssh(1)*.',
        action='append')
    parser.add_argument(
        '--authorize-session',
        help="""\
        If provided, sends OAuth credentials to the current Cloud Shell session
        on behalf of the user. When this completes, the session will be
        authorized to run various Google Cloud command line tools without
        requiring the user to manually authenticate.
        """,
        action='store_true')

  def Run(self, args):
    if not args.authorize_session:
      log.Print(
          'Automatic authentication with GCP CLI tools in Cloud Shell is '
          'disabled. To enable, please rerun command with '
          '`--authorize-session` flag.')
    command_list = args.command.split(' ') if args.command else ['bash -l']
    project = properties.VALUES.core.project.Get()
    connection_info = util.PrepareEnvironment(args)
    if args.authorize_session:
      util.AuthorizeEnvironment()
    command = ssh.SSHCommand(
        remote=ssh.Remote(host=connection_info.host, user=connection_info.user),
        port=six.text_type(connection_info.port),
        identity_file=connection_info.key,
        remote_command=(['DEVSHELL_PROJECT_ID=' + project]
                        if project else []) + command_list,
        extra_flags=args.ssh_flag,
        tty=not args.command,
        options={'StrictHostKeyChecking': 'no'},
        remainder=args.ssh_args if args.ssh_args else None
    )

    if args.dry_run:
      elems = command.Build(connection_info.ssh_env)
      log.Print(' '.join([six.moves.shlex_quote(elem) for elem in elems]))
    elif args.authorize_session:
      self.done = threading.Event()
      thread = threading.Thread(target=self.Reauthorize, args=())
      thread.daemon = True
      thread.start()
      command.Run(connection_info.ssh_env)
      self.done.set()
    else:
      command.Run(connection_info.ssh_env)

  def Reauthorize(self):
    while not self.done.is_set():
      self.done.wait(
          (util.MIN_CREDS_EXPIRY - datetime.timedelta(minutes=2)).seconds)
      if not self.done.is_set():
        util.AuthorizeEnvironment()