#!/usr/bin/python
#

# Copyright (C) 2010 Google Inc.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301, USA.

"""Tool to setup the SSH configuration on a remote node.

This is needed before we can join the node into the cluster.

"""

# pylint: disable-msg=C0103
# C0103: Invalid name setup-ssh

import getpass
import logging
import paramiko
import os.path
import optparse
import sys

from ganeti import cli
from ganeti import constants
from ganeti import errors
from ganeti import netutils
from ganeti import ssconf
from ganeti import ssh
from ganeti import utils


class RemoteCommandError(errors.GenericError):
  """Exception if remote command was not successful.

  """


class JoinCheckError(errors.GenericError):
  """Exception raised if join check fails.

  """


def _CheckJoin(transport):
  """Checks if a join is safe or dangerous.

  Note: This function relies on the fact, that all
  hosts have the same configuration at compile time of
  Ganeti. So that the constants do not mismatch.

  @param transport: The paramiko transport instance
  @return: True if the join is safe; False otherwise

  """
  sftp = transport.open_sftp_client()
  ss = ssconf.SimpleStore()
  ss_cluster_name_path = ss.KeyToFilename(constants.SS_CLUSTER_NAME)

  cluster_files = {
    ss_cluster_name_path: utils.ReadFile(ss_cluster_name_path),
    constants.NODED_CERT_FILE: utils.ReadFile(constants.NODED_CERT_FILE),
    }

  try:
    remote_noded_file = _ReadSftpFile(sftp, constants.NODED_CERT_FILE)
  except IOError:
    # We can just assume that the file doesn't exist as such error reporting
    # is lacking from paramiko
    #
    # We don't have the noded certificate. As without the cert, the
    # noded is not running, we are on the safe bet to say that this
    # node doesn't belong to a cluster
    return True

  try:
    remote_cluster_name = _ReadSftpFile(sftp, ss_cluster_name_path)
  except IOError:
    # This can indicate that a previous join was not successful
    # So if the noded cert was found and matches we are fine
    return cluster_files[constants.NODED_CERT_FILE] == remote_noded_file

  return (cluster_files[constants.NODED_CERT_FILE] == remote_noded_file and
          cluster_files[ss_cluster_name_path] == remote_cluster_name)


def _RunRemoteCommand(transport, command):
  """Invokes and wait for the command over SSH.

  @param transport: The paramiko transport instance
  @param command: The command to be executed

  """
  chan = transport.open_session()
  chan.set_combine_stderr(True)
  output_handler = chan.makefile("r")
  chan.exec_command(command)

  result = chan.recv_exit_status()
  msg = output_handler.read()

  out_msg = "'%s' exited with status code %s, output %r" % (command, result,
                                                            msg)

  # If result is -1 (no exit status provided) we assume it was not successful
  if result:
    raise RemoteCommandError(out_msg)

  if msg:
    logging.info(out_msg)


def _InvokeDaemonUtil(transport, command):
  """Invokes daemon-util on the remote side.

  @param transport: The paramiko transport instance
  @param command: The daemon-util command to be run

  """
  _RunRemoteCommand(transport, "%s %s" % (constants.DAEMON_UTIL, command))


def _ReadSftpFile(sftp, filename):
  """Reads a file over sftp.

  @param sftp: An open paramiko SFTP client
  @param filename: The filename of the file to read
  @return: The content of the file

  """
  remote_file = sftp.open(filename, "r")
  try:
    return remote_file.read()
  finally:
    remote_file.close()


def _WriteSftpFile(sftp, name, perm, data):
  """SFTPs data to a remote file.

  @param sftp: A open paramiko SFTP client
  @param name: The remote file name
  @param perm: The remote file permission
  @param data: The data to write

  """
  remote_file = sftp.open(name, "w")
  try:
    sftp.chmod(name, perm)
    remote_file.write(data)
  finally:
    remote_file.close()


