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 paramiko
from helper.install_lab import exec_cmd
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
            LOG.info("Waiting 10s")
            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} 2>/dev/null'
        exec_cmd(cmd)

    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}')
    exec_cmd(cmd)


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:
        LOG.info("Waiting 10s")
        time.sleep(10)