From 0a1fd907015d2af7fa6f3d3ba7ebf82867026f2f Mon Sep 17 00:00:00 2001 From: Pierre Letessier Date: Tue, 26 Jul 2016 14:51:36 +0200 Subject: [PATCH 1/3] Add a limit_bandwidth parameter to get and put function. The timeout has to be increased to handle pauses happening at start. The more the bandwidth is limited, the more the timeout should be. --- scp.py | 14 +++++++++++--- test.py | 18 ++++++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/scp.py b/scp.py index c4f2499..31441c5 100644 --- a/scp.py +++ b/scp.py @@ -120,7 +120,7 @@ def __exit__(self, type, value, traceback): self.close() def put(self, files, remote_path=b'.', - recursive=False, preserve_times=False): + recursive=False, preserve_times=False, limit_bandwidth=0): """ Transfer files to remote host. @@ -135,12 +135,16 @@ def put(self, files, remote_path=b'.', @param preserve_times: preserve mtime and atime of transfered files and directories. @type preserve_times: bool + @param limit_bandwidth: Limits the used bandwidth, specified in Kbit/s. + @type limit_bandwidth: int """ self.preserve_times = preserve_times self.channel = self._open() self._pushed = 0 self.channel.settimeout(self.socket_timeout) - scp_command = (b'scp -t ', b'scp -r -t ')[recursive] + scp_command = b'scp -t ' + scp_command += (b'', b' -r ')[recursive] + scp_command += (b'', b' -l '+ bytes(str(limit_bandwidth), 'ascii') + b' ')[limit_bandwidth > 0] self.channel.exec_command(scp_command + self.sanitize(asbytes(remote_path))) self._recv_confirm() @@ -156,7 +160,7 @@ def put(self, files, remote_path=b'.', self.close() def get(self, remote_path, local_path='', - recursive=False, preserve_times=False): + recursive=False, preserve_times=False, limit_bandwidth=0): """ Transfer files from remote host to localhost @@ -171,6 +175,8 @@ def get(self, remote_path, local_path='', @param preserve_times: preserve mtime and atime of transfered files and directories. @type preserve_times: bool + @param limit_bandwidth: Limits the used bandwidth, specified in Kbit/s. + @type limit_bandwidth: int """ if not isinstance(remote_path, (list, tuple)): remote_path = [remote_path] @@ -187,12 +193,14 @@ def get(self, remote_path, local_path='', asunicode(self._recv_dir)) rcsv = (b'', b' -r')[recursive] prsv = (b'', b' -p')[preserve_times] + lmbw = (b'', b' -l '+ bytes(str(limit_bandwidth), 'ascii'))[limit_bandwidth > 0] self.channel = self._open() self._pushed = 0 self.channel.settimeout(self.socket_timeout) self.channel.exec_command(b"scp" + rcsv + prsv + + lmbw + b" -f " + b' '.join(remote_path)) self._recv_all() diff --git a/test.py b/test.py index a9787f5..10e0720 100644 --- a/test.py +++ b/test.py @@ -299,6 +299,24 @@ def test_up_and_down(self): finally: os.chdir(previous) + def test_up_and_down_with_limit(self): + '''send and receive files with the same client''' + previous = os.getcwd() + testfile = os.path.join(self._temp, 'testfile') + testfile_sent = os.path.join(self._temp, 'testfile_sent') + testfile_rcvd = os.path.join(self._temp, 'testfile_rcvd') + try: + os.chdir(self._temp) + with open(testfile, 'w') as f: + f.write("TESTING\n") + with SCPClient(self.ssh.get_transport(), socket_timeout=60.0) as scp: + scp.put(testfile, testfile_sent, limit_bandwidth=10) + scp.get(testfile_sent, testfile_rcvd, limit_bandwidth=10) + + assert open(testfile_rcvd).read() == 'TESTING\n' + finally: + os.chdir(previous) + if __name__ == '__main__': unittest.main() From 23848cea606d588de599a65b3d0e56b036f93876 Mon Sep 17 00:00:00 2001 From: Pierre Letessier Date: Tue, 26 Jul 2016 14:51:36 +0200 Subject: [PATCH 2/3] Add a limit_bandwidth parameter to get and put function. The timeout has to be increased to handle pauses happening at start. The more the bandwidth is limited, the more the timeout should be. --- scp.py | 14 +++++++++++--- test.py | 18 ++++++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/scp.py b/scp.py index c4f2499..ed5e56e 100644 --- a/scp.py +++ b/scp.py @@ -120,7 +120,7 @@ def __exit__(self, type, value, traceback): self.close() def put(self, files, remote_path=b'.', - recursive=False, preserve_times=False): + recursive=False, preserve_times=False, limit_bandwidth=0): """ Transfer files to remote host. @@ -135,12 +135,16 @@ def put(self, files, remote_path=b'.', @param preserve_times: preserve mtime and atime of transfered files and directories. @type preserve_times: bool + @param limit_bandwidth: Limits the used bandwidth, specified in Kbit/s. + @type limit_bandwidth: int """ self.preserve_times = preserve_times self.channel = self._open() self._pushed = 0 self.channel.settimeout(self.socket_timeout) - scp_command = (b'scp -t ', b'scp -r -t ')[recursive] + scp_command = b'scp -t ' + scp_command += (b'', b' -r ')[recursive] + scp_command += (b'', b' -l '+ str(limit_bandwidth).encode() + b' ')[limit_bandwidth > 0] self.channel.exec_command(scp_command + self.sanitize(asbytes(remote_path))) self._recv_confirm() @@ -156,7 +160,7 @@ def put(self, files, remote_path=b'.', self.close() def get(self, remote_path, local_path='', - recursive=False, preserve_times=False): + recursive=False, preserve_times=False, limit_bandwidth=0): """ Transfer files from remote host to localhost @@ -171,6 +175,8 @@ def get(self, remote_path, local_path='', @param preserve_times: preserve mtime and atime of transfered files and directories. @type preserve_times: bool + @param limit_bandwidth: Limits the used bandwidth, specified in Kbit/s. + @type limit_bandwidth: int """ if not isinstance(remote_path, (list, tuple)): remote_path = [remote_path] @@ -187,12 +193,14 @@ def get(self, remote_path, local_path='', asunicode(self._recv_dir)) rcsv = (b'', b' -r')[recursive] prsv = (b'', b' -p')[preserve_times] + lmbw = (b'', b' -l '+ str(limit_bandwidth).encode())[limit_bandwidth > 0] self.channel = self._open() self._pushed = 0 self.channel.settimeout(self.socket_timeout) self.channel.exec_command(b"scp" + rcsv + prsv + + lmbw + b" -f " + b' '.join(remote_path)) self._recv_all() diff --git a/test.py b/test.py index a9787f5..10e0720 100644 --- a/test.py +++ b/test.py @@ -299,6 +299,24 @@ def test_up_and_down(self): finally: os.chdir(previous) + def test_up_and_down_with_limit(self): + '''send and receive files with the same client''' + previous = os.getcwd() + testfile = os.path.join(self._temp, 'testfile') + testfile_sent = os.path.join(self._temp, 'testfile_sent') + testfile_rcvd = os.path.join(self._temp, 'testfile_rcvd') + try: + os.chdir(self._temp) + with open(testfile, 'w') as f: + f.write("TESTING\n") + with SCPClient(self.ssh.get_transport(), socket_timeout=60.0) as scp: + scp.put(testfile, testfile_sent, limit_bandwidth=10) + scp.get(testfile_sent, testfile_rcvd, limit_bandwidth=10) + + assert open(testfile_rcvd).read() == 'TESTING\n' + finally: + os.chdir(previous) + if __name__ == '__main__': unittest.main() From 2a217ec7d48dccc4e3703d939fdee46718d34cbe Mon Sep 17 00:00:00 2001 From: Pierre Letessier Date: Wed, 27 Jul 2016 15:42:13 +0200 Subject: [PATCH 3/3] Use a proper if to test the rate-limit value --- scp.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/scp.py b/scp.py index ed5e56e..995c606 100644 --- a/scp.py +++ b/scp.py @@ -144,7 +144,8 @@ def put(self, files, remote_path=b'.', self.channel.settimeout(self.socket_timeout) scp_command = b'scp -t ' scp_command += (b'', b' -r ')[recursive] - scp_command += (b'', b' -l '+ str(limit_bandwidth).encode() + b' ')[limit_bandwidth > 0] + if limit_bandwidth > 0: + scp_command += b' -l ' + str(limit_bandwidth).encode() + b' ' self.channel.exec_command(scp_command + self.sanitize(asbytes(remote_path))) self._recv_confirm()