rss_client.py

#!/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.
"""

from __future__ import division
import socket
import struct
import time
import sys
import os
import glob

import six


# 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):

    def __init__(self, msg, e=None, filename=None):
        Exception.__init__(self, msg, e, filename)
        self.msg = msg
        self.e = e
        self.filename = filename

    def __str__(self):
        s = self.msg
        if self.e and self.filename:
            s += "    (error: %s,    filename: %s)" % (self.e, self.filename)
        elif self.e:
            s += "    (%s)" % self.e
        elif self.filename:
            s += "    (filename: %s)" % self.filename
        return s


class FileTransferConnectError(FileTransferError):
    pass


class FileTransferTimeoutError(FileTransferError):
    pass


class FileTransferProtocolError(FileTransferError):
    pass


class FileTransferSocketError(FileTransferError):
    pass


class FileTransferServerError(FileTransferError):

    def __init__(self, errmsg):
        FileTransferError.__init__(self, None, errmsg)

    def __str__(self):
        s = "Server said: %r" % self.e
        if self.filename:
            s += "    (filename: %s)" % self.filename
        return s


class FileTransferNotFoundError(FileTransferError):
    pass


class FileTransferClient(object):

    """
    Connect to a RSS (remote shell server) and transfer files.
    """

    def __init__(self, address, port, log_func=None, timeout=20):
        """
        Connect to a server.

        :param address: The server's address
        :param port: The server's port
        :param log_func: If provided, transfer stats will be passed to this
                function during the transfer
        :param timeout: Time duration to wait for connection to succeed
        :raise FileTransferConnectError: Raised if the connection fails
        """
        family = ":" in address and socket.AF_INET6 or socket.AF_INET
        self._socket = socket.socket(family, socket.SOCK_STREAM)
        self._socket.settimeout(timeout)
        try:
            addrinfo = socket.getaddrinfo(address, port, family,
                                          socket.SOCK_STREAM,
                                          socket.IPPROTO_TCP)
            self._socket.connect(addrinfo[0][4])
        except socket.error as e:
            raise FileTransferConnectError("Cannot connect to server at "
                                           "%s:%s" % (address, port), e)
        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))
        self._log_func = log_func
        self._last_time = time.time()
        self._last_transferred = 0
        self.transferred = 0

    def __del__(self):
        self.close()

    def close(self):
        """
        Close the connection.
        """
        self._socket.close()

    def _send(self, sr, timeout=60):
        try:
            if timeout <= 0:
                raise socket.timeout
            self._socket.settimeout(timeout)
            self._socket.sendall(sr)
        except socket.timeout:
            raise FileTransferTimeoutError("Timeout expired while sending "
                                           "data to server")
        except socket.error as e:
            raise FileTransferSocketError("Could not send data to server", e)

    def _receive(self, size, timeout=60):
        strs = []
        end_time = time.time() + timeout
        try:
            while size > 0:
                timeout = end_time - time.time()
                if timeout <= 0:
                    raise socket.timeout
                self._socket.settimeout(timeout)
                data = self._socket.recv(size)
                if not data:
                    raise FileTransferProtocolError("Connection closed "
                                                    "unexpectedly while "
                                                    "receiving data from "
                                                    "server")
                strs.append(data)
                size -= len(data)
        except socket.timeout:
            raise FileTransferTimeoutError("Timeout expired while receiving "
                                           "data from server")
        except socket.error as e:
            raise FileTransferSocketError("Error receiving data from server",
                                          e)
        return b"".join(strs)

    def _report_stats(self, sr):
        if self._log_func:
            dt = time.time() - self._last_time
            if dt >= 1:
                transferred = self.transferred / 1048576.
                speed = (self.transferred - self._last_transferred) / dt
                speed /= 1048576.
                self._log_func("%s %.3f MB (%.3f MB/sec)" %
                               (sr, transferred, speed))
                self._last_time = time.time()
                self._last_transferred = self.transferred

    def _send_packet(self, sr, timeout=60):
        self._send(struct.pack("=I", len(sr)))
        self._send(sr, timeout)
        self.transferred += len(sr) + 4
        self._report_stats("Sent")

    def _receive_packet(self, timeout=60):
        """return bytes"""
        size = struct.unpack("=I", self._receive(4))[0]
        sr = self._receive(size, timeout)
        self.transferred += len(sr) + 4
        self._report_stats("Received")
        return sr

    def _send_file_chunks(self, filename, timeout=60):
        if self._log_func:
            self._log_func("Sending file %s" % filename)
        f = open(filename, "rb")
        try:
            try:
                end_time = time.time() + timeout
                while True:
                    data = f.read(CHUNKSIZE)
                    self._send_packet(data, end_time - time.time())
                    if len(data) < CHUNKSIZE:
                        break
            except FileTransferError as e:
                e.filename = filename
                raise
        finally:
            f.close()

    def _receive_file_chunks(self, filename, timeout=60):
        if self._log_func:
            self._log_func("Receiving file %s" % filename)
        f = open(filename, "wb")
        try:
            try:
                end_time = time.time() + timeout
                while True:
                    data = self._receive_packet(end_time - time.time())
                    f.write(data)
                    if len(data) < CHUNKSIZE:
                        break
            except FileTransferError as e:
                e.filename = filename
                raise
        finally:
            f.close()

    def _send_msg(self, msg, timeout=60):
        self._send(struct.pack("=I", msg))

    def _receive_msg(self, timeout=60):
        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
            six.reraise(*e)
        if msg == RSS_ERROR:
            errmsg = self._receive_packet().decode()
            raise FileTransferServerError(errmsg)
        six.reraise(*e)


class FileUploadClient(FileTransferClient):

    """
    Connect to a RSS (remote shell server) and upload files or directory trees.
    """

    def __init__(self, address, port, log_func=None, timeout=20):
        """
        Connect to a server.

        :param address: The server's address
        :param port: The server's port
        :param log_func: If provided, transfer stats will be passed to this
                function during the transfer
        :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 FileTransferSocketError: Raised if the RSS_UPLOAD message cannot
                be sent to the server
        """
        super(FileUploadClient, self).__init__(
            address, port, log_func, 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).encode())
            self._send_file_chunks(path, end_time - time.time())
        elif os.path.isdir(path):
            self._send_msg(RSS_CREATE_DIR)
            self._send_packet(os.path.basename(path).encode())
            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.encode())
                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(end_time - time.time())
                if msg == RSS_OK:
                    return
                elif msg == RSS_ERROR:
                    errmsg = self._receive_packet().decode()
                    raise FileTransferServerError(errmsg)
                else:
                    # Neither RSS_OK nor RSS_ERROR found
                    raise FileTransferProtocolError("Received unexpected msg")
        except Exception:
            # 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, log_func=None, timeout=20):
        """
        Connect to a server.

        :param address: The server's address
        :param port: The server's port
        :param log_func: If provided, transfer stats will be passed to this
                function during the transfer
        :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, log_func, 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.encode())
            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().decode()
                    if os.path.isdir(dst_path):
                        dst_path = os.path.join(dst_path, filename)
                    self._receive_file_chunks(dst_path, 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().decode()
                    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().decode()
                    raise FileTransferServerError(errmsg)
                else:
                    # Unexpected msg
                    raise FileTransferProtocolError("Received unexpected msg")
        except Exception:
            # In any case, if the transfer failed, close the connection
            self.close()
            raise


def upload(address, port, src_pattern, dst_path, log_func=None, timeout=60,
           connect_timeout=20):
    """
    Connect to server and upload files.

    :see:: FileUploadClient
    """
    client = FileUploadClient(address, port, log_func, connect_timeout)
    client.upload(src_pattern, dst_path, timeout)
    client.close()


def download(address, port, src_pattern, dst_path, log_func=None, timeout=60,
             connect_timeout=20):
    """
    Connect to server and upload files.

    :see:: FileDownloadClient
    """
    client = FileDownloadClient(address, port, log_func, connect_timeout)
    client.download(src_pattern, dst_path, timeout)
    client.close()


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("-v", "--verbose",
                      action="store_true", dest="verbose",
                      help="be verbose")
    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)

    logger = None
    if options.verbose:
        def p(s):
            print(s)
        logger = p

    if options.download:
        download(address, port, src_pattern, dst_path, logger, options.timeout)
    elif options.upload:
        upload(address, port, src_pattern, dst_path, logger, options.timeout)


if __name__ == "__main__":
    main()