Commit ba207b0f authored by JOSSOUD Olivier's avatar JOSSOUD Olivier
Browse files

SFTPDMC. Export SFTP-related functions to new common.sftp module.

parent a9595cdd
import paramiko
import os
import wimcollect.common.logger as logger
def connect(sftp_config: dict, log: logger.Logger, object_id: str) -> paramiko.SFTPClient:
"""Connect to SFTP server.
Notes
-----
SFTP connection opened by this function should be close with the `paramiko.SFTPClient.close()` function.
Parameters
----------
sftp_config: dict
Dict containing SFTP connection information: `host`, `user` and `private_key_path`.
log: wimcollect.common.logger
Logger to record debug/info/error messages.
object_id: str
Object identifier, used as marker for log messages.
Returns
-------
paramiko.SFTPClient
SFTP client.
"""
client = paramiko.SSHClient()
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
host = sftp_config["host"]
user = sftp_config["user"]
try:
client.connect(hostname=host,
username=user,
key_filename=sftp_config["private_key_path"])
except Exception as e:
log.write(object_id, "Failed to connect to " + host + " as " + user + ": " + str(e))
raise e
else:
log.write(object_id, "Connected to [" + host + "] as user: " + user)
return client.open_sftp()
def download_file(sftp_client: paramiko.SFTPClient, source_filepath: str, dest_filepath: str,
log: logger.Logger, object_id: str) -> bool:
"""Download a file from the FTP server.
Parameters
----------
sftp_client: paramiko.SFTPClient
SFTP client.
source_filepath: str
Full path of the distant to-be-downloaded file.
dest_filepath: str
Full path where the downloaded file should be stored.
log: wimcollect.logger
Logger to record debug/info/error messages.
object_id: str
Object identifier, used as marker for log messages.
Returns
-------
bool
`True` if everything went well, otherwise raises Exception.
"""
log.write(object_id, "Downloading " + source_filepath + " ...")
try:
sftp_client.get(source_filepath, dest_filepath)
except FileNotFoundError:
msg = "File not found: " + source_filepath
log.write(object_id, msg)
raise FileNotFoundError(msg)
# Check if the destination file exists
success = os.path.exists(dest_filepath)
if success:
log.write(object_id, "File [" + source_filepath + "] downloaded in [" + dest_filepath + "].")
else:
msg = "Failed to download [" + source_filepath + "] in [" + dest_filepath + "]."
log.write(object_id, msg)
raise Exception(msg)
return success
......@@ -2,10 +2,10 @@ import os
import re
import configobj
import datetime
import paramiko
import wimcollect.common.utils as utils
import wimcollect.common.logger as logger
import wimcollect.common.sftp as sftp
class Collector(utils.LogConfig):
......@@ -17,7 +17,7 @@ class Collector(utils.LogConfig):
self.distant_base_dir = self.config[self.object_id]["distant_base_dir"]
def download_picarro(self, day: datetime.date):
"""Download Picarro data file from SSH server.
"""Download Picarro data file from SFTP server.
The distant files will be deleted from the server if the transfer is successful.
......@@ -27,16 +27,21 @@ class Collector(utils.LogConfig):
"""
self.logger.write(self.object_id, "Download DMC Picarro data.")
sftp_client = self.__sftp_connect__()
# Build source file path
picarro_id = self.config[self.object_id]["picarro_id"]
picarro_number = re.sub("[^0-9]", "", picarro_id)
source_filepath = self.distant_base_dir + "/" + picarro_number \
+ "/DMC_" + picarro_id + "_" + day.strftime("%Y%m%d") + ".zip"
# Build destination file path
dest_filepath = self.__get_dest_filepath__(day, source_filepath)
success = self.__sftp_download__(sftp_client, source_filepath, dest_filepath)
# Download file
sftp_client = sftp.connect(self.config[self.object_id], self.logger, self.object_id)
success = sftp.download_file(sftp_client, source_filepath, dest_filepath, self.logger, self.object_id)
sftp_client.close()
# Re-compress, from ZIP to LZMA
if success:
self.logger.write(self.object_id, day.strftime("%Y-%m-%d") + ": Re-compressing from ZIP to LZMA...")
lzma_filepath = utils.recompress_file(dest_filepath)
......@@ -45,8 +50,6 @@ class Collector(utils.LogConfig):
else:
self.logger.write(self.object_id, day.strftime("%Y-%m-%d") + ": Done. Archive file: " + lzma_filepath)
sftp_client.close()
def __get_dest_filepath__(self, day: datetime.date, source_filepath: str) -> str:
# Build destination directory
dest_dir = os.path.join(self.config["LOCAL"]["base_dir"],
......@@ -62,25 +65,4 @@ class Collector(utils.LogConfig):
return dest_filepath
def __sftp_connect__(self) -> paramiko.SFTPClient:
client = paramiko.SSHClient()
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
try:
client.connect(hostname=self.config[self.object_id]["host"],
username=self.config[self.object_id]["user"],
key_filename=self.config[self.object_id]["private_key_path"])
except Exception as e:
self.logger.write(self.object_id, "Failed to connect to server:\n" + str(e))
raise e
else:
return client.open_sftp()
def __sftp_download__(self, sftp_client: paramiko.SFTPClient, source_filepath: str, dest_filepath: str) -> bool:
try:
sftp_client.get(source_filepath, dest_filepath)
except FileNotFoundError:
msg = "File not found: " + source_filepath
self.logger.write(self.object_id, msg)
raise FileNotFoundError(msg)
success = os.path.exists(dest_filepath)
return success
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment