From 54d6c6e0f33d1a6dcdadbff29785d2da86350979 Mon Sep 17 00:00:00 2001
From: Benjamin Renard <brenard@easter-eggs.com>
Date: Thu, 30 Jun 2022 14:34:25 +0200
Subject: [PATCH] SFTP client: add get_file and open_file methods

---
 mylib/scripts/sftp_test.py | 67 +++++++++++++++++++++++++++++---------
 mylib/sftp.py              | 11 +++++--
 2 files changed, 61 insertions(+), 17 deletions(-)

diff --git a/mylib/scripts/sftp_test.py b/mylib/scripts/sftp_test.py
index 9cdf5dd..1e10933 100644
--- a/mylib/scripts/sftp_test.py
+++ b/mylib/scripts/sftp_test.py
@@ -1,10 +1,13 @@
 # -*- coding: utf-8 -*-
 
 """ Test SFTP client """
+import atexit
 import tempfile
 import logging
 import sys
 import os
+import random
+import string
 
 import getpass
 
@@ -46,24 +49,58 @@ def main(argv=None):  # pylint: disable=too-many-locals,too-many-statements
     log.info('Initialize Email client')
     sftp = SFTPClient(options=options, just_try=options.just_try)
     sftp.connect()
+    atexit.register(sftp.close)
 
     log.debug('Create tempory file')
-    tmp_file = tempfile.NamedTemporaryFile()  # pylint: disable=consider-using-with
-    log.debug('Temporary file path: "%s"', tmp_file.name)
-    tmp_file.write(b'Juste un test.')
+    test_content = b'Juste un test.'
+    tmp_dir = tempfile.TemporaryDirectory()  # pylint: disable=consider-using-with
+    tmp_file = os.path.join(
+        tmp_dir.name,
+        f'tmp{"".join(random.choice(string.ascii_lowercase) for i in range(8))}'
+    )
+    log.debug('Temporary file path: "%s"', tmp_file)
+    with open(tmp_file, 'wb') as file_desc:
+        file_desc.write(test_content)
+
     log.debug(
-        'Upload file %s to SFTP server (in %s)', tmp_file.name,
+        'Upload file %s to SFTP server (in %s)', tmp_file,
         options.upload_path if options.upload_path else "remote initial connection directory")
-    if not sftp.upload_file(tmp_file.name, options.upload_path):
+    if not sftp.upload_file(tmp_file, options.upload_path):
         log.error('Fail to upload test file on SFTP server')
-    else:
-        log.info('Test file uploaded on SFTP server')
-        remote_filepath = (
-            os.path.join(options.upload_path, os.path.basename(tmp_file.name))
-            if options.upload_path else os.path.basename(tmp_file.name)
-        )
-        if sftp.remove_file(remote_filepath):
-            log.info('Test file removed on SFTP server')
+        sys.exit(1)
+
+    log.info('Test file uploaded on SFTP server')
+    remote_filepath = (
+        os.path.join(options.upload_path, os.path.basename(tmp_file))
+        if options.upload_path else os.path.basename(tmp_file)
+    )
+
+    with tempfile.NamedTemporaryFile() as tmp_file2:
+        log.info('Retrieve test file to %s', tmp_file2.name)
+        if not sftp.get_file(remote_filepath, tmp_file2.name):
+            log.error('Fail to retrieve test file')
         else:
-            log.error('Fail to remove test file on SFTP server')
-    sftp.close()
+            with open(tmp_file2.name, 'rb') as file_desc:
+                content = file_desc.read()
+                log.debug('Read content: %s', content)
+                if test_content == content:
+                    log.info('Content file retrieved match with uploaded one')
+                else:
+                    log.error('Content file retrieved doest not match with uploaded one')
+
+    try:
+        log.info('Remotly open test file %s', remote_filepath)
+        file_desc = sftp.open_file(remote_filepath)
+        content = file_desc.read()
+        log.debug('Read content: %s', content)
+        if test_content == content:
+            log.info('Content of remote file match with uploaded one')
+        else:
+            log.error('Content of remote file doest not match with uploaded one')
+    except Exception:  # pylint: disable=broad-except
+        log.exception('An exception occurred remotly opening test file %s', remote_filepath)
+
+    if sftp.remove_file(remote_filepath):
+        log.info('Test file removed on SFTP server')
+    else:
+        log.error('Fail to remove test file on SFTP server')
diff --git a/mylib/sftp.py b/mylib/sftp.py
index c5db741..0a2d9ec 100644
--- a/mylib/sftp.py
+++ b/mylib/sftp.py
@@ -101,6 +101,14 @@ class SFTPClient(ConfigurableObject):
             log.debug("Fail to retreive remote directory, use empty string instead")
             self.initial_directory = ""
 
+    def get_file(self, remote_filepath, local_filepath):
+        log.debug("Retreive file '%s' to '%s'", remote_filepath, local_filepath)
+        return self.sftp_client.get(remote_filepath, local_filepath) is None
+
+    def open_file(self, remote_filepath, mode='r'):
+        log.debug("Remotly open file '%s'", remote_filepath)
+        return self.sftp_client.open(remote_filepath, mode=mode)
+
     def upload_file(self, filepath, remote_directory=None):
         """ Upload a file on SFTP server """
         self.connect()
@@ -124,8 +132,7 @@ class SFTPClient(ConfigurableObject):
         if self._get_option('just_try'):
             log.debug("Just - try mode: do not really remove file '%s'", filepath)
             return True
-        self.sftp_client.remove(filepath)
-        return True
+        return self.sftp_client.remove(filepath) is None
 
     def close(self):
         """ Close SSH/SFTP connection """