| #!/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): |
| 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 |
| """ |
| self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
| self._socket.settimeout(timeout) |
| try: |
| self._socket.connect((address, port)) |
| except socket.error, 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, str, timeout=60): |
| try: |
| if timeout <= 0: |
| raise socket.timeout |
| self._socket.settimeout(timeout) |
| self._socket.sendall(str) |
| except socket.timeout: |
| raise FileTransferTimeoutError("Timeout expired while sending " |
| "data to server") |
| except socket.error, 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, e: |
| raise FileTransferSocketError("Error receiving data from server", |
| e) |
| return "".join(strs) |
| |
| |
| def _report_stats(self, str): |
| 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)" % |
| (str, transferred, speed)) |
| self._last_time = time.time() |
| self._last_transferred = self.transferred |
| |
| |
| def _send_packet(self, str, timeout=60): |
| self._send(struct.pack("=I", len(str))) |
| self._send(str, timeout) |
| self.transferred += len(str) + 4 |
| self._report_stats("Sent") |
| |
| |
| def _receive_packet(self, timeout=60): |
| size = struct.unpack("=I", self._receive(4))[0] |
| str = self._receive(size, timeout) |
| self.transferred += len(str) + 4 |
| self._report_stats("Received") |
| return str |
| |
| |
| 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, 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, 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 |
| raise e[0], e[1], e[2] |
| if msg == RSS_ERROR: |
| errmsg = self._receive_packet() |
| raise FileTransferServerError(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, 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)) |
| 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)) |
| 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(end_time - time.time()) |
| if msg == RSS_OK: |
| return |
| elif msg == RSS_ERROR: |
| errmsg = self._receive_packet() |
| raise FileTransferServerError(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, 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) |
| 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, 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(errmsg) |
| else: |
| # Unexpected msg |
| raise FileTransferProtocolError("Received unexpected msg") |
| except: |
| # 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() |
