remote.py

"""
Utils to run commands on remote server

Copyright (C) 2010-2020 Cr@ns <roots@crans.org>
Authors : Daniel Stan <daniel.stan@crans.org>
          Vincent Le Gallic <legallic@crans.org>
          Alexandre Iooss <erdnaxe@crans.org>
SPDX-License-Identifier: GPL-3.0-or-later
"""

from functools import lru_cache
from getpass import getpass
import json
import logging
from pathlib import Path

from paramiko.client import SSHClient
from paramiko.config import SSHConfig
from paramiko.ssh_exception import SSHException, PasswordRequiredException, AuthenticationException

from .locale import _

# Local logger
log = logging.getLogger(__name__)


@lru_cache()
def create_ssh_client(host, password=None):
    """
    Create a SSH client with paramiko module
    """
    # Create SSH client with system host keys and agent
    client = SSHClient()

    # Load config file and use the right username
    try:
        config = SSHConfig()
        config.parse(Path.home().joinpath(".ssh/config").open())
        username = config.lookup(host).get("user", None)
    except FileNotFoundError:
        username=None

    # Load system private keys
    client.load_system_host_keys()

    # Connect
    try:
        client.connect(host, username=username, password=password)
    except PasswordRequiredException:
        password = getpass("SSH password: ")
        return create_ssh_client(host, password)
    except AuthenticationException:
        log.error(_("SSH authentication failed."))
        exit(1)
    except SSHException:
        log.error(_("An error occured during SSH connection, debug with -vv"))
        raise

    return client


def remote_command(options, command, arg=None, stdin_contents=None):
    """
    Execute remote command and return output
    """
    if "host" not in options.serverdata:
        log.error("Missing parameter `host` in active server configuration")
        exit(1)
    client = create_ssh_client(str(options.serverdata['host']))

    # Build command
    if "remote_cmd" not in options.serverdata:
        log.error("Missing parameter `remote_cmd` in active server configuration")
        exit(1)
    remote_cmd = options.serverdata['remote_cmd'] + " " + command
    if arg:
        remote_cmd += " " + arg

    # Run command and timeout after 10s
    log.info(_("Running command `%s`") % remote_cmd)
    stdin, stdout, stderr = client.exec_command(remote_cmd, timeout=10)

    # Write
    if stdin_contents is not None:
        log.info(_("Writing to stdin: %s") % stdin_contents)
        stdin.write(json.dumps(stdin_contents))
        stdin.flush()

    # If the server is not expected to exit, then exit now
    if stdin_contents:
        return

    # Return code == 0 if success
    ret = stdout.channel.recv_exit_status()
    if ret != 0:
        err = ""
        if stderr.channel.recv_stderr_ready():
            err = stderr.read()
        log.error(_("Wrong server return code %s, error is %s") % (ret, err))
        exit(ret)

    # Decode directly read buffer
    try:
        answer = json.load(stdout)
    except ValueError:
        log.error(_("Error while parsing JSON"))
        exit(42)

    log.debug(_("Server returned %s") % answer)
    return answer