virtual-deployment/virtualbox/pybox/utils/sftp.py

#!/usr/bin/python3
#
# SPDX-License-Identifier: Apache-2.0
#

"""
This module provides functionality to send files and directories to remote servers using the
rsync and paramiko libraries.
"""

import getpass
import os
import time
import subprocess
import paramiko
from utils.install_log import LOG


def sftp_send(source, destination, client_dict):
    """
    Send files to remote server
    """

    remote_host = client_dict["remote_host"]
    username = client_dict["username"]
    sftp_client = None

    LOG.info("Connecting to server %s with username %s", remote_host, username)

    ssh_client = paramiko.SSHClient()
    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    # TODO(WEI): need to make this timeout handling better
    retry = 0
    while retry < 8:
        try:
            ssh_client.connect(
                remote_host,
                port=client_dict["remote_port"],
                username=username,
                password=client_dict["password"],
                look_for_keys=False,
                allow_agent=False
            )
            sftp_client = ssh_client.open_sftp()
            retry = 8
        except Exception:  # pylint: disable=W0703
            LOG.info("******* try again")
            retry += 1
            time.sleep(10)

    LOG.info("Sending file from %s to %s", source, destination)
    if sftp_client:
        sftp_client.put(source, destination)
        LOG.info("Done")
        sftp_client.close()
    ssh_client.close()


# pylint: disable=R0801
def send_dir(params_dict):
    """
    Send directory `source` to remote host `remote_host` at port `remote_port` and destination
    `destination` using `rsync` over `ssh`.

    Args:
        params_dict (dict): A dictionary containing the following keys:
            - source (str): The local directory to be sent.
            - remote_host (str): The IP address or domain name of the remote host.
            - remote_port (int): The port number of the remote host to connect to.
            - destination (str): The remote directory to copy `source` into.
            - username (str): The username for the SSH connection.
            - password (str): The password for the SSH connection.
            - follow_links (bool, optional): Whether to follow symbolic links when
                copying files. Default is True.
            - clear_known_hosts (bool, optional): Whether to clear the known_hosts file
                before making the SSH connection. Default is True.

    Raises:
        Exception: If there is an error in `rsync`, raises an exception with the return code.

    Note:
        This method only works from a Linux environment.
    """

    source = params_dict['source']
    remote_host = params_dict['remote_host']
    remote_port = params_dict['remote_port']
    destination = params_dict['destination']
    username = params_dict['username']
    password = params_dict['password']
    follow_links = params_dict.get('follow_links', True)
    clear_known_hosts = params_dict.get('clear_known_hosts', True)

    # Only works from linux for now
    if not source.endswith('/') or not source.endswith('\\'):
        source = source + '/'

    follow_links = "L" if follow_links else ""
    if clear_known_hosts:
        if remote_host == '127.0.0.1':
            keygen_arg = f"[127.0.0.1]:{remote_port}"
        else:
            keygen_arg = remote_host
        cmd = f'ssh-keygen -f "/home/{getpass.getuser()}/.ssh/known_hosts" -R {keygen_arg}'
        LOG.info("CMD: %s", cmd)
        with subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE) as process:
            for line in iter(process.stdout.readline, b''):
                LOG.info("%s", line.decode("utf-8").strip())
            process.wait()

    LOG.info('Running rsync of dir: %s -> %s@%s:%s', source, username, remote_host, destination)
    cmd = (f'rsync -av{follow_links} --rsh="/usr/bin/sshpass -p {password} '
           f'ssh -p {remote_port} -o StrictHostKeyChecking=no -l {username}" '
           f'{source}* {username}@{remote_host}:{destination}')
    LOG.info("CMD: %s", cmd)

    with subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE) as process:
        for line in iter(process.stdout.readline, b''):
            LOG.info("%s", line.decode("utf-8").strip())
        process.wait()
    if process.returncode:
        raise Exception(f"Error in rsync, return code:{process.returncode}")  # pylint: disable=E0012, W0719


def send_dir_fallback(source, remote_host, destination, username, password):
    """
    Send directory contents to remote server, usually controller-0
    Note: does not send nested directories only files.
    args:
    - source: full path to directory
    e.g. /localhost/loadbuild/jenkins/latest_build/
    - Remote host: name of host to log into, controller-0 by default
    e.g. myhost.com
    - destination: where to store the file on host: /home/myuser/
    """

    LOG.info("Connecting to server %s with username %s", remote_host, username)
    ssh_client = paramiko.SSHClient()
    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh_client.connect(
        remote_host,
        username=username,
        password=password,
        look_for_keys=False,
        allow_agent=False
    )
    sftp_client = ssh_client.open_sftp()
    send_img = False
    for items in os.listdir(source):
        path = source + items
        if os.path.isfile(path):
            if items.endswith('.img'):
                remote_path = destination + 'images/' + items
                LOG.info("Sending file from %s to %s", path, remote_path)
                sftp_client.put(path, remote_path)
                send_img = True
            elif items.endswith('.iso'):
                pass
            else:
                remote_path = destination + items
                LOG.info("Sending file from %s to %s", path, remote_path)
                sftp_client.put(path, remote_path)
    LOG.info("Done")
    sftp_client.close()
    ssh_client.close()
    if send_img:
        time.sleep(10)