# Copyright (C) 2008 The Android Open Source Project
#
# 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.

"""Common SSH management logic."""

import functools
import os
import re
import signal
import subprocess
import sys
import tempfile
try:
  import threading as _threading
except ImportError:
  import dummy_threading as _threading
import time

import platform_utils
from repo_trace import Trace


_ssh_proxy_path = None
_ssh_sock_path = None
_ssh_clients = []


def _run_ssh_version():
  """run ssh -V to display the version number"""
  return subprocess.check_output(['ssh', '-V'], stderr=subprocess.STDOUT).decode()


def _parse_ssh_version(ver_str=None):
  """parse a ssh version string into a tuple"""
  if ver_str is None:
    ver_str = _run_ssh_version()
  m = re.match(r'^OpenSSH_([0-9.]+)(p[0-9]+)?\s', ver_str)
  if m:
    return tuple(int(x) for x in m.group(1).split('.'))
  else:
    return ()


@functools.lru_cache(maxsize=None)
def version():
  """return ssh version as a tuple"""
  try:
    return _parse_ssh_version()
  except subprocess.CalledProcessError:
    print('fatal: unable to detect ssh version', file=sys.stderr)
    sys.exit(1)


def proxy():
  global _ssh_proxy_path
  if _ssh_proxy_path is None:
    _ssh_proxy_path = os.path.join(
        os.path.dirname(__file__),
        'git_ssh')
  return _ssh_proxy_path


def add_client(p):
  _ssh_clients.append(p)


def remove_client(p):
  try:
    _ssh_clients.remove(p)
  except ValueError:
    pass


def _terminate_clients():
  global _ssh_clients
  for p in _ssh_clients:
    try:
      os.kill(p.pid, signal.SIGTERM)
      p.wait()
    except OSError:
      pass
  _ssh_clients = []


_master_processes = []
_master_keys = set()
_ssh_master = True
_master_keys_lock = None


def init():
  """Should be called once at the start of repo to init ssh master handling.

  At the moment, all we do is to create our lock.
  """
  global _master_keys_lock
  assert _master_keys_lock is None, "Should only call init once"
  _master_keys_lock = _threading.Lock()


def _open_ssh(host, port=None):
  global _ssh_master

  # Bail before grabbing the lock if we already know that we aren't going to
  # try creating new masters below.
  if sys.platform in ('win32', 'cygwin'):
    return False

  # Acquire the lock.  This is needed to prevent opening multiple masters for
  # the same host when we're running "repo sync -jN" (for N > 1) _and_ the
  # manifest <remote fetch="ssh://xyz"> specifies a different host from the
  # one that was passed to repo init.
  _master_keys_lock.acquire()
  try:

    # Check to see whether we already think that the master is running; if we
    # think it's already running, return right away.
    if port is not None:
      key = '%s:%s' % (host, port)
    else:
      key = host

    if key in _master_keys:
      return True

    if not _ssh_master or 'GIT_SSH' in os.environ:
      # Failed earlier, so don't retry.
      return False

    # We will make two calls to ssh; this is the common part of both calls.
    command_base = ['ssh',
                    '-o', 'ControlPath %s' % sock(),
                    host]
    if port is not None:
      command_base[1:1] = ['-p', str(port)]

    # Since the key wasn't in _master_keys, we think that master isn't running.
    # ...but before actually starting a master, we'll double-check.  This can
    # be important because we can't tell that that 'git@myhost.com' is the same
    # as 'myhost.com' where "User git" is setup in the user's ~/.ssh/config file.
    check_command = command_base + ['-O', 'check']
    try:
      Trace(': %s', ' '.join(check_command))
      check_process = subprocess.Popen(check_command,
                                       stdout=subprocess.PIPE,
                                       stderr=subprocess.PIPE)
      check_process.communicate()  # read output, but ignore it...
      isnt_running = check_process.wait()

      if not isnt_running:
        # Our double-check found that the master _was_ infact running.  Add to
        # the list of keys.
        _master_keys.add(key)
        return True
    except Exception:
      # Ignore excpetions.  We we will fall back to the normal command and print
      # to the log there.
      pass

    command = command_base[:1] + ['-M', '-N'] + command_base[1:]
    try:
      Trace(': %s', ' '.join(command))
      p = subprocess.Popen(command)
    except Exception as e:
      _ssh_master = False
      print('\nwarn: cannot enable ssh control master for %s:%s\n%s'
            % (host, port, str(e)), file=sys.stderr)
      return False

    time.sleep(1)
    ssh_died = (p.poll() is not None)
    if ssh_died:
      return False

    _master_processes.append(p)
    _master_keys.add(key)
    return True
  finally:
    _master_keys_lock.release()


def close():
  global _master_keys_lock

  _terminate_clients()

  for p in _master_processes:
    try:
      os.kill(p.pid, signal.SIGTERM)
      p.wait()
    except OSError:
      pass
  del _master_processes[:]
  _master_keys.clear()

  d = sock(create=False)
  if d:
    try:
      platform_utils.rmdir(os.path.dirname(d))
    except OSError:
      pass

  # We're done with the lock, so we can delete it.
  _master_keys_lock = None


URI_SCP = re.compile(r'^([^@:]*@?[^:/]{1,}):')
URI_ALL = re.compile(r'^([a-z][a-z+-]*)://([^@/]*@?[^/]*)/')


def preconnect(url):
  m = URI_ALL.match(url)
  if m:
    scheme = m.group(1)
    host = m.group(2)
    if ':' in host:
      host, port = host.split(':')
    else:
      port = None
    if scheme in ('ssh', 'git+ssh', 'ssh+git'):
      return _open_ssh(host, port)
    return False

  m = URI_SCP.match(url)
  if m:
    host = m.group(1)
    return _open_ssh(host)

  return False

def sock(create=True):
  global _ssh_sock_path
  if _ssh_sock_path is None:
    if not create:
      return None
    tmp_dir = '/tmp'
    if not os.path.exists(tmp_dir):
      tmp_dir = tempfile.gettempdir()
    if version() < (6, 7):
      tokens = '%r@%h:%p'
    else:
      tokens = '%C'  # hash of %l%h%p%r
    _ssh_sock_path = os.path.join(
        tempfile.mkdtemp('', 'ssh-', tmp_dir),
        'master-' + tokens)
  return _ssh_sock_path