Re: [KVM-AUTOTEST PATCH v4] [RFC] KVM test: add python client for rss file transfer services

[Lucas Meneghel Rodrigues <lmr@redhat.com>]

On Sun, 2010-07-04 at 16:42 +0300, Michael Goldish wrote:
> See details in docstrings in rss_file_transfer.py.
> See protocol details in deps/rss.cpp.

I've gone through both rss.cpp code and rss_file_transfer.py code, and
must say it's very good. The only comment I'd make about this module is
that the command line syntax for it when using it as a stand alone
program could work similarly to utilities like rsync and scp, due to the
familiar model of usage, and also, the module/program could be called
rsscp instead of rss_file_transfer. Of course this is a minor comment,
and can be changed at a later time.

I'll commit this patchset shortly, together with a patch making
guest_test able to use file transfer already, as well as publish a new
winutils.iso.

> Changes from v3:
> - Protocol change: instead of sending a file as one big packet, send it in
>   multiple chunks.  See details in deps/rss.cpp.
> 
> Changes from v2:
> - Raise FileTransferNotFoundError if no files/dirs are transferred (due to
>   a bad path or wildcard pattern)
> - Make all connection related errors in the base class raise
>   FileTransferConnectError
> 
> Changes from v1:
> - Use glob() instead of iglob() (Python 2.4 doesn't like iglob())
> - Change a few comments
> 
> Signed-off-by: Michael Goldish <mgoldish@redhat.com>
> ---
>  client/tests/kvm/rss_file_transfer.py |  431 +++++++++++++++++++++++++++++++++
>  1 files changed, 431 insertions(+), 0 deletions(-)
>  create mode 100755 client/tests/kvm/rss_file_transfer.py
> 
> diff --git a/client/tests/kvm/rss_file_transfer.py b/client/tests/kvm/rss_file_transfer.py
> new file mode 100755
> index 0000000..c584397
> --- /dev/null
> +++ b/client/tests/kvm/rss_file_transfer.py
> @@ -0,0 +1,431 @@
> +#!/usr/bin/python
> +"""
> +Client for file transfer services offered by RSS (Remote Shell Server).
> +
> +@author: Michael Goldish (mgoldish@redhat.com)
> +@copyright: 2008-2010 Red Hat Inc.
> +"""
> +
> +import socket, struct, time, sys, os, glob
> +
> +# Globals
> +CHUNKSIZE = 65536
> +
> +# Protocol message constants
> +RSS_MAGIC           = 0x525353
> +RSS_OK              = 1
> +RSS_ERROR           = 2
> +RSS_UPLOAD          = 3
> +RSS_DOWNLOAD        = 4
> +RSS_SET_PATH        = 5
> +RSS_CREATE_FILE     = 6
> +RSS_CREATE_DIR      = 7
> +RSS_LEAVE_DIR       = 8
> +RSS_DONE            = 9
> +
> +# See rss.cpp for protocol details.
> +
> +
> +class FileTransferError(Exception):
> +    pass
> +
> +
> +class FileTransferConnectError(FileTransferError):
> +    pass
> +
> +
> +class FileTransferTimeoutError(FileTransferError):
> +    pass
> +
> +
> +class FileTransferProtocolError(FileTransferError):
> +    pass
> +
> +
> +class FileTransferSendError(FileTransferError):
> +    pass
> +
> +
> +class FileTransferServerError(FileTransferError):
> +    pass
> +
> +
> +class FileTransferNotFoundError(FileTransferError):
> +    pass
> +
> +
> +class FileTransferClient(object):
> +    """
> +    Connect to a RSS (remote shell server) and transfer files.
> +    """
> +
> +    def __init__(self, address, port, timeout=10):
> +        """
> +        Connect to a server.
> +
> +        @param address: The server's address
> +        @param port: The server's port
> +        @param timeout: Time duration to wait for connection to succeed
> +        @raise FileTransferConnectError: Raised if the connection fails
> +        @raise FileTransferProtocolError: Raised if an incorrect magic number
> +                is received
> +        """
> +        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
> +        self._socket.settimeout(timeout)
> +        try:
> +            self._socket.connect((address, port))
> +        except socket.error:
> +            raise FileTransferConnectError("Could not connect to server")
> +        try:
> +            if self._receive_msg(timeout) != RSS_MAGIC:
> +                raise FileTransferConnectError("Received wrong magic number")
> +        except FileTransferTimeoutError:
> +            raise FileTransferConnectError("Timeout expired while waiting to "
> +                                           "receive magic number")
> +        self._send(struct.pack("=i", CHUNKSIZE))
> +
> +
> +    def __del__(self):
> +        self.close()
> +
> +
> +    def close(self):
> +        """
> +        Close the connection.
> +        """
> +        self._socket.close()
> +
> +
> +    def _send(self, str):
> +        try:
> +            self._socket.sendall(str)
> +        except socket.error:
> +            raise FileTransferSendError("Could not send data to server")
> +
> +
> +    def _receive(self, size, timeout=10):
> +        strs = []
> +        end_time = time.time() + timeout
> +        while size > 0:
> +            try:
> +                self._socket.settimeout(max(0.0001, end_time - time.time()))
> +                data = self._socket.recv(size)
> +            except socket.timeout:
> +                raise FileTransferTimeoutError("Timeout expired while "
> +                                               "receiving data from server")
> +            except socket.error:
> +                raise FileTransferProtocolError("Error receiving data from "
> +                                                "server")
> +            if not data:
> +                raise FileTransferProtocolError("Connection closed "
> +                                                "unexpectedly")
> +            strs.append(data)
> +            size -= len(data)
> +        return "".join(strs)
> +
> +
> +    def _send_packet(self, str):
> +        self._send(struct.pack("=I", len(str)))
> +        self._send(str)
> +
> +
> +    def _receive_packet(self, timeout=10):
> +        size = struct.unpack("=I", self._receive(4))[0]
> +        return self._receive(size, timeout)
> +
> +
> +    def _send_file_chunks(self, filename, timeout=30):
> +        f = open(filename, "rb")
> +        try:
> +            end_time = time.time() + timeout
> +            while time.time() < end_time:
> +                data = f.read(CHUNKSIZE)
> +                self._send_packet(data)
> +                if len(data) < CHUNKSIZE:
> +                    break
> +            else:
> +                raise FileTransferTimeoutError("Timeout expired while sending "
> +                                               "file %s" % filename)
> +        finally:
> +            f.close()
> +
> +
> +    def _receive_file_chunks(self, filename, timeout=30):
> +        f = open(filename, "wb")
> +        try:
> +            end_time = time.time() + timeout
> +            while True:
> +                try:
> +                    data = self._receive_packet(end_time - time.time())
> +                except FileTransferTimeoutError:
> +                    raise FileTransferTimeoutError("Timeout expired while "
> +                                                   "receiving file %s" %
> +                                                   filename)
> +                except FileTransferProtocolError:
> +                    raise FileTransferProtocolError("Error receiving file %s" %
> +                                                    filename)
> +                f.write(data)
> +                if len(data) < CHUNKSIZE:
> +                    break
> +        finally:
> +            f.close()
> +
> +
> +    def _send_msg(self, msg, timeout=10):
> +        self._send(struct.pack("=I", msg))
> +
> +
> +    def _receive_msg(self, timeout=10):
> +        s = self._receive(4, timeout)
> +        return struct.unpack("=I", s)[0]
> +
> +
> +    def _handle_transfer_error(self):
> +        # Save original exception
> +        e = sys.exc_info()
> +        try:
> +            # See if we can get an error message
> +            msg = self._receive_msg()
> +        except FileTransferError:
> +            # No error message -- re-raise original exception
> +            raise e[0], e[1], e[2]
> +        if msg == RSS_ERROR:
> +            errmsg = self._receive_packet()
> +            raise FileTransferServerError("Server said: %s" % errmsg)
> +        raise e[0], e[1], e[2]
> +
> +
> +class FileUploadClient(FileTransferClient):
> +    """
> +    Connect to a RSS (remote shell server) and upload files or directory trees.
> +    """
> +
> +    def __init__(self, address, port, timeout=10):
> +        """
> +        Connect to a server.
> +
> +        @param address: The server's address
> +        @param port: The server's port
> +        @param timeout: Time duration to wait for connection to succeed
> +        @raise FileTransferConnectError: Raised if the connection fails
> +        @raise FileTransferProtocolError: Raised if an incorrect magic number
> +                is received
> +        @raise FileTransferSendError: Raised if the RSS_UPLOAD message cannot
> +                be sent to the server
> +        """
> +        super(FileUploadClient, self).__init__(address, port, timeout)
> +        self._send_msg(RSS_UPLOAD)
> +
> +
> +    def _upload_file(self, path, end_time):
> +        if os.path.isfile(path):
> +            self._send_msg(RSS_CREATE_FILE)
> +            self._send_packet(os.path.basename(path))
> +            self._send_file_chunks(path, max(0, end_time - time.time()))
> +        elif os.path.isdir(path):
> +            self._send_msg(RSS_CREATE_DIR)
> +            self._send_packet(os.path.basename(path))
> +            for filename in os.listdir(path):
> +                self._upload_file(os.path.join(path, filename), end_time)
> +            self._send_msg(RSS_LEAVE_DIR)
> +
> +
> +    def upload(self, src_pattern, dst_path, timeout=600):
> +        """
> +        Send files or directory trees to the server.
> +        The semantics of src_pattern and dst_path are similar to those of scp.
> +        For example, the following are OK:
> +            src_pattern='/tmp/foo.txt', dst_path='C:\\'
> +                (uploads a single file)
> +            src_pattern='/usr/', dst_path='C:\\Windows\\'
> +                (uploads a directory tree recursively)
> +            src_pattern='/usr/*', dst_path='C:\\Windows\\'
> +                (uploads all files and directory trees under /usr/)
> +        The following is not OK:
> +            src_pattern='/tmp/foo.txt', dst_path='C:\\Windows\\*'
> +                (wildcards are only allowed in src_pattern)
> +
> +        @param src_pattern: A path or wildcard pattern specifying the files or
> +                directories to send to the server
> +        @param dst_path: A path in the server's filesystem where the files will
> +                be saved
> +        @param timeout: Time duration in seconds to wait for the transfer to
> +                complete
> +        @raise FileTransferTimeoutError: Raised if timeout expires
> +        @raise FileTransferServerError: Raised if something goes wrong and the
> +                server sends an informative error message to the client
> +        @note: Other exceptions can be raised.
> +        """
> +        end_time = time.time() + timeout
> +        try:
> +            try:
> +                self._send_msg(RSS_SET_PATH)
> +                self._send_packet(dst_path)
> +                matches = glob.glob(src_pattern)
> +                for filename in matches:
> +                    self._upload_file(os.path.abspath(filename), end_time)
> +                self._send_msg(RSS_DONE)
> +            except FileTransferTimeoutError:
> +                raise
> +            except FileTransferError:
> +                self._handle_transfer_error()
> +            else:
> +                # If nothing was transferred, raise an exception
> +                if not matches:
> +                    raise FileTransferNotFoundError("Pattern %s does not "
> +                                                    "match any files or "
> +                                                    "directories" %
> +                                                    src_pattern)
> +                # Look for RSS_OK or RSS_ERROR
> +                msg = self._receive_msg(max(0, end_time - time.time()))
> +                if msg == RSS_OK:
> +                    return
> +                elif msg == RSS_ERROR:
> +                    errmsg = self._receive_packet()
> +                    raise FileTransferServerError("Server said: %s" % errmsg)
> +                else:
> +                    # Neither RSS_OK nor RSS_ERROR found
> +                    raise FileTransferProtocolError("Received unexpected msg")
> +        except:
> +            # In any case, if the transfer failed, close the connection
> +            self.close()
> +            raise
> +
> +
> +class FileDownloadClient(FileTransferClient):
> +    """
> +    Connect to a RSS (remote shell server) and download files or directory trees.
> +    """
> +
> +    def __init__(self, address, port, timeout=10):
> +        """
> +        Connect to a server.
> +
> +        @param address: The server's address
> +        @param port: The server's port
> +        @param timeout: Time duration to wait for connection to succeed
> +        @raise FileTransferConnectError: Raised if the connection fails
> +        @raise FileTransferProtocolError: Raised if an incorrect magic number
> +                is received
> +        @raise FileTransferSendError: Raised if the RSS_UPLOAD message cannot
> +                be sent to the server
> +        """
> +        super(FileDownloadClient, self).__init__(address, port, timeout)
> +        self._send_msg(RSS_DOWNLOAD)
> +
> +
> +    def download(self, src_pattern, dst_path, timeout=600):
> +        """
> +        Receive files or directory trees from the server.
> +        The semantics of src_pattern and dst_path are similar to those of scp.
> +        For example, the following are OK:
> +            src_pattern='C:\\foo.txt', dst_path='/tmp'
> +                (downloads a single file)
> +            src_pattern='C:\\Windows', dst_path='/tmp'
> +                (downloads a directory tree recursively)
> +            src_pattern='C:\\Windows\\*', dst_path='/tmp'
> +                (downloads all files and directory trees under C:\\Windows)
> +        The following is not OK:
> +            src_pattern='C:\\Windows', dst_path='/tmp/*'
> +                (wildcards are only allowed in src_pattern)
> +
> +        @param src_pattern: A path or wildcard pattern specifying the files or
> +                directories, in the server's filesystem, that will be sent to
> +                the client
> +        @param dst_path: A path in the local filesystem where the files will
> +                be saved
> +        @param timeout: Time duration in seconds to wait for the transfer to
> +                complete
> +        @raise FileTransferTimeoutError: Raised if timeout expires
> +        @raise FileTransferServerError: Raised if something goes wrong and the
> +                server sends an informative error message to the client
> +        @note: Other exceptions can be raised.
> +        """
> +        dst_path = os.path.abspath(dst_path)
> +        end_time = time.time() + timeout
> +        file_count = 0
> +        dir_count = 0
> +        try:
> +            try:
> +                self._send_msg(RSS_SET_PATH)
> +                self._send_packet(src_pattern)
> +            except FileTransferError:
> +                self._handle_transfer_error()
> +            while True:
> +                msg = self._receive_msg()
> +                if msg == RSS_CREATE_FILE:
> +                    # Receive filename and file contents
> +                    filename = self._receive_packet()
> +                    if os.path.isdir(dst_path):
> +                        dst_path = os.path.join(dst_path, filename)
> +                    self._receive_file_chunks(
> +                            dst_path, max(0, end_time - time.time()))
> +                    dst_path = os.path.dirname(dst_path)
> +                    file_count += 1
> +                elif msg == RSS_CREATE_DIR:
> +                    # Receive dirname and create the directory
> +                    dirname = self._receive_packet()
> +                    if os.path.isdir(dst_path):
> +                        dst_path = os.path.join(dst_path, dirname)
> +                    if not os.path.isdir(dst_path):
> +                        os.mkdir(dst_path)
> +                    dir_count += 1
> +                elif msg == RSS_LEAVE_DIR:
> +                    # Return to parent dir
> +                    dst_path = os.path.dirname(dst_path)
> +                elif msg == RSS_DONE:
> +                    # Transfer complete
> +                    if not file_count and not dir_count:
> +                        raise FileTransferNotFoundError("Pattern %s does not "
> +                                                        "match any files or "
> +                                                        "directories that "
> +                                                        "could be downloaded" %
> +                                                        src_pattern)
> +                    break
> +                elif msg == RSS_ERROR:
> +                    # Receive error message and abort
> +                    errmsg = self._receive_packet()
> +                    raise FileTransferServerError("Server said: %s" % errmsg)
> +                else:
> +                    # Unexpected msg
> +                    raise FileTransferProtocolError("Received unexpected msg")
> +        except:
> +            # In any case, if the transfer failed, close the connection
> +            self.close()
> +            raise
> +
> +
> +def main():
> +    import optparse
> +
> +    usage = "usage: %prog [options] address port src_pattern dst_path"
> +    parser = optparse.OptionParser(usage=usage)
> +    parser.add_option("-d", "--download",
> +                      action="store_true", dest="download",
> +                      help="download files from server")
> +    parser.add_option("-u", "--upload",
> +                      action="store_true", dest="upload",
> +                      help="upload files to server")
> +    parser.add_option("-t", "--timeout",
> +                      type="int", dest="timeout", default=3600,
> +                      help="transfer timeout")
> +    options, args = parser.parse_args()
> +    if options.download == options.upload:
> +        parser.error("you must specify either -d or -u")
> +    if len(args) != 4:
> +        parser.error("incorrect number of arguments")
> +    address, port, src_pattern, dst_path = args
> +    port = int(port)
> +
> +    if options.download:
> +        client = FileDownloadClient(address, port)
> +        client.download(src_pattern, dst_path, timeout=options.timeout)
> +        client.close()
> +    elif options.upload:
> +        client = FileUploadClient(address, port)
> +        client.upload(src_pattern, dst_path, timeout=options.timeout)
> +        client.close()
> +
> +
> +if __name__ == "__main__":
> +    main()