def SetupSSH(transport):
  """Sets the SSH up on the other side.

  @param transport: The paramiko transport instance

  """
  priv_key, pub_key, auth_keys = ssh.GetUserFiles(constants.GANETI_RUNAS)
  keyfiles = [
    (constants.SSH_HOST_DSA_PRIV, 0600),
    (constants.SSH_HOST_DSA_PUB, 0644),
    (constants.SSH_HOST_RSA_PRIV, 0600),
    (constants.SSH_HOST_RSA_PUB, 0644),
    (priv_key, 0600),
    (pub_key, 0644),
    ]

  sftp = transport.open_sftp_client()

  filemap = dict((name, (utils.ReadFile(name), perm))
                 for (name, perm) in keyfiles)

  auth_path = os.path.dirname(auth_keys)

  try:
    sftp.mkdir(auth_path, 0700)
  except IOError:
    # Sadly paramiko doesn't provide errno or similiar
    # so we can just assume that the path already exists
    logging.info("Path %s seems already to exist on remote node. Ignoring.",
                 auth_path)

  for name, (data, perm) in filemap.iteritems():
    _WriteSftpFile(sftp, name, perm, data)

  authorized_keys = sftp.open(auth_keys, "a+")
  try:
    # Due to the way SFTPFile and BufferedFile are implemented,
    # opening in a+ mode and then issuing a read(), readline() or
    # iterating over the file (which uses read() internally) will see
    # an empty file, since the paramiko internal file position and the
    # OS-level file-position are desynchronized; therefore, we issue
    # an explicit seek to resynchronize these; writes should (note
    # should) still go to the right place
    authorized_keys.seek(0, 0)
    # We don't have to close, as the close happened already in AddAuthorizedKey
    utils.AddAuthorizedKey(authorized_keys, filemap[pub_key][0])
  finally:
    authorized_keys.close()

  _InvokeDaemonUtil(transport, "reload-ssh-keys")


def ParseOptions():
  """Parses options passed to program.

  """
  program = os.path.basename(sys.argv[0])

  parser = optparse.OptionParser(usage=("%prog [--debug|--verbose] [--force]"
                                        " <node> <node...>"), prog=program)
  parser.add_option(cli.DEBUG_OPT)
  parser.add_option(cli.VERBOSE_OPT)
  parser.add_option(cli.NOSSH_KEYCHECK_OPT)
  default_key = ssh.GetUserFiles(constants.GANETI_RUNAS)[0]
  parser.add_option(optparse.Option("-f", dest="private_key",
                                    default=default_key,
                                    help="The private key to (try to) use for"
                                    "authentication "))
  parser.add_option(optparse.Option("--key-type", dest="key_type",
                                    choices=("rsa", "dsa"), default="dsa",
                                    help="The private key type (rsa or dsa)"))
  parser.add_option(optparse.Option("-j", "--force-join", dest="force_join",
                                    action="store_true", default=False,
                                    help="Force the join of the host"))

  (options, args) = parser.parse_args()

  return (options, args)


def SetupLogging(options):
  """Sets up the logging.

  @param options: Parsed options

  """
  fmt = "%(asctime)s: %(threadName)s "
  if options.debug or options.verbose:
    fmt += "%(levelname)s "
  fmt += "%(message)s"

  formatter = logging.Formatter(fmt)

  file_handler = logging.FileHandler(constants.LOG_SETUP_SSH)
  stderr_handler = logging.StreamHandler()
  stderr_handler.setFormatter(formatter)
  file_handler.setFormatter(formatter)
  file_handler.setLevel(logging.INFO)

  if options.debug:
    stderr_handler.setLevel(logging.DEBUG)
  elif options.verbose:
    stderr_handler.setLevel(logging.INFO)
  else:
    stderr_handler.setLevel(logging.WARNING)

  root_logger = logging.getLogger("")
  root_logger.setLevel(logging.NOTSET)
  root_logger.addHandler(stderr_handler)
  root_logger.addHandler(file_handler)

  # This is the paramiko logger instance
  paramiko_logger = logging.getLogger("paramiko")
  paramiko_logger.addHandler(file_handler)
  # We don't want to debug Paramiko, so filter anything below warning
  paramiko_logger.setLevel(logging.WARNING)


def LoadPrivateKeys(options):
  """Load the list of available private keys.

  It loads the standard ssh key from disk and then tries to connect to
  the ssh agent too.

  @rtype: list
  @return: a list of C{paramiko.PKey}

  """
  if options.key_type == "rsa":
    pkclass = paramiko.RSAKey
  elif options.key_type == "dsa":
    pkclass = paramiko.DSSKey
  else:
    logging.critical("Unknown key type %s selected (choose either rsa or dsa)",
                     options.key_type)
    sys.exit(1)

  try:
    private_key = pkclass.from_private_key_file(options.private_key)
  except (paramiko.SSHException, EnvironmentError), err:
    logging.critical("Can't load private key %s: %s", options.private_key, err)
    sys.exit(1)

  try:
    agent = paramiko.Agent()
    agent_keys = agent.get_keys()
  except paramiko.SSHException, err:
    # this will only be seen when the agent is broken/uses invalid
    # protocol; for non-existing agent, get_keys() will just return an
    # empty tuple
    logging.warning("Can't connect to the ssh agent: %s; skipping its use",
                    err)
    agent_keys = []

  return [private_key] + list(agent_keys)


