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 -*- # -*- coding: utf-8 -*-
""" Test SFTP client """ """ Test SFTP client """
import atexit
import tempfile import tempfile
import logging import logging
import sys import sys
import os import os
import random
import string
import getpass import getpass
@ -46,24 +49,58 @@ def main(argv=None): # pylint: disable=too-many-locals,too-many-statements
log.info('Initialize Email client') log.info('Initialize Email client')
sftp = SFTPClient(options=options, just_try=options.just_try) sftp = SFTPClient(options=options, just_try=options.just_try)
sftp.connect() sftp.connect()
atexit.register(sftp.close)
log.debug('Create tempory file') log.debug('Create tempory file')
tmp_file = tempfile.NamedTemporaryFile() # pylint: disable=consider-using-with test_content = b'Juste un test.'
log.debug('Temporary file path: "%s"', tmp_file.name) tmp_dir = tempfile.TemporaryDirectory() # pylint: disable=consider-using-with
tmp_file.write(b'Juste un test.') 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( 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") 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') log.error('Fail to upload test file on SFTP server')
else: sys.exit(1)
log.info('Test file uploaded on SFTP server') log.info('Test file uploaded on SFTP server')
remote_filepath = ( remote_filepath = (
os.path.join(options.upload_path, 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.name) 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): if sftp.remove_file(remote_filepath):
log.info('Test file removed on SFTP server') log.info('Test file removed on SFTP server')
else: else:
log.error('Fail to remove test file on SFTP server') 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") log.debug("Fail to retreive remote directory, use empty string instead")
self.initial_directory = "" 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): def upload_file(self, filepath, remote_directory=None):
""" Upload a file on SFTP server """ """ Upload a file on SFTP server """
self.connect() self.connect()
@ -124,8 +132,7 @@ class SFTPClient(ConfigurableObject):
if self._get_option('just_try'): if self._get_option('just_try'):
log.debug("Just - try mode: do not really remove file '%s'", filepath) log.debug("Just - try mode: do not really remove file '%s'", filepath)
return True return True
self.sftp_client.remove(filepath) return self.sftp_client.remove(filepath) is None
return True
def close(self): def close(self):
""" Close SSH/SFTP connection """ """ Close SSH/SFTP connection """