SFTP client: add get_file and open_file methods
All checks were successful
ci/woodpecker/tag/woodpecker Pipeline was successful

This commit is contained in:
Benjamin Renard 2022-06-30 14:34:25 +02:00
parent 9511b31a79
commit 54d6c6e0f3
2 changed files with 61 additions and 17 deletions

View file

@ -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:
sys.exit(1)
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)
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:
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')
sftp.close()

View file

@ -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 """