def LoginViaKeys(transport, username, keys):
  """Try to login on the given transport via a list of keys.

  @param transport: the transport to use
  @param username: the username to login as
  @type keys: list
  @param keys: list of C{paramiko.PKey} to use for authentication
  @rtype: boolean
  @return: True or False depending on whether the login was
      successfull or not

  """
  for private_key in keys:
    try:
      transport.auth_publickey(username, private_key)
      fpr = ":".join("%02x" % ord(i) for i in private_key.get_fingerprint())
      if isinstance(private_key, paramiko.AgentKey):
        logging.debug("Authentication via the ssh-agent key %s", fpr)
      else:
        logging.debug("Authenticated via public key %s", fpr)
      return True
    except paramiko.SSHException:
      continue
  else:
    # all keys exhausted
    return False


def LoadKnownHosts():
  """Load the known hosts.

  @return: paramiko.util.load_host_keys dict

  """
  homedir = utils.GetHomeDir(constants.GANETI_RUNAS)
  known_hosts = os.path.join(homedir, ".ssh", "known_hosts")

  try:
    return paramiko.util.load_host_keys(known_hosts)
  except EnvironmentError:
    # We didn't found the path, silently ignore and return an empty dict
    return {}


def main():
  """Main routine.

  """
  (options, args) = ParseOptions()

  SetupLogging(options)

  errs = 0

  all_keys = LoadPrivateKeys(options)

  passwd = None
  username = constants.GANETI_RUNAS
  ssh_port = netutils.GetDaemonPort("ssh")
  host_keys = LoadKnownHosts()

  # Below, we need to join() the transport objects, as otherwise the
  # following happens:
  # - the main thread finishes
  # - the atexit functions run (in the main thread), and cause the
  #   logging file to be closed
  # - a tiny bit later, the transport thread is finally ending, and
  #   wants to log one more message, which fails as the file is closed
  #   now

  for host in args:
    transport = paramiko.Transport((host, ssh_port))
    transport.start_client()
    server_key = transport.get_remote_server_key()
    keytype = server_key.get_name()

    our_server_key = host_keys.get(host, {}).get(keytype, None)
    if options.ssh_key_check:
      if not our_server_key:
        hexified_key = ssh.FormatParamikoFingerprint(
            paramiko.util.hexify(server_key.get_fingerprint()))
        msg = ("Unable to verify hostkey of host %s: %s. Do you want to accept"
               " it?" % (host, hexified_key))

        if cli.AskUser(msg):
          our_server_key = server_key

      if our_server_key != server_key:
        logging.error("Unable to verify identity of host. Aborting")
        transport.close()
        transport.join()
        # TODO: Run over all hosts, fetch the keys and let them verify from the
        #       user beforehand then proceed with actual work later on
        raise paramiko.SSHException("Unable to verify identity of host")

    try:
      if LoginViaKeys(transport, username, all_keys):
        logging.info("Authenticated to %s via public key", host)
      else:
        logging.warning("Authentication to %s via public key failed, trying"
                        " password", host)
        if passwd is None:
          passwd = getpass.getpass(prompt="%s password:" % username)
        transport.auth_password(username=username, password=passwd)
        logging.info("Authenticated to %s via password", host)
    except paramiko.SSHException, err:
      logging.error("Connection or authentication failed to host %s: %s",
                    host, err)
      transport.close()
      # this is needed for compatibility with older Paramiko or Python
      # versions
      transport.join()
      continue
    try:
      try:
        if not _CheckJoin(transport):
          if options.force_join:
            logging.warning("Host %s failed join check, forced to continue",
                            host)
          else:
            raise JoinCheckError(("Host %s failed join check; Please verify"
                                  " that the host was not previously joined"
                                  " to another cluster and use --force-join"
                                  " to continue") % host)
        SetupSSH(transport)
      except errors.GenericError, err:
        logging.error("While doing setup on host %s an error occurred: %s",
                      host, err)
        errs += 1
    finally:
      transport.close()
      # this is needed for compatibility with older Paramiko or Python
      # versions
      transport.join()

    if errs > 0:
      sys.exit(1)


if __name__ == "__main__":
  main()