From patchwork Sun Jul 4 13:42:21 2010 Content-Type: text/plain; charset="utf-8" MIME-Version: 1.0 Content-Transfer-Encoding: 7bit X-Patchwork-Submitter: Michael Goldish X-Patchwork-Id: 110148 Received: from vger.kernel.org (vger.kernel.org [209.132.180.67]) by demeter.kernel.org (8.14.4/8.14.3) with ESMTP id o64DgsPU005698 for ; Sun, 4 Jul 2010 13:42:54 GMT Received: (majordomo@vger.kernel.org) by vger.kernel.org via listexpand id S1757884Ab0GDNmp (ORCPT ); Sun, 4 Jul 2010 09:42:45 -0400 Received: from mx1.redhat.com ([209.132.183.28]:44044 "EHLO mx1.redhat.com" rhost-flags-OK-OK-OK-OK) by vger.kernel.org with ESMTP id S1757879Ab0GDNmn (ORCPT ); Sun, 4 Jul 2010 09:42:43 -0400 Received: from int-mx03.intmail.prod.int.phx2.redhat.com (int-mx03.intmail.prod.int.phx2.redhat.com [10.5.11.16]) by mx1.redhat.com (8.13.8/8.13.8) with ESMTP id o64Dgg8X010057 (version=TLSv1/SSLv3 cipher=DHE-RSA-AES256-SHA bits=256 verify=OK); Sun, 4 Jul 2010 09:42:42 -0400 Received: from ns3.rdu.redhat.com (ns3.rdu.redhat.com [10.11.255.199]) by int-mx03.intmail.prod.int.phx2.redhat.com (8.13.8/8.13.8) with ESMTP id o64DgfP3013062; Sun, 4 Jul 2010 09:42:41 -0400 Received: from localhost.localdomain (dhcp-1-188.tlv.redhat.com [10.35.1.188]) by ns3.rdu.redhat.com (8.13.8/8.13.8) with ESMTP id o64DgXbP028739; Sun, 4 Jul 2010 09:42:39 -0400 From: Michael Goldish To: autotest@test.kernel.org, kvm@vger.kernel.org Cc: Michael Goldish Subject: [KVM-AUTOTEST PATCH v4] [RFC] KVM test: add python client for rss file transfer services Date: Sun, 4 Jul 2010 16:42:21 +0300 Message-Id: <1278250941-30933-2-git-send-email-mgoldish@redhat.com> In-Reply-To: <1278250941-30933-1-git-send-email-mgoldish@redhat.com> References: <1278250941-30933-1-git-send-email-mgoldish@redhat.com> X-Scanned-By: MIMEDefang 2.67 on 10.5.11.16 Sender: kvm-owner@vger.kernel.org Precedence: bulk List-ID: X-Mailing-List: kvm@vger.kernel.org X-Greylist: IP, sender and recipient auto-whitelisted, not delayed by milter-greylist-4.2.3 (demeter.kernel.org [140.211.167.41]); Sun, 04 Jul 2010 13:42:55 +0000 (UTC) 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()