From ebf1c1a2e1dc2e60d7cfd2d33df7b32c9ba618a5 Mon Sep 17 00:00:00 2001
From: Jonathan Paynter <jonathan.paynter@arm.com>
Date: Thu, 10 Sep 2020 13:45:20 +0100
Subject: [PATCH] utils/ssh: Add paramiko based scp transfers

Using scp over paramiko allows scp transfers to be treated similarly to
sftp transfers, instead of requiring subprocesses, and provides
the ability to monitor an scp transfer using a callback as can be done
using sftp.
---
 devlib/utils/ssh.py | 33 +++++++++++++++++++++------------
 setup.py            |  1 +
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/devlib/utils/ssh.py b/devlib/utils/ssh.py
index 05bd623..1a85208 100644
--- a/devlib/utils/ssh.py
+++ b/devlib/utils/ssh.py
@@ -36,6 +36,7 @@ from future.utils import raise_from
 
 from paramiko.client import SSHClient, AutoAddPolicy, RejectPolicy
 import paramiko.ssh_exception
+from scp import SCPClient
 # By default paramiko is very verbose, including at the INFO level
 logging.getLogger("paramiko").setLevel(logging.WARNING)
 
@@ -446,6 +447,10 @@ class SshConnection(SshConnectionBase):
         sftp.get_channel().settimeout(timeout)
         return sftp
 
+    
+    def _get_scp(self, timeout):
+        return SCPClient(self.client.get_transport(), socket_timeout=timeout)
+
     @classmethod
     def _push_file(cls, sftp, src, dst):
         try:
@@ -560,21 +565,25 @@ class SshConnection(SshConnectionBase):
 
     def push(self, sources, dest, timeout=30):
         # If using scp, use implementation from base class
-        if self.use_scp:
-            super().push(sources, dest, timeout)
-        else:
-            with _handle_paramiko_exceptions(), self._get_sftp(timeout) as sftp:
-                for source in sources:
-                    self._push_path(sftp, source, dest)
+        with _handle_paramiko_exceptions():
+            if self.use_scp:
+                scp = self._get_scp(timeout)
+                scp.put(sources, dest, recursive=True)
+            else:
+                with self._get_sftp(timeout) as sftp:
+                    for source in sources:
+                        self._push_path(sftp, source, dest)
 
     def pull(self, sources, dest, timeout=30):
         # If using scp, use implementation from base class
-        if self.use_scp:
-            super().pull(sources, dest, timeout)
-        else:
-            with _handle_paramiko_exceptions(), self._get_sftp(timeout) as sftp:
-                for source in sources:
-                    self._pull_path(sftp, source, dest)
+        with _handle_paramiko_exceptions():
+            if self.use_scp:
+                scp = self._get_scp(timeout)
+                scp.get(sources, dest, recursive=True)
+            else:
+                with self._get_sftp(timeout) as sftp:
+                    for source in sources:
+                        self._pull_path(sftp, source, dest)
 
     def execute(self, command, timeout=None, check_exit_code=True,
                 as_root=False, strip_colors=True, will_succeed=False): #pylint: disable=unused-argument
diff --git a/setup.py b/setup.py
index ca30eb1..84bf3df 100644
--- a/setup.py
+++ b/setup.py
@@ -83,6 +83,7 @@ params = dict(
         'pexpect>=3.3',  # Send/recieve to/from device
         'pyserial',  # Serial port interface
         'paramiko', # SSH connection
+        'scp', # SSH connection file transfers
         'wrapt',  # Basic for construction of decorator functions
         'future', # Python 2-3 compatibility
         'enum34;python_version<"3.4"', # Enums for Python < 3.4