Merge pull request #182 from multiflexi/ftp_publisher

minor improvement of FTP publisher

src/publishers/publishers/ftp_publisher.py

@@ -1,15 +1,31 @@
+"""Publisher for publishing to FTP and SFTP server.
+
+Raises:
+    Exception: _description_
+"""
 import datetime
 import ftplib
 import os
 from base64 import b64decode
 from urllib.parse import urlsplit
 import paramiko
+import mimetypes
+
 
 from .base_publisher import BasePublisher
 from shared.schema.parameter import Parameter, ParameterType
 
 
 class FTPPublisher(BasePublisher):
+    """FTP Publisher class.
+
+    Arguments:
+        BasePublisher -- Publisher base class
+
+    Raises:
+        Exception: _description_
+    """
+
     type = "FTP_PUBLISHER"
     name = "FTP Publisher"
     description = "Publisher for publishing to FTP server"
@@ -19,27 +35,23 @@ class FTPPublisher(BasePublisher):
     parameters.extend(BasePublisher.parameters)
 
     def publish(self, publisher_input):
+        """Publish to FTP or SFTP server.
 
-        try:
-            ftp_url = publisher_input.parameter_values_map['FTP_URL']
+        Arguments:
+            publisher_input -- intput data for publisher
 
+        Raises:
+            Exception: _description_
+        """
+        try:
+            ftp_url = publisher_input.parameter_values_map["FTP_URL"]
             mime_type = publisher_input.mime_type[:]
-
-            filename = ''
-            if mime_type[:] == 'application/pdf':
-                filename = 'file_' + datetime.datetime.now().strftime("%d-%m-%Y_%H:%M") + '.pdf'
-            elif mime_type[:] == 'text/plain':
-                filename = 'file_' + datetime.datetime.now().strftime("%d-%m-%Y_%H:%M") + '.txt'
-            elif mime_type[:] == 'application/json':
-                filename = 'file_' + datetime.datetime.now().strftime("%d-%m-%Y_%H:%M") + '.json'
-            elif mime_type[:] == 'text/html':
-                filename = 'file_' + datetime.datetime.now().strftime("%d-%m-%Y_%H:%M") + '.html'
-
+            file_extension = mimetypes.guess_extension(mime_type)
+            filename = f"file_{datetime.datetime.now().strftime('%d-%m-%Y_%H:%M')}{file_extension}"
             data = publisher_input.data[:]
-
             bytes_data = b64decode(data, validate=True)
 
-            f = open(filename, 'wb')
+            f = open(filename, "wb")
             f.write(bytes_data)
             f.close()
 
@@ -51,23 +63,23 @@ def publish(self, publisher_input):
 
             remote_path = ftp_data.path + filename
 
-            if ftp_data.scheme == 'sftp':
+            if ftp_data.scheme == "sftp":
                 ssh_port = ftp_data.port if ftp_data.port else 22
                 ssh = paramiko.SSHClient()
                 ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
                 ssh.connect(hostname=ftp_hostname, port=ssh_port, username=ftp_username, password=ftp_password)
                 sftp = ssh.open_sftp()
                 sftp.put(filename, remote_path)
                 sftp.close()
-            elif ftp_data.scheme == 'ftp':
+            elif ftp_data.scheme == "ftp":
                 ftp_port = ftp_data.port if ftp_data.port else 21
                 ftp = ftplib.FTP()
                 ftp.connect(host=ftp_hostname, port=ftp_port)
                 ftp.login(ftp_username, ftp_password)
-                ftp.storbinary('STOR ' + remote_path, open(filename, 'rb'))
+                ftp.storbinary("STOR " + remote_path, open(filename, "rb"))
                 ftp.quit()
             else:
-              raise Exception("Schema '{}' not supported, choose 'ftp' or 'sftp'".format(ftp_data.scheme))
+                raise Exception("Schema '{}' not supported, choose 'ftp' or 'sftp'".format(ftp_data.scheme))
         except Exception as error:
             BasePublisher.print_exception(self, error)
         finally: