Best Python code snippet using molecule_python
remote_operations.py
Source:remote_operations.py  
1#!/usr/bin/env python32"""Remote access utilities, via ssh & scp."""3import optparse4import os5import posixpath6import re7import shlex8import sys9import time10import subprocess11# Get relative imports to work when the package is not installed on the PYTHONPATH.12if __name__ == "__main__" and __package__ is None:13    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))14_IS_WINDOWS = sys.platform == "win32" or sys.platform == "cygwin"15_OPERATIONS = ["shell", "copy_to", "copy_from"]16_SSH_CONNECTION_ERRORS = [17    "Connection refused",18    "Connection timed out during banner exchange",19    "Permission denied",20    "System is booting up.",21    "ssh_exchange_identification: read: Connection reset by peer",22]23def posix_path(path):24    """Return posix path, used on Windows since scp requires posix style paths."""25    # If path is already quoted, we need to remove the quotes before calling26    path_quote = "\'" if path.startswith("\'") else ""27    path_quote = "\"" if path.startswith("\"") else path_quote28    if path_quote:29        path = path[1:-1]30    drive, new_path = os.path.splitdrive(path)31    if drive:32        new_path = posixpath.join("/cygdrive", drive.split(":")[0], *re.split("/|\\\\", new_path))33    return "{quote}{path}{quote}".format(quote=path_quote, path=new_path)34class RemoteOperations(object):  # pylint: disable=too-many-instance-attributes35    """Class to support remote operations."""36    def __init__(  # pylint: disable=too-many-arguments37            self, user_host, ssh_connection_options=None, ssh_options=None, scp_options=None,38            retries=0, retry_sleep=0, debug=False, shell_binary="/bin/bash", use_shell=False):39        """Initialize RemoteOperations."""40        self.user_host = user_host41        self.ssh_connection_options = ssh_connection_options if ssh_connection_options else ""42        self.ssh_options = ssh_options if ssh_options else ""43        self.scp_options = scp_options if scp_options else ""44        self.retries = retries45        self.retry_sleep = retry_sleep46        self.debug = debug47        self.shell_binary = shell_binary48        self.use_shell = use_shell49        # Check if we can remotely access the host.50        self._access_code, self._access_buff = self._remote_access()51    def _call(self, cmd):52        if self.debug:53            print(cmd)54        # If use_shell is False we need to split the command up into a list.55        if not self.use_shell:56            cmd = shlex.split(cmd)57        # Use a common pipe for stdout & stderr for logging.58        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,59                                   shell=self.use_shell)60        buff_stdout, _ = process.communicate()61        return process.poll(), buff_stdout.decode("utf-8", "replace")62    def _remote_access(self):63        """Check if a remote session is possible."""64        cmd = "ssh {} {} {} date".format(self.ssh_connection_options, self.ssh_options,65                                         self.user_host)66        attempt_num = 067        buff = ""68        while True:69            ret, buff = self._call(cmd)70            # Ignore any connection errors before sshd has fully initialized.71            if not ret and not any(ssh_error in buff for ssh_error in _SSH_CONNECTION_ERRORS):72                return ret, buff73            attempt_num += 174            if attempt_num > self.retries:75                break76            if self.debug:77                print("Failed remote attempt {}, retrying in {} seconds".format(78                    attempt_num, self.retry_sleep))79            time.sleep(self.retry_sleep)80        return ret, buff81    def _perform_operation(self, cmd):82        return self._call(cmd)83    def access_established(self):84        """Return True if initial access was established."""85        return not self._access_code86    def access_info(self):87        """Return the return code and output buffer from initial access attempt(s)."""88        return self._access_code, self._access_buff89    @staticmethod90    def ssh_error(message):91        """Return True if the error message is generated from the ssh client.92        This can help determine if an error is due to a remote operation failing or an ssh93        related issue, like a connection issue.94        """95        return message.startswith("ssh:")96    def operation(  # pylint: disable=too-many-branches97            self, operation_type, operation_param, operation_dir=None):98        """Execute Main entry for remote operations. Returns (code, output).99        'operation_type' supports remote shell and copy operations.100        'operation_param' can either be a list or string of commands or files.101        'operation_dir' is '.' if unspecified for 'copy_*'.102        """103        if not self.access_established():104            return self.access_info()105        # File names with a space must be quoted, since we permit the106        # the file names to be either a string or a list.107        if operation_type.startswith("copy") and isinstance(operation_param, str):108            operation_param = shlex.split(operation_param, posix=not _IS_WINDOWS)109        cmds = []110        if operation_type == "shell":111            if operation_dir is not None:112                operation_param = "cd {}; {}".format(operation_dir, operation_param)113            dollar = ""114            if re.search("\"|'", operation_param):115                # To ensure any quotes in operation_param are handled correctly when116                # invoking the operation_param, escape with \ and add $ in the front.117                # See https://stackoverflow.com/questions/8254120/118                #   how-to-escape-a-single-quote-in-single-quote-string-in-bash119                operation_param = "{}".format(operation_param.replace("'", r"\'"))120                operation_param = "{}".format(operation_param.replace("\"", r"\""))121                dollar = "$"122            cmd = "ssh {} {} {} {} -c \"{}'{}'\"".format(self.ssh_connection_options,123                                                         self.ssh_options, self.user_host,124                                                         self.shell_binary, dollar, operation_param)125            cmds.append(cmd)126        elif operation_type == "copy_to":127            cmd = "scp -r {} {} ".format(self.ssh_connection_options, self.scp_options)128            # To support spaces in the filename or directory, we quote them one at a time.129            for copy_file in operation_param:130                # Quote file on Posix.131                quote = "\"" if not _IS_WINDOWS else ""132                cmd += "{quote}{file}{quote} ".format(quote=quote, file=posix_path(copy_file))133            operation_dir = operation_dir if operation_dir else ""134            cmd += " {}:{}".format(self.user_host, posix_path(operation_dir))135            cmds.append(cmd)136        elif operation_type == "copy_from":137            operation_dir = operation_dir if operation_dir else "."138            if not os.path.isdir(operation_dir):139                raise ValueError("Local directory '{}' does not exist.".format(operation_dir))140            # We support multiple files being copied from the remote host141            # by invoking scp for each file specified.142            # Note - this is a method which scp does not support directly.143            for copy_file in operation_param:144                copy_file = posix_path(copy_file)145                cmd = "scp -r {} {} {}:".format(self.ssh_connection_options, self.scp_options,146                                                self.user_host)147                # Quote (on Posix), and escape the file if there are spaces.148                # Note - we do not support other non-ASCII characters in a file name.149                quote = "\"" if not _IS_WINDOWS else ""150                if " " in copy_file:151                    copy_file = re.escape("{quote}{file}{quote}".format(152                        quote=quote, file=copy_file))153                cmd += "{} {}".format(copy_file, posix_path(operation_dir))154                cmds.append(cmd)155        else:156            raise ValueError("Invalid operation '{}' specified, choose from {}.".format(157                operation_type, _OPERATIONS))158        final_ret = 0159        buff = ""160        for cmd in cmds:161            ret, new_buff = self._perform_operation(cmd)162            buff += new_buff163            final_ret = final_ret or ret164        return final_ret, buff165    def shell(self, operation_param, operation_dir=None):166        """Provide helper for remote shell operations."""167        return self.operation(operation_type="shell", operation_param=operation_param,168                              operation_dir=operation_dir)169    def copy_to(self, operation_param, operation_dir=None):170        """Provide helper for remote copy_to operations."""171        return self.operation(operation_type="copy_to", operation_param=operation_param,172                              operation_dir=operation_dir)173    def copy_from(self, operation_param, operation_dir=None):174        """Provide helper for remote copy_from operations."""175        return self.operation(operation_type="copy_from", operation_param=operation_param,176                              operation_dir=operation_dir)177def main():  # pylint: disable=too-many-branches,too-many-statements178    """Execute Main program."""179    parser = optparse.OptionParser(description=__doc__)180    control_options = optparse.OptionGroup(parser, "Control options")181    shell_options = optparse.OptionGroup(parser, "Shell options")182    copy_options = optparse.OptionGroup(parser, "Copy options")183    parser.add_option(184        "--userHost", dest="user_host", default=None,185        help=("User and remote host to execute commands on [REQUIRED]."186              " Examples, 'user@1.2.3.4' or 'user@myhost.com'."))187    parser.add_option(188        "--operation", dest="operation", default="shell", choices=_OPERATIONS,189        help=("Remote operation to perform, choose one of '{}',"190              " defaults to '%default'.".format(", ".join(_OPERATIONS))))191    control_options.add_option(192        "--sshConnectionOptions", dest="ssh_connection_options", default=None, action="append",193        help=("SSH connection options which are common to ssh and scp."194              " More than one option can be specified either"195              " in one quoted string or by specifying"196              " this option more than once. Example options:"197              " '-i $HOME/.ssh/access.pem -o ConnectTimeout=10"198              " -o ConnectionAttempts=10'"))199    control_options.add_option(200        "--sshOptions", dest="ssh_options", default=None, action="append",201        help=("SSH specific options."202              " More than one option can be specified either"203              " in one quoted string or by specifying"204              " this option more than once. Example options:"205              " '-t' or '-T'"))206    control_options.add_option(207        "--scpOptions", dest="scp_options", default=None, action="append",208        help=("SCP specific options."209              " More than one option can be specified either"210              " in one quoted string or by specifying"211              " this option more than once. Example options:"212              " '-l 5000'"))213    control_options.add_option(214        "--retries", dest="retries", type=int, default=0,215        help=("Number of retries to attempt for operation,"216              " defaults to '%default'."))217    control_options.add_option(218        "--retrySleep", dest="retry_sleep", type=int, default=10,219        help=("Number of seconds to wait between retries,"220              " defaults to '%default'."))221    control_options.add_option("--debug", dest="debug", action="store_true", default=False,222                               help="Provides debug output.")223    control_options.add_option("--verbose", dest="verbose", action="store_true", default=False,224                               help="Print exit status and output at end.")225    shell_options.add_option(226        "--commands", dest="remote_commands", default=None, action="append",227        help=("Commands to excute on the remote host. The"228              " commands must be separated by a ';' and can either"229              " be specifed in a quoted string or by specifying"230              " this option more than once. A ';' will be added"231              " between commands when this option is specifed"232              " more than once."))233    shell_options.add_option(234        "--commandDir", dest="command_dir", default=None,235        help=("Working directory on remote to execute commands"236              " form. Defaults to remote login directory."))237    copy_options.add_option(238        "--file", dest="files", default=None, action="append",239        help=("The file to copy to/from remote host. To"240              " support spaces in the file, each file must be"241              " specified using this option more than once."))242    copy_options.add_option(243        "--remoteDir", dest="remote_dir", default=None,244        help=("Remote directory to copy to, only applies when"245              " operation is 'copy_to'. Defaults to the login"246              " directory on the remote host."))247    copy_options.add_option(248        "--localDir", dest="local_dir", default=".",249        help=("Local directory to copy to, only applies when"250              " operation is 'copy_from'. Defaults to the"251              " current directory, '%default'."))252    parser.add_option_group(control_options)253    parser.add_option_group(shell_options)254    parser.add_option_group(copy_options)255    (options, _) = parser.parse_args()256    if not getattr(options, "user_host", None):257        parser.print_help()258        parser.error("Missing required option")259    if options.operation == "shell":260        if not getattr(options, "remote_commands", None):261            parser.print_help()262            parser.error("Missing required '{}' option '{}'".format(options.operation,263                                                                    "--commands"))264        operation_param = ";".join(options.remote_commands)265        operation_dir = options.command_dir266    else:267        if not getattr(options, "files", None):268            parser.print_help()269            parser.error("Missing required '{}' option '{}'".format(options.operation, "--file"))270        operation_param = options.files271        if options.operation == "copy_to":272            operation_dir = options.remote_dir273        else:274            operation_dir = options.local_dir275    if not options.ssh_connection_options:276        ssh_connection_options = None277    else:278        ssh_connection_options = " ".join(options.ssh_connection_options)279    if not options.ssh_options:280        ssh_options = None281    else:282        ssh_options = " ".join(options.ssh_options)283    if not options.scp_options:284        scp_options = None285    else:286        scp_options = " ".join(options.scp_options)287    remote_op = RemoteOperations(288        user_host=options.user_host, ssh_connection_options=ssh_connection_options,289        ssh_options=ssh_options, scp_options=scp_options, retries=options.retries,290        retry_sleep=options.retry_sleep, debug=options.debug)291    ret_code, buff = remote_op.operation(options.operation, operation_param, operation_dir)292    if options.verbose:293        print("Return code: {} for command {}".format(ret_code, sys.argv))294        print(buff)295    sys.exit(ret_code)296if __name__ == "__main__":...Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.
You could also refer to video tutorials over LambdaTest YouTube channel to get step by step demonstration from industry experts.
Get 100 minutes of automation test minutes FREE!!
